Gradient Accumulation Results

[AakashKumarNain]

I am noticing a very weird pattern with gradient accumulation on my end. Without gradient accumulation, the training proceeds smoothly, but when gradient accumulation is enabled, the optimization gets stuck at a loss which does not improve after that point. I am not sure what can be the underlying issue here, but a few info on the training run:

1. The model is instantiated with the same seed in both the cases
2. The dataloader is deterministic and produces the same sequences in both the cases
3. The learning rate for different layers, and the weight decay for different layers is exactly the same.

Here is the optimizer definition:

def make_adamw_layerwise_like_nanochat(
    params,
    *,
    d_model: int,                    
    other_peak_lr: float,           
    other_min_lr: float,            
    total_train_steps: int,
    warmup_steps: int = 30,
    b1: float = 0.9,                
    b2: float = 0.95,           
    embedding_lr: float = 3e-3,
    unembedding_lr: float = 3e-4,

):
    """
    Parameter groups:
      - params.embed      -> AdamW with nanochat embedding LR (scaled by (d_model/768)^-0.5)
      - params.lm_head    -> AdamW with nanochat unembedding LR (scaled by (d_model/768)^-0.5)
      - everything else   -> AdamW with warmup+cosine schedule
    """
    emb_lr = embedding_lr * other_peak_lr
    unemb_lr = 1.0 * other_peak_lr


    schedules = {
        "embed": optax.constant_schedule(emb_lr),
        "lm_head": optax.constant_schedule(unemb_lr),
        "other": optax.warmup_cosine_decay_schedule(
            init_value=other_min_lr,
            peak_value=other_peak_lr,
            warmup_steps=warmup_steps,
            decay_steps=max(1, total_train_steps - warmup_steps),
            end_value=other_min_lr,
        ),
    }

    def _path_names(path):
        out = []
        for k in path:
            if isinstance(k, GetAttrKey):
                out.append(k.name)
            elif isinstance(k, SequenceKey):
                out.append(str(k.idx))
            elif isinstance(k, DictKey):
                out.append(str(k.key))
            else:
                out.append(str(k))
        return out

    def label_fn(path, leaf):
        # Top-level fields in your GPT pytree: embed, blocks, lm_head
        names = _path_names(path)
        top = names[0] if names else ""
        if top == "embed":
            return "embed"
        if top == "lm_head":
            return "lm_head"
        return "other"

    param_labels = jax.tree_util.tree_map_with_path(label_fn, params)

    def make_adamw(lr_schedule, weight_decay=0.0):
        return optax.adamw(
            learning_rate=lr_schedule,
            b1=b1,
            b2=b2,
            weight_decay=weight_decay,
            mu_dtype=jnp.float32
        )

    def log_step_and_lr(label, schedule):
        def init_fn(params):
            return {"count": jnp.array(0, dtype=jnp.int32)}
        def update_fn(updates, state, params=None):
            step = state["count"]
            lr = schedule(step)
            jax.debug.print("[{label}] step {step} lr {lr}", label=label, step=step, lr=lr)
            return updates, {"count": step + 1}
        return optax.GradientTransformation(init_fn, update_fn)

    tx = optax.multi_transform(
        {
            "embed": make_adamw(schedules["embed"]),
            "lm_head": make_adamw(schedules["lm_head"]),
            "other": make_adamw(schedules["other"], weight_decay=0.001),
        },
        param_labels,
    )
    return tx

Initialization of the optimier w/wo gradient accumulation:

# In case gradient accum is enabled. Default value for grad_acum_steps=4
optim = optax.MultiSteps(optim, every_k_schedule=grad_accum_steps)
optim_state = optim.init(model)

Here is how the optimizer state is updated inside the training loop:

@partial(jax.jit, static_argnames=("optim",), donate_argnums=(0, 4))
def train_step_accum(params, x_batch, y_batch, freqs, optim_state, optim):
    (loss, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(
        params, x_batch, y_batch, freqs
    )
    _, optim_state = optim.update(grads, optim_state, params)
    return params, loss, optim_state


@partial(jax.jit, static_argnames=("optim",), donate_argnums=(0, 4))
def train_step(params, x_batch, y_batch, freqs, optim_state, optim):
    (loss, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(
        params, x_batch, y_batch, freqs
    )
    updates, optim_state = optim.update(grads, optim_state, params)
    updated_params = optax.apply_updates(params, updates)
    return updated_params, loss, optim_state


for step in range(total_train_steps):
    ...
    for micro_step in range(grad_accum_steps):
        x, y = get_next_batch(starts, ends, bsz, seqlen, tokens, data_sharding)
        if micro_step < grad_accum_steps - 1:
            model, loss, optim_state = train_step_accum(model, x, y, freqs, optim_state, optim)
            train_step_loss += loss
        else:
            model, loss, optim_state = train_step(model, x, y, freqs, optim_state, optim)
            train_step_loss += loss
    avg_train_loss = train_step_loss / grad_accum_steps
    ...

Can someone please explain the reason behind this behavior? I would have expected a more stable training and faster convergence with gradient accumulation, but right now it is totally broken

[vroulet]

Thanks for reporting the bug.
This may be issues in dtypes.
Could you check if the dtypes are the same inside the optimizer in each case?
What precision are you using?

[AakashKumarNain]

This may be issues in dtypes.

How come? I do not see any reason why would they be different.

What precision are you using?

bf16 for model weights, fp32 for logits, momentum accumulation also in fp32

Another thing to note, If I just use train_step(...) and not a combination of train_step_accum(...) and train_step(...) (as shown in the above code), the loss converges towards the same value but a bit slowly. But in that case, I won't be able to increase the throughput of my training run and squeeze every ounce of my accelerator. What do you suggest that I should do here?

[vroulet]

Multisteps can be converting dtypes in its logic and I was thinking it could be the culprit (so it would be a bug on our side)
Can you check that the dtypes in the state match in both cases?

Just a few more questions looking back at your code:
Why do you differentiate train_step and train_step_accum? You a priori do not need this distinction, although I don't think there is anything wrong doing it. Can you quickly test by using just train_step? I.e., removing the if logic you have in your code.

Thanks again for reporting this!

[AakashKumarNain]

Multisteps can be converting dtypes in its logic and I was thinking it could be the culprit (so it would be a bug on our side)
Can you check that the dtypes in the state match in both cases?

Yes, I will report that

Why do you differentiate train_step and train_step_accum? You a priori do not need this distinction, although I don't think there is anything wrong doing it.

I make two branches so that during accumulation we can avoid the no-op update and writing all the weights to HBM at each accumulation step.

[AakashKumarNain]

@vroulet not sure if this is what you wanted me to check, but got a toy example (this is generated btw) to check it.

import jax
import jax.numpy as jnp
import optax

def predict(params, X):
    # X: [N, D]
    # w: [D], b: []
    w, b = params["w"], params["b"]
    return X @ w + b  # [N]

def mse_loss(params, X, y):
    y_pred = predict(params, X)
    return jnp.mean((y_pred - y) ** 2)

def make_train_step(optimizer):
    @jax.jit
    def train_step(params, opt_state, X, y):
        loss, grads = jax.value_and_grad(mse_loss)(params, X, y)
        updates, opt_state =  optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss
    return train_step

def make_train_step_accum(optimizer):
    @jax.jit
    def train_step_accum(params, opt_state, X, y):
        loss, grads = jax.value_and_grad(mse_loss)(params, X, y)
        updates, opt_state =  optimizer.update(grads, opt_state, params)
        return params, opt_state, loss
    return train_step_accum

def print_ms_state(tag, params_before, params_after, opt_state, grads=None):
    print(f"\n[{tag}]")
    try:
        print("  params_before dtypes -> w:", params_before["w"].dtype, " b:", params_before["b"].dtype)
        print("  params_after  dtypes -> w:", params_after["w"].dtype, " b:", params_after["b"].dtype)
    except Exception as e:
        print("  params dtype print error:", e)

    try:
        mini = int(opt_state.mini_step)
        grad_step = int(opt_state.gradient_step)
        print("  MultiStepsState detected")
        print("  mini_step:", mini, " gradient_step:", grad_step)
        try:
            aw = opt_state.acc_grads["w"]
            ab = opt_state.acc_grads["b"]
            print("  acc_grads dtypes    -> w:", aw.dtype, " b:", ab.dtype)
        except Exception:
            leaves, _ = jax.tree_util.tree_flatten(opt_state.acc_grads)
            summary = []
            for x in leaves[:4]:
                d = getattr(x, "dtype", None)
                summary.append(str(d) if d is not None else type(x).__name__)
            print("  acc_grads leaf dtype sample:", summary)
    except Exception:
        print("  Non-MultiSteps optimizer state")
        leaves, _ = jax.tree_util.tree_flatten(opt_state)
        summary = []
        for x in leaves[:6]:
            d = getattr(x, "dtype", None)
            summary.append(str(d) if d is not None else type(x).__name__)
        print("  opt_state leaf dtype sample:", summary)
        if grads is not None:
            try:
                print("  grads dtypes        -> w:", grads["w"].dtype, " b:", grads["b"].dtype)
            except Exception:
                gleaves, _ = jax.tree_util.tree_flatten(grads)
                gsummary = []
                for x in gleaves[:4]:
                    d = getattr(x, "dtype", None)
                    gsummary.append(str(d) if d is not None else type(x).__name__)
                print("  grads leaf dtype sample:", gsummary)


def main():
    key = jax.random.PRNGKey(0)
    N, D = 1024, 3
    key, k1, k2, k3 = jax.random.split(key, 4)
    X = jax.random.normal(k1, (N, D)).astype(jnp.bfloat16)
    true_w = jnp.array([2.0, -1.0, 0.5]).astype(jnp.bfloat16)
    true_b = 0.7
    noise = 0.1 * jax.random.normal(k2, (N,)).astype(jnp.bfloat16)
    y = X @ true_w + true_b + noise

    # Initialize parameters
    params = {
        "w": 0.01 * jax.random.normal(k3, (D,)).astype(jnp.bfloat16),
        "b": jnp.array(0.0).astype(jnp.bfloat16),
    }

    # Optimizer
    grad_accum_steps = 4
    optimizer = optax.adamw(learning_rate=1e-2, b1=0.9, b2=0.95, mu_dtype=jnp.float32)
    optimizer = optax.MultiSteps(optimizer, every_k_schedule=grad_accum_steps)
    opt_state = optimizer.init(params)
    train_step = make_train_step(optimizer)
    train_step_accum = make_train_step_accum(optimizer)

    dbg_grads0 = jax.grad(mse_loss)(params, X, y)
    print_ms_state("after_init", params, params, opt_state, grads=dbg_grads0)

    for step in range(2000):
        for micro_step in range(grad_accum_steps):
            params_before = params
            if micro_step < grad_accum_steps - 1:
                params, opt_state, loss = train_step_accum(params, opt_state, X, y)
            else:
                params, opt_state, loss = train_step(params, opt_state, X, y)
            if step == 0 and micro_step in (0, 1, grad_accum_steps - 1):
                print_ms_state(f"after_micro_{micro_step}", params_before, params, opt_state)

    # === (Optional) after training completes ===
    print_ms_state("after_training", params, params, opt_state)

if __name__ == "__main__":
    main()

It outputs the following:

[after_init]
  params_before dtypes -> w: bfloat16  b: bfloat16
  params_after  dtypes -> w: bfloat16  b: bfloat16
  MultiStepsState detected
  mini_step: 0  gradient_step: 0
  acc_grads dtypes    -> w: bfloat16  b: bfloat16

[after_micro_0]
  params_before dtypes -> w: bfloat16  b: bfloat16
  params_after  dtypes -> w: bfloat16  b: bfloat16
  MultiStepsState detected
  mini_step: 1  gradient_step: 0
  acc_grads dtypes    -> w: float32  b: float32

[after_micro_1]
  params_before dtypes -> w: bfloat16  b: bfloat16
  params_after  dtypes -> w: bfloat16  b: bfloat16
  MultiStepsState detected
  mini_step: 2  gradient_step: 0
  acc_grads dtypes    -> w: float32  b: float32

[after_micro_3]
  params_before dtypes -> w: bfloat16  b: bfloat16
  params_after  dtypes -> w: bfloat16  b: bfloat16
  MultiStepsState detected
  mini_step: 0  gradient_step: 1
  acc_grads dtypes    -> w: float32  b: float32

[after_training]
  params_before dtypes -> w: bfloat16  b: bfloat16
  params_after  dtypes -> w: bfloat16  b: bfloat16
  MultiStepsState detected
  mini_step: 0  gradient_step: 2000
  acc_grads dtypes    -> w: float32  b: float32

Do let me know if this is what you wanted from my side.

[AakashKumarNain]

@vroulet any updates on this?

[rdyro]

The dtypes look correct (accumulation happens in float32).

Can you clarify again? The training gets stuck in local minima if the accumulation is configured to use the same mini-batch? (as in gradient accumulation is done over micro-batches within the same mini-batch?)

I'm trying to understand if the issue is numerical differences in gradient accumulation or if it's learning dynamics. If it's the latter, there's maybe no reason to expect the accumulated version to reach the same minimum? Does this sound right?

[AakashKumarNain]

Sorry for the late response @rdyro Let me clarify again. If I use optax.MultiStep and a single function say train_step as shown in the code, then it works fine. Calling multisteps within the training loop does achieve gradient accumulation, but it does not have give any benefit to squeeze the accelerator. Now, there are two ways I can think of to increase the throughput there:

1. Segregate accumulation and update steps

@partial(jax.jit, static_argnames=("optim",), donate_argnums=(0, 4))
def train_step_accum(params, x_batch, y_batch, freqs, optim_state, optim):
    (loss, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(
        params, x_batch, y_batch, freqs
    )
    _, optim_state = optim.update(grads, optim_state, params)
    return params, loss, optim_state


@partial(jax.jit, static_argnames=("optim",), donate_argnums=(0, 4))
def train_step(params, x_batch, y_batch, freqs, optim_state, optim):
    (loss, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(
        params, x_batch, y_batch, freqs
    )
    updates, optim_state = optim.update(grads, optim_state, params)
    updated_params = optax.apply_updates(params, updates)
    return updated_params, loss, optim_state


for step in range(total_train_steps):
    ...
    for micro_step in range(grad_accum_steps):
        x, y = get_next_batch(starts, ends, bsz, seqlen, tokens, data_sharding)
        if micro_step < grad_accum_steps - 1:
            model, loss, optim_state = train_step_accum(model, x, y, freqs, optim_state, optim)
            train_step_loss += loss
        else:
            model, loss, optim_state = train_step(model, x, y, freqs, optim_state, optim)
            train_step_loss += loss
    avg_train_loss = train_step_loss / grad_accum_steps
    ...

This code is correct, but the convergence is very poor compared to calling only train_step in the loop.

2. This method is even better, but MultiStep is casting the params and is causing a conflict.

ef make_train_step(optimizer, grad_accum_steps: int):
    def accum_step(carry, micro_batch):
        """Single micro-step: compute grads and accumulate via MultiSteps."""
        params, opt_state = carry
        X_micro, y_micro = micro_batch
        loss, grads = jax.value_and_grad(mse_loss)(params, X_micro, y_micro)
        # MultiSteps accumulates internally; on the final micro-step it emits
        # the real update, otherwise it returns a zero-tree (no-op on params).
        updates, new_opt_state = optimizer.update(grads, opt_state, params)
        new_params = optax.apply_updates(params, updates)
        return (new_params, new_opt_state), loss

    @jax.jit
    def train_step(params, opt_state, X, y):
        # Split batch into grad_accum_steps micro-batches along axis 0.
        # X: [N, D] -> [grad_accum_steps, N // grad_accum_steps, D]
        # y: [N]    -> [grad_accum_steps, N // grad_accum_steps]
        X_splits = X.reshape(grad_accum_steps, X.shape[0] // grad_accum_steps, *X.shape[1:])
        y_splits = y.reshape(grad_accum_steps, y.shape[0] // grad_accum_steps)

        (new_params, new_opt_state), losses = jax.lax.scan(
            accum_step,
            init=(params, opt_state),
            xs=(X_splits, y_splits),
        )
        return new_params, new_opt_state, losses.mean()

    return train_step

It throws this error:

Traceback (most recent call last):
  File "/home/ubuntu/optax_bug.py", line 221, in <module>
    main()
  File "/home/ubuntu/optax_bug.py", line 212, in main
    params, opt_state, loss = train_step(params, opt_state, X, y)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/optax_bug.py", line 176, in train_step
    (new_params, new_opt_state), losses = jax.lax.scan(
                                          ^^^^^^^^^^^^^
TypeError: scan body function carry input and carry output must have equal types, but they differ:

  * the input carry component carry[1].acc_grads['b'] has type bfloat16[] but the corresponding output carry component has type float32[], so the dtypes do not match;
  * the input carry component carry[1].acc_grads['w'] has type bfloat16[3] but the corresponding output carry component has type float32[3], so the dtypes do not match.

Revise the function so that all output types match the corresponding input types.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Let me know if you need any more info

[rdyro]

I created a PR for fixing this dtype instability (so scan should work). Can you give it a try?

If this scan approach still gives you bad performance, please let me know, let's see if we can fix that.