

For additional information on tf32 vs other precisions, please refer to the following benchmarks: 
RTX-3090 and
A100.
Flash Attention 2
You can speedup the training throughput by using Flash Attention 2 integration in transformers. Check out the appropriate section in the single GPU section to learn more about how to load a model with Flash Attention 2 modules. 
Optimizer choice
The most common optimizer used to train transformer models is Adam or AdamW (Adam with weight decay). Adam achieves 
good convergence by storing the rolling average of the previous gradients; however, it adds an additional memory 
footprint of the order of the number of model parameters. To remedy this, you can use an alternative optimizer. 
For example if you have NVIDIA/apex installed for NVIDIA GPUs, or ROCmSoftwarePlatform/apex for AMD GPUs, adamw_apex_fused will give you the
fastest training experience among all supported AdamW optimizers.
[Trainer] integrates a variety of optimizers that can be used out of box: adamw_hf, adamw_torch, adamw_torch_fused, 
adamw_apex_fused, adamw_anyprecision, adafactor, or adamw_bnb_8bit. More optimizers can be plugged in via a third-party implementation.
Let's take a closer look at two alternatives to AdamW optimizer:
1. adafactor which is available in [Trainer]
2. adamw_bnb_8bit is also available in Trainer, but a third-party integration is provided below for demonstration.
For comparison, for a 3B-parameter model, like “google-t5/t5-3b”: 
* A standard AdamW optimizer will need 24GB of GPU memory because it uses 8 bytes for each parameter (83 => 24GB)
* Adafactor optimizer will need more than 12GB. It uses slightly more than 4 bytes for each parameter, so 43 and then some extra.
* 8bit BNB quantized optimizer will use only (2*3) 6GB if all optimizer states are quantized.
Adafactor
Adafactor doesn't store rolling averages for each element in weight matrices. Instead, it keeps aggregated information 
(sums of rolling averages row- and column-wise), significantly reducing its footprint. However, compared to Adam, 
Adafactor may have slower convergence in certain cases.
You can switch to Adafactor by setting optim="adafactor" in [TrainingArguments]:
py
training_args = TrainingArguments(per_device_train_batch_size=4, optim="adafactor", **default_args)
Combined with other approaches (gradient accumulation, gradient checkpointing, and mixed precision training) 
you can notice up to 3x improvement while maintaining the throughput! However, as mentioned before, the convergence of 
Adafactor can be worse than Adam. 
8-bit Adam
Instead of aggregating optimizer states like Adafactor, 8-bit Adam keeps the full state and quantizes it. Quantization 
means that it stores the state with lower precision and dequantizes it only for the optimization. This is similar to the 
idea behind mixed precision training.
To use adamw_bnb_8bit, you simply need to set optim="adamw_bnb_8bit" in [TrainingArguments]:
py
training_args = TrainingArguments(per_device_train_batch_size=4, optim="adamw_bnb_8bit", **default_args)
However, we can also use a third-party implementation of the 8-bit optimizer for demonstration purposes to see how that can be integrated.
First, follow the installation guide in the GitHub repo to install the bitsandbytes library 
that implements the 8-bit Adam optimizer.
Next you need to initialize the optimizer. This involves two steps: 
* First, group the model's parameters into two groups - one where weight decay should be applied, and the other one where it should not. Usually, biases and layer norm parameters are not weight decayed. 
* Then do some argument housekeeping to use the same parameters as the previously used AdamW optimizer.