support for bf16 in lightning trainer

[yuvalkirstain]

🚀 Feature

Many models (e.g. T5) were trained using bf16, and as such, using fp16 results with nans. Now pytorch supports bf16. I think that supporting bf16 is a must have feature (and due to my ongoing research that builds upon Pytorch Lightning it will be amazing if it will be integrated fast :) )

Motivation

Many models (especially in NLP like T5) cannot be trained with fp16 when using pytorch-lightning. This makes some of them practically unusable (due to the memory consumption), and others just train slow.

Pitch

I want to pass an argument to the trainer that says precision=bf16 and it will train it with bf16.

[leezu]

With pytorch/pytorch#61002 and nccl pytorch/pytorch#61799 the following patch can setup PL for bf16 training if the user calls th.set_autocast_gpu_dtype(th.bfloat16). Adding a separate API that automatically calls th.set_autocast_gpu_dtype(th.bfloat16) based on a config option may be helpful.

modified   pytorch_lightning/plugins/precision/native_amp.py
@@ -42,6 +42,12 @@ class NativeMixedPrecisionPlugin(MixedPrecisionPlugin):
         model: 'pl.LightningModule',
         closure_loss: torch.Tensor,
     ) -> torch.Tensor:
+        if hasattr(torch, 'get_autocast_gpu_dtype') and torch.get_autocast_gpu_dtype() == torch.bfloat16:
+            from pytorch_lightning.core.lightning import warning_cache
+            warning_cache.warn(
+                "Skipping torch.cuda.amp.GradScaler() in NativeMixedPrecisionPlugin as torch.bfloat16 is used "
+            )
+            return super().pre_backward(model, closure_loss)
         closure_loss = self.scaler.scale(closure_loss)
         return super().pre_backward(model, closure_loss)
 
@@ -58,6 +64,8 @@ class NativeMixedPrecisionPlugin(MixedPrecisionPlugin):
                 f"native PyTorch amp and lbfgs are not compatible (optimizer {optimizer_idx})."
                 " To request, please file a Github issue in PyTorch and tag @mcarilli"
             )
+        if hasattr(torch, 'get_autocast_gpu_dtype') and torch.get_autocast_gpu_dtype() == torch.bfloat16:
+            return True  # make_optimizer_step outside of the precision plugin
         result = True
         if model.automatic_optimization:
             result = lambda_closure()

[yuvalkirstain]

Thanks for the quick reply @leezu !
Can you please specify the exact steps needed to train a model with bf16 using PL ?
Thanks so much!

[leezu]

@yuvalkirstain apply above patch to pytorch-lightning AND in your training script, call th.set_autocast_gpu_dtype(th.bfloat16) before you begin training. Your training script must be setup for precision=16 training with native amp. Note that this is a small hack to enable bf16 training. Proper support should be added to Pytorch-Lightning

[yuvalkirstain]

Thanks @leezu !
Do I need to pip install pytorch-nightly as well (I guess so)? How do I import th?

[leezu]

th refers to torch. You need pytorch version that includes pytorch/pytorch@324673a If you do multi-gpu training, you need to compile pytorch from source should also apply the commits from pytorch/pytorch#61799 before compiling

[carmocca]

Awesome @leezu !

Feel free to open a PR when that is available in a torch release.

[tchaton]

Hey @carmocca @awaelchli, what do you think about Trainer(precision=torch.bfloat16).

[carmocca]

what do you think about Trainer(precision=torch.bfloat16)

The problem with passing the torch dtype is that then users will expect to be able to pass any torch dtype, which is not supported right now