Introducing Faster Training with Lightning and Brain Float16

Lightning introduces BFloat16 support with a single flag to the Trainer, for faster and stable training in lower precision.

Image By Phoeby Naren

PyTorch Lightning v1.5 marks a major leap of reliability to support the increasingly complex demands of the leading AI organizations and prestigious research labs that rely on Lightning to develop and deploy AI at scale.

To better support our fast-growing community, PyTorch Lightning aims at becoming the simplest, most flexible framework for expediting any kind of deep learning research to production.

In Lightning 1.5, you can now use BFloat16 to speed up training keeping high numerical accuracy, without harming your model’s performance. All it takes is Trainer(precision="bf16") .

Keep reading to learn about BFloat16 and how it works under the hood!

Simplifying Mixed Precision Training

BFloat16 maintains the same dynamic range as Float32 while using half the memory. This differs from typical Float16 precision which sacrifices more of the exponent. The Google Research team recently demonstrated that BFloat16 maintains stable training without any additional logic while providing improvements in throughput and memory usage.

Image By Google

With PyTorch 1.10, torch.bloat16 support was added for both CPUs & GPUs using Automatic Mixed Precision (AMP).

Using Lightning to Train Google Transformers

Google released a variety of transformer models trained with TPUs (for example, multilingual-T5). Since these models were trained with BFloat16, they require BFloat16 or Float32 for fine-tuning and running inference.

Up to now Lightning users had to sacrifice performance by upscaling to Float32 training. With Lightning 1.5 we can enable BFloat16 and fine-tune these models how they were meant to be fine-tuned!

By passing Trainer(precision="bf16"), we can fine-tune our Multilingual T5 Transformer model using Lightning Flash.

Next Steps

The Lightning Team is more than ever committed to providing the best experience possible to anyone doing optimization with PyTorch and the PyTorch Lightning API being already stable, breaking changes will be minimal.

If you’re interested in helping out with these efforts, find us on slack!

Built by the PyTorch Lightning creators, let us introduce you to our platform Grid.ai which enables you to scale your model training without worrying about infrastructure, similarly as Lightning automates the training.

Find out more with our docs, and you can get started with Grid.ai for free with just a GitHub or Google Account.