Skip to content

Allocate the gradient accumulation carry in config.grad_dtype #5039

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
NuojCheng wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Allocate the gradient accumulation carry in config.grad_dtype #5039

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions src/maxtext/utils/gradient_accumulation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,11 @@ def convert_to_bf16(param):
ga_params = params

ga_params = jax.tree.map(_maybe_shard_with_name, ga_params, ga_params_shardings)

def _to_grad_dtype(arr):
"""Downcast float32 leaves to config.grad_dtype."""
return arr.astype(config.grad_dtype) if arr.dtype == jnp.float32 else arr

if is_nnx:
grad_func = nnx.value_and_grad(_loss_fn, argnums=0, has_aux=True)
else:
Expand Down Expand Up @@ -133,7 +138,9 @@ def accumulate_gradient(acc_grad_and_loss, data):
acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"]
acc_grad_and_loss["indexer_loss"] += aux["indexer_loss"]
acc_grad_and_loss["mtp_loss"] += aux["mtp_loss"]
acc_grad_and_loss["grad"] = jax.tree_util.tree_map(lambda x, y: x + y, cur_batch_gradient, acc_grad_and_loss["grad"])
acc_grad_and_loss["grad"] = jax.tree_util.tree_map(
lambda x, y: y + x.astype(y.dtype), cur_batch_gradient, acc_grad_and_loss["grad"]
)
Comment on lines +141 to +143

@gemini-code-assist gemini-code-assist Bot Aug 28, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Performing the addition in float32 and then casting the result back to y.dtype is numerically more precise than casting x to y.dtype before the addition.

When y is bfloat16 and x is float32, doing y + x.astype(y.dtype) performs the addition in bfloat16 (7-bit mantissa), which can lead to significant rounding errors or underflow (swamping) over multiple accumulation steps. By doing (y + x).astype(y.dtype), the addition is performed in float32 precision (23-bit mantissa) and only the final sum is rounded back to bfloat16. This preserves the carry's dtype for jax.lax.scan while maintaining much higher numerical precision.

Suggested change
acc_grad_and_loss["grad"] = jax.tree_util.tree_map(
lambda x, y: y + x.astype(y.dtype), cur_batch_gradient, acc_grad_and_loss["grad"]
)
acc_grad_and_loss["grad"] = jax.tree_util.tree_map(
lambda x, y: (y + x).astype(y.dtype), cur_batch_gradient, acc_grad_and_loss["grad"]
)

acc_grad_and_loss["total_weights"] += aux["total_weights"]
return acc_grad_and_loss, aux

Expand All @@ -145,7 +152,7 @@ def reshape_to_microbatch_accumulations(batch_arr):
return jnp.swapaxes(reshaped_batch_arr, 0, 1)

data = jax.tree_util.tree_map(reshape_to_microbatch_accumulations, data)
init_grad = jax.tree_util.tree_map(jnp.zeros_like, ga_params)
init_grad = jax.tree_util.tree_map(lambda p: _to_grad_dtype(jnp.zeros_like(p)), ga_params)
init_grad = jax.tree.map(_maybe_shard_with_name, init_grad, grad_shardings)
init_grad_and_loss = {
"loss": 0.0, # accumulates xent_sum across microbatches
Expand Down Expand Up @@ -183,7 +190,7 @@ def reshape_to_microbatch_accumulations(batch_arr):
config.gradient_accumulation_steps if getattr(config, "use_tunix_gradient_accumulation", False) else denominator
)
raw_grads = jax.tree_util.tree_map(
lambda arr: jnp.where(has_weights, arr / divisor, jnp.zeros_like(arr)),
lambda arr: jnp.where(has_weights, arr / divisor, jnp.zeros_like(arr)).astype(arr.dtype),
raw_grads,
)
aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr
Expand Down
1 change: 1 addition & 0 deletions tests/unit/gradient_accumulation_nnx_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ class _Cfg:
debug_sharding: bool = False
training_objective: str = "causal_lm"
use_tunix_gradient_accumulation: bool = False
grad_dtype: jnp.dtype = jnp.float32


class _TinyNNX(nnx.Module):
Expand Down
Loading