Allocate the gradient accumulation carry in config.grad_dtype

[NuojCheng]

Description

Allocate the gradient accumulation scan carry in config.grad_dtype instead of always float32.

FIXES: b/505094303

Why does turning on gradient accumulation increase HBM?

Short answer: at a matched global batch size it does not — GA is a large win. But enabling GA does add a fixed cost that is independent of the number of microbatches, and today half of that fixed cost is avoidable.

All numbers below are from AOT compilation (maxtext.trainers.pre_train.train_compile, compile_topology=v5e-8) of a 4-layer, base_emb_dim=512, vocab_size=32000, max_target_length=512 model, with --xla_dump_to HLO dumps.

1. At matched global batch, GA saves memory (expected, working as intended)

config	global batch	peak HBM
gradient_accumulation_steps=1, per_device_batch_size=4	4	361.28 MiB
gradient_accumulation_steps=4, per_device_batch_size=1	4	148.71 MiB

Activations shrink by 1/K, which is the whole point.

2. But GA adds a constant overhead, independent of K

config	peak HBM
ga=1, pdbs=1	113.82 MiB
ga=2, pdbs=1	165.88 MiB
ga=4, pdbs=1	165.93 MiB

ga=2 and ga=4 cost the same, so this is a one-time cost of enabling GA, not a per-microbatch cost.

The cause is visible in the optimized HLO. Without GA, XLA interleaves the backward pass with the optimizer update, so only one tensor's gradient is live at a time and no full-size gradient tree is ever materialized. With GA, jax.lax.scan must keep the entire accumulated gradient tree live for the duration of the loop. That is inherent to accumulation.

3. Half of that overhead was avoidable

The accumulator was created with jnp.zeros_like(ga_params), i.e. always float32, regardless of config.grad_dtype. grad_dtype was applied only after the loop, by the trainer:

# src/maxtext/trainers/pre_train/train.py
raw_grads = jax.tree_util.tree_map(
    lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x,
    raw_grads,
)

On the non-GA path that cast happens straight out of the backward pass, so a grad_dtype=bfloat16 recipe never holds a float32 gradient tree. On the GA path the float32 tree is held for the whole scan and only narrowed at the very end — so GA used 2x the gradient memory of the equivalent non-GA step for the same recipe.

The while carry of the GA scan in the pre-fix HLO shows it directly — a full float32 gradient tree alongside the hoisted bf16 parameter copies:

%while.246 = (s32[], f32[512], f32[64,4,2048], f32[64,4,2048], f32[2048,4,64],
              f32[512,4], f32[512,4], f32[64,4,8,64], f32[8,4,64,64], f32[64,4,8,64],
              f32[64,4,8,64], f32[64,32000], f32[32000,64], ...,
              bf16[32000,64], bf16[64,4,2048], ..., bf16[64,32000], ...) while(...)
              metadata={op_name="jit(train_step)/while"}

The change

• Allocate init_grad in config.grad_dtype rather than always float32.
• Accumulate into the carry's dtype (y + x.astype(y.dtype)) — required anyway, since adding a float32 microbatch gradient to a narrower accumulator would promote the result and change the carry type, which scan rejects.
• Cast back after the post-loop division. divisor is a float32 scalar, so arr / divisor promotes every leaf back to float32 and would materialize a full float32 gradient tree immediately after the loop, undoing the saving.

This aligns the GA path with the non-GA path, and with the existing shard_optimizer_over_data (ZeRO-1) path, which already casts params to bf16 before the loop and therefore already accumulated in bf16.

Results

grad_dtype=bfloat16, same model as above:

metric	before	after	delta
temp_size_in_bytes (ga=2)	135,877,632	99,719,168	-26.6%
peak HBM (ga=2)	165.88 MiB	148.63 MiB	-10.4%
peak HBM (ga=4)	165.93 MiB	148.71 MiB	-10.4%

The saving is 2 bytes x (parameters per device shard), so it scales with model size.

After the fix the scan carry is bf16 throughout — f32[32000,64], f32[64,32000], f32[64,4,2048], f32[2048,4,64], f32[64,4,8,64] and f32[8,4,64,64] are all gone.

No change at the default

grad_dtype defaults to float32, so the default path is untouched. AOT memory statistics are byte-identical before and after for every grad_dtype=float32 config tested:

config	before	after
ga=1, grad_dtype=float32	42,166,784	42,166,784
ga=2, grad_dtype=float32	136,329,216	136,329,216
ga=1, grad_dtype=bfloat16	43,263,488	43,263,488

What this does not fix

The remaining GA overhead (~36 MiB on the test model) is:

• the bf16 gradient accumulator itself — inherent to accumulation;
• a bf16 copy of the parameter tree that XLA hoists out of the loop by loop-invariant code motion (the weight_dtype=float32 -> dtype=bfloat16 cast, which without GA is recomputed per step and never held). Eliminating that would mean casting ga_params before the loop the way shard_optimizer_over_data does, which changes gradient numerics, so it is left out of this PR.

Also checked and ruled out: ga_params in the scan carry. Hoisting it into a closure produces byte-identical output — XLA already performs that motion.

Tests

1. Existing GA integration tests (v5e-4, 2 passed):

python3 -m pytest tests/integration/gradient_accumulation_test.py -q
# test_grad_accumulate_same_loss, test_sft_grad_accumulate_same_loss

These assert that GA and non-GA produce the same loss.

2. Numerical parity between grad_dtype=float32 and grad_dtype=bfloat16 under GA. 8-step training run on TPU, identical loss to 3 decimal places at every step (10.865 / 10.851 / 10.839 / 10.828 / 10.819 / 10.813 / 10.810 / 10.808):

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  base_emb_dim=256 base_num_query_heads=4 base_num_kv_heads=4 base_mlp_dim=1024 \
  base_num_decoder_layers=4 head_dim=64 max_target_length=256 vocab_size=32000 \
  decoder_block=simple dataset_type=synthetic enable_checkpointing=false \
  enable_goodput_recording=false gradient_clipping_threshold=0 \
  per_device_batch_size=1 gradient_accumulation_steps=4 grad_dtype=<float32|bfloat16> \
  steps=8 run_name=r

3. AOT memory comparison (the tables above). Reproduce with:

XLA_FLAGS="--xla_dump_to=/tmp/hlo" \
python3 -m maxtext.trainers.pre_train.train_compile src/maxtext/configs/base.yml \
  compile_topology=v5e-8 compile_topology_num_slices=1 \
  base_emb_dim=512 base_num_query_heads=8 base_num_kv_heads=8 base_mlp_dim=2048 \
  base_num_decoder_layers=4 head_dim=64 vocab_size=32000 max_target_length=512 \
  gradient_accumulation_steps=<K> per_device_batch_size=<B> \
  dtype=bfloat16 grad_dtype=<float32|bfloat16> enable_checkpointing=false \
  dataset_type=synthetic run_name=ga

Peak HBM is the Total bytes line of /tmp/hlo/module_*jit_train_step*-memory-usage-report.txt.

4. tests/unit/gradient_accumulation_nnx_test.py — 6 passed (JAX_PLATFORMS=cpu). The stub _Cfg in that test enumerates the config fields the function reads, so it gains grad_dtype: jnp.dtype = jnp.float32.

5. Lint: pyink and pylint clean. The one remaining pylint warning, too-many-positional-arguments, is pre-existing and the function signature is unchanged.

No large-scale workload link: the change is dtype-only and gated on a config default that is unchanged, and the effect is fully characterized by the AOT memory statistics above.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request optimizes gradient accumulation memory usage by storing the accumulator in config.grad_dtype instead of always using float32. It introduces a downcasting helper, initializes the accumulator with this dtype, and ensures gradients are cast back to the accumulator's dtype after division. Feedback suggests performing the accumulation addition in float32 before casting to the accumulator's dtype, rather than casting the microbatch gradient beforehand, to prevent numerical precision loss and underflow during accumulation.

[gemini-code-assist]

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"]
	)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!