Xla: Mixed precision training slower than float32

[myleott]

❓ Questions and Help

Background
I'd like to get mixed precision training working in pytorch/xla. I'm following the general recipe introduced in [1], where activations and gradients are in bfloat16 and optimization is done on a float32 master copy of the weights. I've used this same recipe successfully on GPUs for several years (substituting float16 for bfloat16). This same recipe seems to be used in Tensorflow as well [2][3], roughly:
1) cast input to bfloat16 and do forward pass
2) cast model outputs to float32 to compute the loss
3) do backward pass in bfloat16
4) all-reduce bfloat16 gradients
5) cast bfloat16 gradients to float32
6) do optimization step in float32 (updating a master copy of the weights stored in float32)
7) cast updated float32 master weights back to bfloat16

How this affects model accuracy
In practice the above recipe results in comparable model quality as training in pure float32. I've found this both on GPUs and TPUs.

How this affects training speed
The above mixed precision training recipe adds significant overhead. However, V100 GPUs have Tensor Cores which make float16 operations significantly (5-9x) faster. So in practice we typically see 2-3x speedup from mixed precision training.

Unfortunately, on TPUs I'm seeing the above mixed precision recipe trains much slower than even pure float32 training. Is this expected?

Instructions to reproduce
I've created a small mixed precision benchmark script which measures the time for 50 steps (after a 50 step warmup). I get the following results:

$ python test_mixed_precision.py --mode=fp32   # 6.0 seconds
$ python test_mixed_precision.py --mode=bf16   # 5.3 seconds (~12% faster than pure float32)
$ python test_mixed_precision.py --mode=mixed  # 28.7 seconds

[1] https://arxiv.org/abs/1710.03740
[2] https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus
[3] https://www.tensorflow.org/guide/keras/mixed_precision

[dlibenzi]

Hi Myle,

the gain between F32 and BF16 on TPUs is more limited WRT GPUs, as TPU already do certain operations in BF16 automatically.
Unless you set the high precision flag on, which we do only for tests:

https://github.com/pytorch/xla/blob/e4e6318657b2ea5e6fc132888f8cb2534ea3d2ed/test/pytorch_test_base.py#L350

Most of the F32/BF16 gains come from memory bandwidth savings in moving tensor data to/from HBM.
The other gain comes from input feeding bandwidth if you feed F32 data (ie, vision models), where input bound models can become not so, when feeding BF16.

That does not mean that there could be a model where hand tuning can lead better results, but in general if pure F32 and pure BF16 performance are pretty close, I would not bother with mixed precision (and use either pure BF16 or pure F32 - depending on the convergence data).

[bohoro]

Hi,

Mixed precision is important to many of our training workflows can we relook at Myle's ask?

[myleott]

More precisely, mixed precision reduces memory usage, which is important for us to train larger models.

[taylanbil]

I ran the repro w/ 1 core:

• e2e wall time 20s for fp32, 18s for bf16, 64s for mixed
• ExecuteTime's are 9.3 s | 8s | 54s

I'll get the hlo's next.

[taylanbil]

We were able to isolate the slowness to the code that copies gradients back to fp32. Changing that implementation, I'm able to achieve performance that sits in the middle of fp32 and bf16:

# fp32
initializing dataloader
initializing model
num model params: 156661760
initializing optimizer
initializing paraloader
beginning warmup
end warmup, begin measurement
end measurement, time for rank 0: 4.96778678894043

real    0m19.384s
user    0m5.112s
sys     0m1.788s

# bf16
initializing dataloader
initializing model
num model params: 156661760
initializing optimizer
initializing paraloader
beginning warmup
end warmup, begin measurement
end measurement, time for rank 0: 4.355741500854492

real    0m18.024s
user    0m40.840s
sys     0m1.532s

# mixed
initializing dataloader
initializing model
num model params: 156661760
initializing optimizer
initializing paraloader
beginning warmup
end warmup, begin measurement
end measurement, time for rank 0: 4.80745530128479

real    0m19.647s
user    0m39.988s
sys     0m1.604s

Here's the code changes that gets me the above result; essentially the only meaningful change is, not to keep the fp32 weights in a flat array and use offset and flatten to copy back and forth, but to keep the original shapes intact.

$ diff test_mixed_precision-ORIG.py test_mixed_precision.py
10a11
> import torch_xla.debug.metrics as met
27c28
<         self.fp32_optimizer = optim_cls([self.fp32_params], **kwargs)
---
>         self.fp32_optimizer = optim_cls(self.fp32_params, **kwargs)
38c39,40
<         self.fp32_params.grad.data.zero_()
---
>         for p in self.fp32_params:
>             p.grad.data.zero_()
44,56c46,55
<     total_param_size = sum(p.data.numel() for p in bf16_params)
<     fp32_buf = torch.zeros(
<         total_param_size, dtype=torch.float32, device=bf16_params[0].device
<     )
<
<     offset = 0
<     for p in bf16_params:
<         numel = p.data.numel()
<         fp32_buf.data[offset:offset+numel].copy_(p.data.view(-1))
<         offset += numel
<
<     fp32_params = torch.nn.Parameter(fp32_buf)
<     fp32_params.grad = fp32_buf.new_zeros(total_param_size)
---
>     device = bf16_params[0].device
>     fp32_buf = [
>         torch.zeros(p.shape, dtype=torch.float32, device=device)
>         for p in bf16_params
>     ]
>     for p16, p32 in zip(bf16_params, fp32_buf):
>         p32.copy_(p16.data)
>     fp32_params = [torch.nn.Parameter(p32) for p32 in fp32_buf]
>     for p32 in fp32_params:
>         p32.grad = p32.new_zeros(p32.shape)
61,67c60,62
<     offset = 0
<     for p in bf16_params:
<         assert p.requires_grad and p.grad is not None
<         grad_data = p.grad.data
<         numel = grad_data.numel()
<         fp32_params.grad.data[offset:offset+numel].copy_(grad_data.view(-1))
<         offset += numel
---
>     for p16, p32 in zip(bf16_params, fp32_params):
>         assert p16.requires_grad and p16.grad is not None
>         p32.grad.data.copy_(p16.grad.data)
71,76c66,68
<     offset = 0
<     for p in bf16_params:
<         assert p.requires_grad
<         numel = p.numel()
<         p.data.copy_(fp32_params.data[offset:offset+numel].view_as(p))
<         offset += numel
---
>     for p16, p32 in zip(bf16_params, fp32_params):
>         assert p32.requires_grad and p32.grad is not None
>         p16.data.copy_(p32.data)
100c92
<                 nn.Dropout(),
---
>                 #nn.Dropout(),
110c102
<                 nn.Dropout(0.1),
---
>                 #nn.Dropout(0.1),
151c143
<     print("initializing model")
---
>     xm.master_print("initializing model")
154c146
<     print("num model params: {}".format(sum(p.numel() for p in model.parameters())))
---
>     xm.master_print("num model params: {}".format(sum(p.numel() for p in model.parameters())))
156c148
<     print("initializing optimizer")
---
>     xm.master_print("initializing optimizer")
162c154
<     print("initializing paraloader")
---
>     xm.master_print("initializing paraloader")
165c157
<     print("beginning warmup")
---
>     xm.master_print("beginning warmup")
166a159,161
>         #if i == 1:
>         #    xm.master_print(met.metrics_report())
>         #metsumm(i)
168c163
<             print("end warmup, begin measurement")
---
>             xm.master_print("end warmup, begin measurement")
178d172
<
181d174
<         optimizer.zero_grad()
185c178
<             print(
---
>             xm.master_print(
189a183,185
>         optimizer.zero_grad()
>     #xm.master_print(met.metrics_report())
>
190a187,194
> def metsumm(stepno=''):
>     import torch_xla.debug.metrics as met
>     x = met.metrics_report().split('\n')
>     for i, line in enumerate(x):
>         if 'CompileTime' in line or 'aten::' in line:
>             key = line.split()[-1]
>             value = x[i+1].split()[-1]
>             print('step {}, key {}, value {}'.format(stepno, key, value))
193c197,198
<     xmp.spawn(main, args=(), nprocs=8)
---
>     xmp.spawn(main, args=(), nprocs=1)
>     #xmp.spawn(main, args=(), nprocs=8)

[myleott]

Thanks! I confirmed after this fix that speeds for mixed precision are now somewhere between pure bf16 and pure fp32, as expected :)