MultiSteps

[GeophyAI]

Hi, I'm using MultiSteps for gradient accumulation. When the batchsize is 8 and without MultiSteps, the code cost ~13G GPU memory. When I try to use MultiSteps with every_k_schedule set to 8, i.e. only one sample at a time, I got an OOM. My codes are something like following:

@jax.jit
def step(params):

  batch_data = ...

  for i in range(STEP_PER_EPOCH):

     y=predict(x)

    loss, gradient = compute_gradient(params)
    updates, opt_state = opt.update(gradient, opt_state)

    params = optax.apply_updates(params, updates)

  return params

Do I have to split an upadte function from step?

[vroulet]

In any case, gradient accumulation induces an increase of memory, it's rather well explained in this paper. For e.g. sgd without momentum, this means that MultiSteps should double the memory requirements.
That said, there may be memory usage leaks in the implementation that I'll investigate (the code is using jnp.where in place of jax.lax.cond, maybe @mtthss can chime in on that).

[GeophyAI]

Is there a flexiable way to accumulate the gradient rather than using Multisteps, just like gradient_batch += gradient without increasing too much memory? I tried several ways to do it, but everytime I got an OOM.

[GeophyAI]

I have solved OOM after moving the batch gradient calculation from jax.jit warpped function.