Train with gradient accumulation#


Gradient accumulation is only available for transformer models in weight-streaming execution.

Gradient accumulation is a technique that allows training with larger effective batch sizes than can physically fit in the available device memory. As illustrated in Fig. 10, a batch size exceeding the device memory capacity can be divided into smaller sub-batches (minibatches). Each sub-batch is computed separately, and the resulting gradients are accumulated across sub-batches before the final update to the network weights occurs. This way, gradient accumulation emulates a bigger batch size by running multiple smaller sub-batches and combining the results.

If gradient accumulation is enabled, the compiler will drop the model to a smaller sub-batch size either if the original batch size does not fit into device memory or the compiler estimates that a lower sub-batch will achieve significantly better samples/second performance.


Fig. 10 Gradient accumulation computation#

How to enable#

To enable gradient accumulation, set use_cs_grad_accum: True in the runconfig section of the model params yaml file.

The software stack will automatically choose an appropriate sub-batch size that cleanly divides the batch size, users are not expected to manually choose a sub-batch size.

Known issues and limitations#

The current known limitations include:

  • Vision models (CNNs) are not supported.

  • Batch normalization is not supported.

  • Support is limited to networks whose gradients and statistics can be accumulated across sub-batches.

  • Gradient accumulation support for T5 model is experimental and you may experience extended compilation times.

  • The compiler selects the largest sub-batch size that can be compiled and evenly divides the original model batch size.