Gradient Accumulation in Axlearn

[apoorvtintin]

Gradient accumulation allows training with higher batch sizes without scaling out.

Added a new learner type learner.klass: 'axlearn.common.learner.AccumulatedLearner'

At a high level the optimization does the following:

1. Input batch is split into even microbatches.
2. Creates a buffer for gradients and metrics.
3. Runs forward and backward pass for each microbatch in a loop summing up the gradients and aggregating metrics.
4. Average gradients across microbatches and normalize metrics.

Configuration changes:

• Number of microbatches are specified during configuration through option micriobatches in the learner.

[ruomingp]

Could you explain why this is needed? We find usually it's more efficient to use either a larger mesh or a smaller batch size.

[apghml]

Axlearn already has metric accumulation classes that are used by evalers. Could those be reused here instead of defining new classes?

[apoorvtintin]

These metric accumulation classes are stateless so that they are usable as carry by jax.lax.scan unlike the ones in the evaler, I can make the class structure similar though.

[apghml]

Does the existing learner class use this? If not, we should try to be consistent with its API.

[apoorvtintin]

Could you please be more specific, is the concern naming of members?

[apghml]

I wonder if instead of having a separate learner for microbatching, it would be more flexible to have a generic way of wrapping a ForwardFn so that it uses Jax.lax.map to run the microbatches. Beyond avoiding the need to add a new learner, it would also allow for other microbetching uses outside of learner, eg inference or in second order optimizers.

[apoorvtintin]

jax.lax.map gives no guarantees of sequential execution of microbatches which is the key quality of gradient accumulation.

[apghml]

I’m ooo this week. I have left some preliminary comments for now.

[apghml]

Closing since gradient accumulation functionality has been implemented via #614