

First we wrap the dataset in a DataLoader. 
Then we can enable gradient checkpointing by calling the model's [~PreTrainedModel.gradient_checkpointing_enable] method. 
When we initialize the Accelerator 
we can specify if we want to use mixed precision training and it will take care of it for us in the [prepare] call. 
During the prepare 
call the dataloader will also be distributed across workers should we use multiple GPUs. We use the same 8-bit optimizer from the earlier example.
Finally, we can add the main training loop. Note that the backward call is handled by 🤗 Accelerate. We can also see
how gradient accumulation works: we normalize the loss, so we get the average at the end of accumulation and once we have 
enough steps we run the optimization. 
Implementing these optimization techniques with 🤗 Accelerate only takes a handful of lines of code and comes with the 
benefit of more flexibility in the training loop. For a full documentation of all features have a look at the 
Accelerate documentation.
Efficient Software Prebuilds
PyTorch's pip and conda builds come prebuilt with the cuda toolkit 
which is enough to run PyTorch, but it is insufficient if you need to build cuda extensions.
At times, additional efforts may be required to pre-build some components. For instance, if you're using libraries like apex that 
don't come pre-compiled. In other situations figuring out how to install the right cuda toolkit system-wide can be complicated. 
To address these scenarios PyTorch and NVIDIA released a new version of NGC docker container which already comes with 
everything prebuilt. You just need to install your programs on it, and it will run out of the box.
This approach is also useful if you want to tweak the pytorch source and/or make a new customized build.
To find the docker image version you want start with PyTorch release notes, 
choose one of the latest monthly releases. Go into the release's notes for the desired release, check that the environment's 
components are matching your needs (including NVIDIA Driver requirements!) and then at the very top of that document go 
to the corresponding NGC page. If for some reason you get lost, here is the index of all PyTorch NGC images.
Next follow the instructions to download and deploy the docker image.
Mixture of Experts
Some recent papers reported a 4-5x training speedup and a faster inference by integrating
Mixture of Experts (MoE) into the Transformer models.
Since it has been discovered that more parameters lead to better performance, this technique allows to increase the 
number of parameters by an order of magnitude without increasing training costs.
In this approach every other FFN layer is replaced with a MoE Layer which consists of many experts, with a gated function 
that trains each expert in a balanced way depending on the input token's position in a sequence.