Examples how to use fori_loop for gradient accumulation & more clear exceptions

[qGentry]

Hey folks,

It would be nice to have some examples in Flax NNX documentation on how to do gradient accumulation with nnx.fori_loop.

I've tried lots of different combination of "how to pass arguments, what I can add to closure and what I can't" myself and haven't succeed.

My train step looks something like this:

# TrainSession is NamedTuple with 'model: nnx.Module' and 'optimizer: nnx.Optimizer' fields

def get_training_step(
    forward_pass_cfg: ForwardPassConfigBase,
    batch_sharding: PyTree,
    num_minibatches: int = 1,
) -> Callable[
    [PyTree, TrainSession],
    tuple[DeviceArray, PyTree],
]:
    forward_pass_fn = get_forward_pass(forward_pass_cfg)

    @nnx.jit(donate_argnums=(1,))
    def training_step(
        batch: PyTree,
        train_session: TrainSession,
    ) -> tuple[DeviceArray, PyTree]:

        def _loop_body(
            minibatch_idx,
            carry: tuple[nnx.State, DeviceArray, TrainSession],
        ) -> tuple[nnx.State, DeviceArray, TrainSession]:
            g_accum, loss_accum, train_session = carry
            minibatch = get_sharded_minibatch(
                batch=batch,
                batch_sharding=batch_sharding,
                minibatch_idx=minibatch_idx,
                num_minibatches=num_minibatches,
            )

            def _loss_fn(model) -> tuple[DeviceArray, PyTree]:
                return forward_pass_fn(
                    model=model,
                    batch=minibatch,
                    step=train_session.optimizer.step.value,
                    config=forward_pass_cfg,
                )

            (loss, outputs), grads = nnx.value_and_grad(_loss_fn, has_aux=True)(
                train_session.model
            )
            g_accum = jax.tree.map(
                lambda gm, g: gm + g, g_accum, grads
            )
            loss_accum = loss_accum + loss
            return g_accum, loss_accum, train_session

        g_accum = jax.tree.map(jnp.zeros_like, nnx.state(train_session.model))
        g_accum, loss_accum, train_session = nnx.fori_loop(
            lower=0, 
            upper=num_minibatches,
            body_fun=_loop_body,
            init_val=(g_accum, 0.0, train_session)
        )

        g_accum = jax.tree.map(lambda g: g / num_minibatches, g_accum)
        loss_accum = loss_accum / num_minibatches

        train_session.optimizer.update(train_session.model, g_accum)
        return loss_accum, {}

    return training_step

And I'm getting the following error:

ValueError: nnx.fori_loop requires body function's input and output to have the same reference and pytree structure, but they differ. If the mismatch comes from `outer_index` field, you might have modified reference structure within the body function, which is not allowed.

But I'm not sure I what does it mean and how can I fix it.

It would be nice if flax provided examples on documentation on how to approach it properly, what can be done, what can't be done, how to work with closures, etc.

[vfdev-5]

@qGentry I'll check your code reproducer.
For the info, you can also do the gradient accumulation with optax multisteps: https://optax.readthedocs.io/en/latest/_collections/examples/gradient_accumulation.html

[qGentry]

Hi @vfdev-5 !

Thanks for looking into this issue. Optax implementation of grad accum doesn't really work for me (if I understand correctly, orbax suggest to perform 'number_of_steps * grad_accum_step' instead of ''number_of_steps' training steps). The reason behind it is that I want to eventually use vjp3 for zero-overhead gradient accumulation, have a lot of logic for proper loss scaling in microbatches of single batch, etc.

[vfdev-5]

I could make a small executable reproducer from your code and looks like a bug when passing g_accum structure.
We had something similar previously with nnx.scan complaining about graphdef change in the carry (it is fixed now).
Maybe, you can meanwhile you scan for your case:

import jax
import jax.numpy as jnp
import flax.nnx as nnx

class Model(nnx.Module):
    def __init__(self):
        self.l1 = nnx.Linear(10, 10, rngs=nnx.Rngs(0))
        self.l2 = nnx.Linear(10, 5, rngs=nnx.Rngs(0))

    def __call__(self, x):
        x = self.l1(x)
        x = nnx.relu(x)
        x = self.l2(x)
        return x


def fwd_fn(carry):
    m, g_accum, counter = carry
    def _loss_fn(model_):
        y_preds = model_(x)
        return y_preds.mean()

    loss, grads = nnx.value_and_grad(_loss_fn)(m)
    g_accum = jax.tree.map(
        lambda gm, g: gm + g, g_accum, grads
    )
    counter = counter + 1
    return m, g_accum, counter


model = Model()
x = jax.random.normal(jax.random.key(0), (32, 10))
g_accum = jax.tree.map(jnp.zeros_like, nnx.state(model))

scan_fwd_fn = nnx.scan(fwd_fn, in_axes=nnx.Carry, out_axes=nnx.Carry, length=4) 

model, g_accum, c = scan_fwd_fn((model, g_accum, 0))
print(c)
# Array(4, dtype=int32, weak_type=True)

and something like that with your code:

Code

import jax
import jax.numpy as jnp
import flax.nnx as nnx


class Model(nnx.Module):
    def __init__(self):
        self.l1 = nnx.Linear(10, 10, rngs=nnx.Rngs(0))
        self.l2 = nnx.Linear(10, 5, rngs=nnx.Rngs(0))

    def __call__(self, x):
        x = self.l1(x)
        x = nnx.relu(x)
        x = self.l2(x)
        return x

mymodel = Model()
train_session = {
    "model": mymodel, 
}
batch_ = jnp.ones((32, 10))

@nnx.jit(donate_argnums=(1,), static_argnums=(2, ))
def training_step(batch, train_session, num_minibatches):
    bs = len(batch)

    def _loop_body(carry):
        g_accum, loss_accum, minibatch_idx, train_session = carry
        minibatch = jax.lax.dynamic_slice_in_dim(batch, minibatch_idx * bs, bs, axis=0)
        model = train_session["model"]

        def _loss_fn(model_):
            y_preds = model_(minibatch)
            return y_preds.sum(), y_preds

        (loss, outputs), grads = nnx.value_and_grad(_loss_fn, has_aux=True)(model)
        g_accum = jax.tree.map(
            lambda gm, g: gm + g, g_accum, grads
        )
        loss_accum = loss_accum + loss
        minibatch_idx = minibatch_idx + 1
        return g_accum, loss_accum, minibatch_idx, train_session

    g_accum = jax.tree.map(jnp.zeros_like, nnx.state(train_session["model"]))
    g_accum, loss_accum, _, train_session = nnx.scan(
        _loop_body,
        in_axes=nnx.Carry, 
        out_axes=nnx.Carry, 
        length=num_minibatches,
    )((g_accum, 0.0, 0, train_session))

    g_accum = jax.tree.map(lambda g: g / num_minibatches, g_accum)
    loss_accum = loss_accum / num_minibatches
    return loss_accum, {}

_ = training_step(batch_, train_session, 4)

---

For maintainers, here is a min reproducer of the failure with fori_loop:

import jax
import jax.numpy as jnp
import flax.nnx as nnx


def fwd_fn(i, carry):
    g_accum = carry
    grads = jax.tree.map(jnp.ones_like, g_accum)
    g_accum = jax.tree.map(lambda gm, g: gm + g, g_accum, grads)
    return g_accum


model = nnx.Linear(10, 10, rngs=nnx.Rngs(0), use_bias=False)
g_accum = jax.tree.map(jnp.zeros_like, nnx.state(model))
g_accum = nnx.fori_loop(0, 2, fwd_fn, g_accum)