Gradient Accumulation & Optax MultiSteps

[agemagician]

Hello,

I am testing "optax.MultiSteps" for gradient accumulation on Colab TPU V2, but every time I use anything above 1, I get an OOM.

It seems that it increases the memory requirements equivalently to increase the batch size, which should not be the case.

My understanding is that all I need to do is to use it, is to call it after the optimizer and increase the training batch size by multiplying it by the gradient accumulation steps.
My understanding is also that I should be able to increase the batch size by using gradient accumulation.
In my case, if the per_device_train_batch_size is 8, then I can only set gradient_accumulation_steps to 1, and if the per_device_train_batch_size is 2, then I can then increase the gradient_accumulation_steps to 4.

 train_batch_size = (
        int(training_args.per_device_train_batch_size)
        * jax.device_count()
        * training_args.gradient_accumulation_steps
    )
optimizer = optax.adamw(
            learning_rate=learning_schedule_fn,
            b1=training_args.optim_beta1,
            b2=training_args.optim_beta2,
            weight_decay=training_args.optim_weight_decay,
            mask=decay_mask_fn,
        )

if training_args.gradient_accumulation_steps > 1:
    optimizer = optax.MultiSteps(
        optimizer, training_args.gradient_accumulation_steps
    )

The rest of the code should be the same.

Is my understanding is correct or there is something else we should take care of while using the MultiSteps function?

I have posted this question on Optax repo, but it will be great to hear your feedback.
Feel free to share another library or a method to apply gradient accumulation efficiently and easily.

[agemagician]

Thanks. The optax team found the problem.

[andsteing]

Great! Could you add a link to the discussion/solution for further reference?

[agemagician]

Great! Could you add a link to the discussion/solution for further reference?

Sure, the problem and the solution are here:
google-deepmind/optax#320