DSv3 improvement by keeping the unreduced gradient tag under ZeRO-1 + explicit sharding, including through the MoE shard_map

[NuojCheng]

Description

Running DeepSeek + ZeRO-1 + gradient accumulation under shard_mode=explicit drops the unreduced tag on some parameter gradients, so those gradients get all-reduced once per microbatch inside the accumulation loop instead of once per step after it.

smoke_train ici_data_parallelism=-1 ici_fsdp_parallelism=1 \
  shard_optimizer_over_data=true model_name=deepseek3-tiny \
  shard_mode=explicit gradient_accumulation_steps=4

The failure is silent — gradients stay numerically identical, because JAX happily converts a fully-reduced value back to a partial one. The only symptom is wasted collective traffic.

Background: how the tag works

Under explicit sharding, gradient accumulation marks the parameters it carries through the scan as reduced over the data axis. A reduced input means "already summed over that axis", so its cotangent comes out unreduced, and the all-reduce is deferred until something reshards away from unreduced — which happens once, after the loop. Any code path that rebuilds a PartitionSpec from the logical axis rules produces an untagged spec, and resharding a parameter onto it silently drops that parameter off the deferred path.

Fix 1 — restore the tag under ZeRO-1 + explicit sharding

A tagged PartitionSpec rejects indexing, slicing, unpacking and iteration; only .partitions reads through it. Several sharding helpers were doing exactly that and so were either raising or quietly dropping the tag.

src/maxtext/utils/sharding.py:

• remove_size_one_mesh_axis and adjust_pspec_for_indivisible_shapes now read .partitions instead of indexing the spec directly.
• get_mesh_axes_used_by_tensor_spec reads through .partitions.
• New batch_mesh_axes helper.
• truncate_out_sharding needed the same treatment; main has since landed an equivalent _truncate_pspec helper independently, so this PR now takes upstream's version and only keeps the accompanying tests.

src/maxtext/utils/gradient_accumulation.py:

• Gating via data_is_only_batch_axis(config, params_shardings), plus update_sharding_for_reduced / update_sharding_for_unreduced and a post-scan jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings).

No changes to train.py.

Fix 2 — keep the tag through the MoE shard_map

Fix 1 alone turned out to be a wall-clock regression on MoE models — 2–4% slower. Root cause: moe.py:sparse_matmul reshards the expert kernels onto a pspec rebuilt from the logical axis rules, which carries no tag. jax.shard_map makes every mesh axis manual, so the kernel cotangent was psum'd over data inside the body — once per microbatch, per layer — and the accumulated gradient was reduced again on the way out of the scan. Double work, correct numerics.

• New carry_reduced_axes(spec, value) in sharding.py, which copies the reduced mesh axes of a value's own spec onto a rebuilt spec. It declines when the tag would overlap the partitions (which PartitionSpec rejects) and when there is no spec to read, so it is a no-op outside explicit sharding.
• sparse_matmul applies it to w0/w1/wo and their biases right after maybe_aqt_partition.
• get_wi_gmm_params / get_wo_gmm_params read .partitions, since a tagged spec refuses indexing.

Verified in isolation that jax.shard_map accepts a reduced-tagged in_spec and then emits zero in-body all-reduce.

Rebase note

Rebased onto main at d2d73155b (30 commits). Two conflicts, both from upstream converging on the same ideas:

• truncate_out_sharding — main landed _truncate_pspec, which carries reduced/unreduced across truncation exactly as this branch did. Resolved in favour of upstream's factoring; only the tests from this branch remain.
• sharding_nnx_test.py — main added equivalent coverage for tag-preserving truncation. Kept upstream's two tests, dropped the redundant one from this branch, kept the new TaggedPartitionSpecTest class (which covers get_mesh_axes_used_by_tensor_spec, remove_size_one_mesh_axis and adjust_pspec_for_indivisible_shapes).

moe.py was heavily refactored upstream (~350 lines, including the new get_routed_moe_shardings) but merged cleanly: the carry_reduced_axes calls still sit between maybe_aqt_partition and the shard_map, whose in_specs still consume those pspecs. All measurements below were re-run after the rebase and are unchanged.

One thing did change upstream: the ds_dp2fsdp4_z1_ga4 sweep case previously died with a ShardingTypeError deep in sharding. main now rejects it up front at config validation with a clear message — `shard_optimizer_over_data` (Zero-1) cannot be combined with FSDP — so that gap is closed and is no longer an open item for this PR.

Tests

New coverage

tests/integration/gradient_accumulation_test.py gains two tests.

test_deepseek_zero1_reduces_gradients_once_per_step (AOT, v5e-8 topology) asserts no cross-replica all-reduce larger than 8 elements sits inside the accumulation loop — the loss and token count are legitimately reduced per microbatch and are scalars, so the bound separates those from parameter gradients. A vacuity guard checks that one does exist outside the loop. Supporting helpers live in tests/utils/hlo_test_utils.py (split_by_entry_loop, cross_replica_all_reduce_sizes).

test_deepseek_zero1_ga_scale_memory_and_step_time actually trains a ~1.29B DeepSeek MoE on real chips, twice — once as shipped, once with data_is_only_batch_axis patched to False to reproduce pre-fix behaviour — and asserts on both step time and peak HBM. HLO placement says nothing about what the deferred reduction is worth; this catches a change that keeps the tag but reduces gradients per microbatch somewhere downstream.

untagged: 0.382786 s/step over 8 steps, peak HBM so far 9.9513 GiB
tagged:   0.316346 s/step over 8 steps, peak HBM so far 9.9518 GiB

Asserts speedup > 1.05x (measured 1.21x), memory ratio < 1.02 (measured 1.00005), and a >4 GiB peak floor so the config cannot be silently shrunk past the point where the thresholds mean anything. Skips below 4 devices or 24 GiB/chip.

One design note: the runs are ordered untagged first, on purpose. peak_bytes_in_use is a high-water mark the TPU allocator never resets, so the second reading is max(both); asserting "second <= first" is then exactly "tagging did not raise the peak", with no way for a regression to hide behind the mark. An earlier subprocess-per-variant design was abandoned because the parent pytest process initialises the TPU when it checks the device count, which then locks the children out.

How to reproduce

# Integration tests (needs >=4 chips with >=24 GiB HBM each)
python -m pytest tests/integration/gradient_accumulation_test.py -q -s

# Unit tests (CPU)
JAX_PLATFORMS=cpu python -m pytest \
  tests/unit/gradient_accumulation_nnx_test.py tests/unit/sharding_compare_test.py \
  tests/unit/sharding_desc_test.py tests/unit/sharding_nnx_test.py tests/unit/sharding_test.py -q

# MoE unit tests
python -m pytest tests/unit/moe_test.py -q

Results

• Full tests/integration/gradient_accumulation_test.py: 4 passed (67 s) on a 2x2 v6e host. Scale numbers reproduce identically when sharing a process with the other tests, confirming no peak-HBM carryover between runs.
• Sharding + gradient-accumulation unit tests (5 files, CPU): 102 passed, 1 skipped, 4 subtests.
• tests/unit/moe_test.py: 32 passed, 42 skipped, 1 failed. The failure is test_gmm_grad_equivalence_tokamax_v2_fp8_dynamic_ep4 (NaN relative-norm) and is pre-existing — it fails identically with moe.py stashed back to main.
• AOT sharding sweep, 7 configs: 6 pass. ds_dp2fsdp4_z1_ga4 is now rejected by main's own config validation as an unsupported combination (see the rebase note above).
• Numerics: losses match to ~1e-5 relative over 6 steps between tagged and untagged.
• pyink --pyink-indentation=2 --line-length=122 clean on every file this PR touches except moe.py, which is already not pyink-clean on main in regions this PR does not touch; reformatting them would bury the diff. pylint --rcfile=pylintrc scores 10.00/10 on the test file, and the two E1128 in moe.py are pre-existing on main.

Measurements

All on a single 2x2 v6e host (4 chips, 31.24 GiB/chip), bf16, synthetic data, median of 8 steady-state steps after 4 warmup steps.

The MoE fix flips the sign of Fix 1

config	after Fix 1 only	after Fix 2
8L / 8E, emb 1024	0.1084 s vs 0.1062 s untagged	0.0958 s vs 0.1062 s	2.0% slower -> 9.8% faster
24L / 16E, ~1.29B	0.3994 s vs 0.3827 s untagged	0.3164 s vs 0.3827 s	4.3% slower -> 17.3% faster

Collective placement at 24L/16E: tagged has 2 in-loop all-reduce ops totalling 0.0 MiB plus 2809 MiB once outside; untagged has 34 MiB in-loop (x23 layers x4 microbatches ~ 9 GiB) plus 509 MiB outside.

Explicit vs auto sharding

The tags only exist on an all-Explicit mesh, so this fix is a no-op under shard_mode=auto. What changes is where explicit mode lands relative to auto:

shard_mode	median step	peak HBM	vs auto
explicit (with fix)	0.3163 s	9.9518 GiB	1.15x faster
auto	0.3650 s	9.9481 GiB	—
explicit (pre-fix)	0.3827 s	9.9513 GiB	0.95x (slower)

Explicit went from 4.9% slower than auto to 13.3% faster. Peak HBM is flat across all three (spread 0.04%), so this is pure traffic, not a memory/compute tradeoff.

On stock deepseek3-tiny (pdbs 1, 0.35 GiB) explicit and auto are indistinguishable — 0.5871 s vs 0.5864 s. The model is far too small for collective traffic to register against fixed per-step overhead, which is why the original bug was invisible as a perf problem on the config it was reported against.

When the fix pays off

One factor at a time from the 24L/16E/pdbs2/GA4 base:

axis	value	untagged	tagged	speedup
GA steps	2	0.2117	0.1911	1.108x
	4	0.3827	0.3164	1.210x
	8	0.7253	0.5670	1.279x
tokens/microbatch	pdbs 1	0.3268	0.2589	1.262x
	pdbs 2	0.3828	0.3163	1.210x
	pdbs 4	0.4995	0.4373	1.142x
experts (12L)	8	0.1489	0.1323	1.125x
	16	0.1993	0.1669	1.194x
	32	0.2958	0.2367	1.250x
depth (16E)	12L	0.1995	0.1671	1.194x
	24L	0.3827	0.3163	1.210x
	36L	0.5608	0.4518	1.241x

The eliminated work is exactly (N-1) redundant all-reduces, so time is N*(c+a) before and N*c + a after, giving a speedup of (c+a)/(c+a/N).

• Microbatch count is the strongest lever, but saturates. Absolute time saved scales linearly in N — 0.021 / 0.066 / 0.158 s at N = 2/4/8, i.e. 1 : 3.2 : 7.7 against a predicted 1 : 3 : 7. Total step time scales with N too, so the ratio approaches a ceiling. Fitting T = o + N*c + {N*a or a} to the three points gives a ~ 22 ms, c ~ 63 ms, o ~ 41 ms fixed overhead, so the ceiling is ~1.35x for this shape. GA=8 is already at 1.28.
• Bigger model is roughly neutral. Depth adds gradient bytes and FLOPs together, so a/c barely moves across 3x the layers. The mild 1.19 -> 1.24 drift is fixed per-step overhead being diluted, not a scaling effect.
• Sparsity is the model-shape lever that matters. 8 -> 32 experts moves 1.13x -> 1.25x, because expert weights add gradient bytes ~linearly while FLOPs only grow with routing overhead (top-k fixed at 2). Absolute saved time scaled 1 : 1.95 : 3.56 with expert count while total step time only scaled 1 : 1.34 : 1.99. A real DeepSeek-V3 shape (256 experts, top-8) sits much further right on this curve than anything that fits on four chips.
• Smaller microbatches help, same bytes-per-FLOP reason: fewer tokens per microbatch means the same all-reduce amortizes over less compute.

Summary: the gain is (1 - 1/N) * a/c where a/c is gradient-bytes-per-microbatch-FLOP. Largest with many microbatches, few tokens each, and a sparse MoE — precisely the regime GA exists to serve.

Not measured: this is 4-way DP on one host. More replicas make each all-reduce costlier, which should push a/c and therefore the win higher, but there is no data on it here. Treat multi-host gains as an expectation, not a number.

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. (The reduced/unreduced tag semantics are subtle and fail silently, so carry_reduced_axes, every .partitions read, and the untagged-first ordering in the scale test each carry a comment explaining why.)
• I have run end-to-end tests tests and provided workload links above if applicable. (Run on a local 2x2 v6e host rather than an xpk workload, so there is no workload link; full commands and results are in the Tests section above.)
• 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. (No doc changes needed — this is an internal sharding correctness fix with no user-facing config or API surface.)

[gemini-code-assist]

Code Review

This pull request optimizes ZeRO-1 + gradient accumulation by tagging the accumulator as reduced and gradients as unreduced over the data axis, allowing the cross-replica all-reduce to run once per step instead of once per microbatch. It updates various sharding helpers and layers (MoE, normalizations) to correctly handle and propagate these tagged PartitionSpecs, which reject direct indexing or iteration. The review feedback highlights two robustness improvements: catching potential AttributeErrors in data_is_only_batch_axis when accessing the mesh of a sharding that lacks it, and defensively handling None specs in get_mesh_axes_used_by_tensor_spec to prevent runtime crashes.

[gemini-code-assist]

In data_is_only_batch_axis, param_shardings[0].mesh is accessed. However, if params_shardings contains None as a leaf (which is common for unpartitioned/replicated parameters or placeholders), or if the sharding is a SingleDeviceSharding (which lacks a mesh attribute), this will raise an AttributeError. We should catch AttributeError in the try...except block to make the helper robust.

Suggested change

	try:
	batch_axes = batch_mesh_axes(param_shardings[0].mesh, rules=config.logical_axis_rules)
	except (KeyError, ValueError):
	# No "activation_batch" rule for this mesh: leave the gradients untagged.
	return False
	try:
	batch_axes = batch_mesh_axes(param_shardings[0].mesh, rules=config.logical_axis_rules)
	except (AttributeError, KeyError, ValueError):
	# No "activation_batch" rule for this mesh: leave the gradients untagged.
	return False

[gemini-code-assist]

In get_mesh_axes_used_by_tensor_spec, if tensor_sharding_spec is None, the function will raise a TypeError when attempting to iterate over it in the list comprehension. Adding a defensive check for None at the beginning of the function will prevent this crash and align with the docstring's mention that the spec can be None.

  if tensor_sharding_spec is None:
    return []
  # Read through `.partitions`: a spec carrying reduced/unreduced axes refuses plain iteration.
  if isinstance(tensor_sharding_spec, P):
    tensor_sharding_spec = tensor_sharding_spec.partitions