Ascend O2 BF16 Optimizer

Background

During LLM training, limited memory is a common problem. Optimizers and gradients consume a large amount of memory, and it is difficult to reduce that usage significantly.

Solution

Use half-precision optimizer parameters and gradients to reduce memory usage significantly.

Usage and Theoretical Benefit

You can enable each feature independently with its corresponding parameter.

N indicates the number of model parameters, and DP indicates the data parallel degree when the distributed parallel optimizer is enabled. The memory reduction achieved by each feature is shown in the table.

Feature Name	Feature Parameter	Memory Reduction (Bytes)
Half-precision optimizer	--o2-optimizer	4N/DP
Half-precision gradient	--o2-gradient	2N

If you do not enable the distributed parallel optimizer, then DP = 1.

Usage Scenarios

1. Enable this feature for comparison in the same training scenario. Compared with the Megatron native mixed-precision scenario, the O2 feature is expected to affect accuracy.
2. At present, the O2 feature supports only training scenarios where the parameter type is BF16.

Usage Notes

1. The half-precision optimizer and half-precision gradient features do not depend on each other. You can enable them separately or together.
2. You cannot use half-precision gradients together with the Matmul-Add gradient accumulation feature because the operator does not support this combination. Therefore, you must use the o2-gradient parameter together with the no-gradient-accumulation-fusion parameter.
3. The O2 feature is theoretically expected not to affect performance. In practice, based on the Llama-2-7B 4K pretraining data, the performance impact is less than 1%.