Support gradient accumulation using a forward fn decorator with lax.scan #614
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds a new forward function decorator into the learner config which allows users to define their own decorators that could transform the forward function behaviour. One such example of this use case is gradient accumulation which is what this PR is about.
Specifically, this PR enables gradient accumulation using a a new forward function decorator called
with_minibatch_steps
which internally implements alax.scan
loop over minibatches sliced with dynamic input batch slicing. The new forward function decorator, calledwith_minibatch_steps
, wraps the forward function and accumulates gradients using ajax.lax.scan
for the number of gradient accumulationsteps
which the user specifies. To use this decorator the learner config is extended with an optionalforward_fn_decorator
parameter which will be specified by the user in the following way if they want to enable gradient accumulation:Example:
Since we scan over the forward function (instead of the value_and_grad func) we need to have a way to compute the gradients in a minibatched manner therefore we wrap the forward function with a
custom_vjp
implementation to compute the gradients and accumulate during the forward phase. Since the gradients are accumulated and computed during the forward phase we simply pass them to the backward phase and use them as is. This allows us to compute and accumulate the gradients in a memory efficient (minibatched) manner.A similar effort for implementing grad accumulation using a scan was mentioned by @apoorvtintin.