Loss Scaling on NPU

Objective

In mixed precision computing, the narrower dynamic range of float16 leads to floating-point overflow/underflow during gradient calculation as well as parameter update failure. Loss scaling can prevent the divergence during mixed precision training.

The loss scaling method multiplies the loss from forward computation by a loss scale factor S. This amplifies gradients during backward gradient computation, minimizing overflow issues caused by small gradient values that cannot be represented in FP16. After the parameter gradient aggregation and before the optimizer updates parameters, the aggregated parameter gradient is multiplied by 1/S.

Dynamic loss scaling checks the gradient floating-point compute exceptions during training and selects the loss scale S adaptively with the gradient change in the training process.

In practice, due to differences in floating-point computation characteristics, the AI processor may differ from the GPU in floating-point exception detection and other related aspects during computation. Therefore, loss scaling must be correctly enabled on the NPU.