AMP for TPUs by cowanmeg · Pull Request #4740 · pytorch/xla

cowanmeg · 2023-03-09T00:09:11Z

With pytorch/pytorch#96370
Enable AMP on TPUs with bfloat16.

Currently, torch_xla.amp.autocast(args...) aliases torch.cuda.amp.autocast(args...)
This change proposes updating this to torch_xla.amp.autocast(device, args...) which will call the appropriate autocast depending on the XLA device.

Alternatively, users can call the appropriate autocast directly with torch.cuda.amp.autocast(args...) for XLA:GPU devices and torch.amp.autocast('xla', args...) for XLA:TPU devices.

HLO dump of MNIST model from test/test_train_mp_mnist_amp.py

ENTRY %IrToHlo.149 (p0.1: f32[10], p1.2: f32[10,50], p2.4: f32[50], p3.5: f32[50,320], p4.7: f32[20], p5.8: f32[20], p6.9: f32[20], p7.10: f32[20], p8.11: f32[20], p9.12: f32[20,10,5,5], p10.13: f32[10], p11.14: f32[10], p12.15: f32[10], p13.16: f32[10], p14.17: f32[10], p15.18: f32[10,1,5,5], p16.19: f32[128,1,28,28]) -> (f32[128,10]) {
  ...
  %p16.19 = f32[128,1,28,28]{0,3,2,1} parameter(16)
  %p15.18 = f32[10,1,5,5]{0,3,2,1} parameter(15)
  %convolution.20 = f32[128,10,24,24]{3,2,1,0} convolution(f32[128,1,28,28]{0,3,2,1} %p16.19, f32[10,1,5,5]{0,3,2,1} %p15.18), window={size=5x5}, dim_labels=bf01_oi01->bf01
  %p14.17 = f32[10]{0} parameter(14)
  %broadcast.21 = f32[128,24,24,10]{3,2,1,0} broadcast(f32[10]{0} %p14.17), dimensions={3}
  %transpose.22 = f32[128,10,24,24]{1,3,2,0} transpose(f32[128,24,24,10]{3,2,1,0} %broadcast.21), dimensions={0,3,1,2}
  %add.23 = f32[128,10,24,24]{3,2,1,0} add(f32[128,10,24,24]{3,2,1,0} %convolution.20, f32[128,10,24,24]{1,3,2,0} %transpose.22)
 ...
  %broadcast.136 = f32[128,10]{1,0} broadcast(f32[128]{0} %reduce.135), dimensions={0}
  %subtract.137 = f32[128,10]{1,0} subtract(f32[128,10]{1,0} %add.129, f32[128,10]{1,0} %broadcast.136)
  %exponential.138 = f32[128,10]{1,0} exponential(f32[128,10]{1,0} %subtract.137)
  %constant.139 = f32[] constant(0)
  %reduce.144 = f32[128]{0} reduce(f32[128,10]{1,0} %exponential.138, f32[] %constant.139), dimensions={1}, to_apply=%AddComputation.140
  %log.145 = f32[128]{0} log(f32[128]{0} %reduce.144)
  %broadcast.146 = f32[128,10]{1,0} broadcast(f32[128]{0} %log.145), dimensions={0}
  %subtract.147 = f32[128,10]{1,0} subtract(f32[128,10]{1,0} %subtract.137, f32[128,10]{1,0} %broadcast.146)
  ROOT %tuple.148 = (f32[128,10]{1,0}) tuple(f32[128,10]{1,0} %subtract.147)
}

}

HLO dump with AMP

ENTRY %IrToHlo.162 (p0.1: f32[10], p1.3: f32[10,50], p2.6: f32[50], p3.8: f32[50,320], p4.11: f32[20], p5.12: f32[20], p6.13: f32[20], p7.14: f32[20], p8.15: f32[20], p9.17: f32[20,10,5,5], p10.19: f32[10], p11.20: f32[10], p12.21: f32[10], p13.22: f32[10], p14.23: f32[10], p15.25: f32[10,1,5,5], p16.27: f32[128,1,28,28]) -> (bf16[128,10]) {
 ...
  %p16.27 = f32[128,1,28,28]{0,3,2,1} parameter(16)
  %convert.28 = bf16[128,1,28,28]{0,3,2,1} convert(f32[128,1,28,28]{0,3,2,1} %p16.27)
  %p15.25 = f32[10,1,5,5]{0,3,2,1} parameter(15)
  %convert.26 = bf16[10,1,5,5]{0,3,2,1} convert(f32[10,1,5,5]{0,3,2,1} %p15.25)
  %convolution.29 = bf16[128,10,24,24]{3,2,1,0} convolution(bf16[128,1,28,28]{0,3,2,1} %convert.28, bf16[10,1,5,5]{0,3,2,1} %convert.26), window={size=5x5}, dim_labels=bf01_oi01->bf01
  %p14.23 = f32[10]{0} parameter(14)
  %convert.24 = bf16[10]{0} convert(f32[10]{0} %p14.23)
  %broadcast.30 = bf16[128,24,24,10]{3,2,1,0} broadcast(bf16[10]{0} %convert.24), dimensions={3}
  %transpose.31 = bf16[128,10,24,24]{1,3,2,0} transpose(bf16[128,24,24,10]{3,2,1,0} %broadcast.30), dimensions={0,3,1,2}
  %add.32 = bf16[128,10,24,24]{3,2,1,0} add(bf16[128,10,24,24]{3,2,1,0} %convolution.29, bf16[128,10,24,24]{1,3,2,0} %transpose.31)
  ...
  %broadcast.149 = bf16[128,10]{1,0} broadcast(bf16[128]{0} %reduce.148), dimensions={0}
  %subtract.150 = bf16[128,10]{1,0} subtract(bf16[128,10]{1,0} %add.142, bf16[128,10]{1,0} %broadcast.149)
  %exponential.151 = bf16[128,10]{1,0} exponential(bf16[128,10]{1,0} %subtract.150)
  %constant.152 = bf16[] constant(0)
  %reduce.157 = bf16[128]{0} reduce(bf16[128,10]{1,0} %exponential.151, bf16[] %constant.152), dimensions={1}, to_apply=%AddComputation.153
  %log.158 = bf16[128]{0} log(bf16[128]{0} %reduce.157)
  %broadcast.159 = bf16[128,10]{1,0} broadcast(bf16[128]{0} %log.158), dimensions={0}
  %subtract.160 = bf16[128,10]{1,0} subtract(bf16[128,10]{1,0} %subtract.150, bf16[128,10]{1,0} %broadcast.159)
  ROOT %tuple.161 = (bf16[128,10]{1,0}) tuple(bf16[128,10]{1,0} %subtract.160)
}

cowanmeg · 2023-03-09T00:11:32Z

@miladm @JackCaoG

JackCaoG · 2023-03-09T00:13:53Z

Do you have a corresponding upstream pr? I saw you use torch.xla.amp.autocast but I don't see where it is defined. Seems like it is part of the torch

miladm · 2023-03-13T22:47:00Z

  loss_fn = nn.NLLLoss()
-  scaler = GradScaler(use_zero_grad=FLAGS.use_zero_grad)
+  if device_hw == 'TPU':
+        autocast = torch.xla.amp.autocast


Question: can you point me to the mapping the explains what op is autocasted and what op is not?

The op list is registered on the pytorch side here: https://github.com/pytorch/pytorch/blob/8455aac198e4ee68003f2b38a5e15631fb82690c/aten/src/ATen/autocast_mode.cpp#L520

miladm · 2023-03-13T22:47:56Z

@chandrasekhard2 @cowanmeg who is exploring the ResNet model test? is this part of this PR or a separate PR?

chandrasekhard2 · 2023-03-13T22:49:27Z

@chandrasekhard2 @cowanmeg who is exploring the ResNet model test? is this part of this PR or a separate PR?

I will send a PR today. Thanks!

miladm · 2023-03-13T22:49:27Z

@chandrasekhard2 @cowanmeg can you please share a few words on how the AMP perf looks like at the moment? it's fine if we need more work to improve numbers.

wonjoo-wj · 2023-03-13T22:50:40Z

nit: I assume this PR needs your PyTorch changes -- you can pin the PyTorch PR in this PR so it builds with your PyTorch PR. Example of PyTorch pinning here.

miladm · 2023-03-14T23:54:01Z

@cowanmeg can we do a operator by operator mapping of the current model vs. what we see in HLO to verify what we added to the upstream pytorch for AMP maps to what emerges on the HLO side?

steventk-g · 2023-03-15T00:02:25Z

+    elif device_hw == 'GPU':
+        autocast = torch.cuda.amp.autocast
+        # GradScaler only used for GPU
+        scaler = GradScaler(use_zero_grad=FLAGS.use_zero_grad)


Why don't we use a scaler for TPU?

That comment was a bit misleading. GradScaler is necessary for float16 loss but not bfloat16. float16 has less dynamic range than float32 since it uses only 5 exponent bits vs 8 exponent bits, so the loss needs to be scaled to prevent small values from disappearing. bfloat16 has the same dynamic range as float32 so doesn't have this problem.

That means TPU uses bfloat16 by default?

For AMP, yes!

chandrasekhard2 · 2023-03-15T08:52:09Z

@chandrasekhard2 @cowanmeg who is exploring the ResNet model test? is this part of this PR or a separate PR?

https://github.com/pytorch/xla/pull/4776/files

vanbasten23 · 2023-04-12T02:22:52Z

+
+// CastPolicy::lower_precision_fp General_DeviceType
+template<DeviceType device_type, class Redispatch, Redispatch* F, class Ret, class... Args>
+struct WrapFunction_<CastPolicy::lower_precision_fp, device_type, Redispatch, F, Ret, guts::typelist::typelist<Args...>> {


wondering how this file is being used in the PR

These are templates that help dispatch an op to its lower precision/fp32/promoted version.

Is the proper place for this file?

steventk-g · 2023-04-12T16:43:34Z

Could you add a few unit tests to the C++ or python side to exercise amp for a simple case, e.g. a single operation?

JackCaoG · 2023-04-21T01:31:51Z

+    if device_hw == 'TPU':
+      scaler = None
+    elif device_hw == 'GPU':
+      scaler = GradScaler(use_zero_grad=FLAGS.use_zero_grad)


Now we have a full AMP story, I think we should add a doc under https://github.com/pytorch/xla/tree/master/docs to explain how to use it on TPU and GPU. We should also mentioned tricks like GradScaler here. WDYT?

I had a https://github.com/pytorch/xla/blob/master/docs/gpu.md#amp-automatic-mixed-precision here but I think it is better to have a standalone doc we can refer to during release.

I agree, will add some documentation.

JackCaoG · 2023-04-21T01:40:09Z

+          enabled=enabled,
+          dtype=torch.float16,
+          cache_enabled=cache_enabled)
+    else:


hmm so CPU and TPU shared the same autocast rule here?

JackCaoG

Mostly LGTM, this is exciting!

I have one ux problem, if user do torch.amp.autocast('xla') on GPU device, is there anyway for us to throw a warning and get them to either use torch_xla's autocast or use cuda directly?

cowanmeg · 2023-04-24T17:39:27Z

Mostly LGTM, this is exciting!

I have one ux problem, if user do torch.amp.autocast('xla') on GPU device, is there anyway for us to throw a warning and get them to either use torch_xla's autocast or use cuda directly?

That code will be in Pytorch, so I don't think we can generate a warning since there is no way distinguish XLA:GPU vs. XLA:TPU that I am aware of.

JackCaoG

LGTM! Let's follow up with the doc update in next pr.

With pytorch/xla#5148, pytorch/xla#4740 With these changes XLA:GPU users should use `torch.cuda.amp.autocast()` for AMP with float16 XLA:TPU users should use `torch.amp.autocast('xla')` for AMP with bfloat16 Pull Request resolved: #96370 Approved by: https://github.com/bdhirsh, https://github.com/malfet

tpu amp

0df4743

cowanmeg added the DO_NOT_MERGE_YET label Mar 9, 2023

cowanmeg mentioned this pull request Mar 9, 2023

[AMP] Support XLA:TPU pytorch/pytorch#96370

Closed

miladm requested review from chandrasekhard2, miladm, steventk-g and wonjoo-wj March 13, 2023 22:43

miladm reviewed Mar 13, 2023

View reviewed changes

steventk-g reviewed Mar 15, 2023

View reviewed changes

cowanmeg added 4 commits March 16, 2023 16:45

Add torch pin

7997da0

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

6ac8262

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

042a631

updates

401eaf6

vanbasten23 reviewed Apr 12, 2023

View reviewed changes

Clean up

041fce2

cowanmeg and others added 5 commits April 12, 2023 09:49

Delete test_fsdp_auto_wrap_amp.py

0aa5011

updates

2157e2c

Update autocast key to pick between Cuda and Xla. Unit tests

e247b43

lint

7d215d1

moving code from pt to ptxla

661759a

cowanmeg changed the title ~~[POC] AMP for TPUs~~ AMP for TPUs Apr 14, 2023

move autocast+test from common.sh to run_tests.sh

e4a0daa

cowanmeg requested a review from miladm April 18, 2023 04:44

JackCaoG self-requested a review April 20, 2023 21:46

JackCaoG reviewed Apr 21, 2023

View reviewed changes

Comment thread test/run_tests.sh

JackCaoG reviewed Apr 21, 2023

View reviewed changes

cowanmeg added 5 commits April 24, 2023 18:14

updates

ff8f10b

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

a0437a5

updates with pytorch

410912e

lint

aabbd12

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

1ce912c

JackCaoG approved these changes Apr 26, 2023

View reviewed changes

JackCaoG added REMOVE_TORCH_PIN and removed DO_NOT_MERGE_YET labels Apr 26, 2023

cowanmeg added 10 commits April 26, 2023 19:53

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

dfe1f27

build autocast_mode

5275334

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

0652dd6

Merge branch 'master' of https://github.com/cowanmeg/xla into amp

089f296

Merge branch 'master' of https://github.com/pytorch/xla into amp

e1243d8

Merge branch 'master' of https://github.com/pytorch/xla into amp

85d607a

Merge branch 'master' of https://github.com/pytorch/xla into amp

9241b0a

Disable bazel remote cache

643691c

experiment with no new files

4eba81f

revert back

ffac43f

cowanmeg closed this Jun 8, 2023

This was referenced Jun 8, 2023

AMP for TPUs v2 #5148

Merged

AMP for TPUs v3 #5161

Merged

Conversation

cowanmeg commented Mar 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cowanmeg commented Mar 9, 2023

Uh oh!

JackCaoG commented Mar 9, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

miladm commented Mar 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chandrasekhard2 commented Mar 13, 2023

Uh oh!

miladm commented Mar 13, 2023

Uh oh!

wonjoo-wj commented Mar 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

miladm commented Mar 14, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chandrasekhard2 commented Mar 15, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

steventk-g commented Apr 12, 2023

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JackCaoG left a comment

Choose a reason for hiding this comment

Uh oh!

cowanmeg commented Apr 24, 2023

Uh oh!

JackCaoG left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

cowanmeg commented Mar 9, 2023 •

edited

Loading

miladm commented Mar 13, 2023 •

edited

Loading

wonjoo-wj commented Mar 13, 2023 •

edited

Loading