Onboard explicit sharding to the Gemma models - #4996
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables explicit sharding support for Gemma, Gemma2, and Gemma3 models under ZeRO-1 and gradient accumulation by explicitly propagating sharding layouts through attention, MLP, and normalization layers, and introducing get_grad_shardings to align gradient layouts with optimizer moments. Feedback points out that mlp_intermediate_sharding in the Gemma models hardcodes "activation_length" instead of using self.activation_axis_names[1], which will cause a ShardingTypeError under explicit sharding. Additionally, a potential sharding mismatch is identified in train.py when parameter_memory_host_offload is enabled alongside ZeRO-1 and explicit sharding.
|
|
||
| mlp_lnx = self.mlp(attn_output, deterministic=deterministic) | ||
| mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names) | ||
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_mlp")) |
There was a problem hiding this comment.
The mlp_intermediate_sharding is hardcoded to use "activation_length" as the sequence/length dimension. However, the layer's activations (and self.activation_axis_names) use "activation_norm_length".
Under shard_mode=explicit, JAX enforces strict sharding type matching. If the intermediate sharding specifies "activation_length" while the input attn_output uses "activation_norm_length", JAX will either raise a ShardingTypeError or perform an unnecessary and expensive resharding (all-to-all) to align the dimensions. To ensure consistency and avoid sharding mismatches, mlp_intermediate_sharding should use self.activation_axis_names[1] instead of the hardcoded "activation_length".
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_mlp")) | |
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", self.activation_axis_names[1], "activation_mlp")) |
There was a problem hiding this comment.
Thanks, but I don't think this one is right — the hardcoded activation_length is deliberate, and the suggested change would silently drop context parallelism from the MLP intermediate.
The intermediate is [batch, seq, mlp], and activation_mlp maps to ['tensor', 'tensor_sequence']. activation_norm_length maps to ['tensor_sequence', 'context', 'context_usp_ulysses'] — it shares tensor_sequence with activation_mlp, which is exactly the point: with tensor-sequence parallelism the normed activation shards its sequence dim because its embed dim is replicated, whereas the MLP intermediate already shards tensor_sequence on the mlp dim, so its sequence dim must not. activation_length (['context', 'context_usp_ulysses']) is the correct partner axis.
Resolving both against the rule table shows what the suggestion actually does:
current : P(('data','fsdp','fsdp_transpose','expert'), ('context','context_usp_ulysses'), ('tensor','tensor_sequence'))
suggested: P(('data','fsdp','fsdp_transpose','expert'), None, ('tensor','tensor_sequence'))
The sequence dim goes from context-sharded to replicated, because tensor_sequence is already consumed by activation_mlp and flax then drops the whole rule entry. That is a real regression under context parallelism, not a fix.
This also matches the two models already onboarded to explicit sharding: llama2.py:192 and deepseek.py:83 use the same ("activation_batch", "activation_length", "activation_mlp") triple. Empirically there is no ShardingTypeError — the Gemma parity test (tests/unit/gemma_explicit_sharding_test.py) and the TPU integration test both pass as written, and gemma3-4b was verified end to end with ici_tensor_parallelism=2.
Leaving as is.
| inputs = self._maybe_shard_with_logical(inputs, self.activation_axis_names) | ||
| inputs = checkpoint_name(inputs, "decoder_layer_input") | ||
| lnx_sharding = create_sharding(self.mesh, self.activation_axis_names) | ||
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_mlp")) |
There was a problem hiding this comment.
The mlp_intermediate_sharding is hardcoded to use "activation_length" as the sequence/length dimension. However, the layer's activations (and self.activation_axis_names) use "activation_norm_length" (or "prefill_activation_norm_length" in prefill mode).
Under shard_mode=explicit, JAX enforces strict sharding type matching. If the intermediate sharding specifies "activation_length" while the input attn_output uses "activation_norm_length" (or "prefill_activation_norm_length"), JAX will either raise a ShardingTypeError or perform an unnecessary and expensive resharding (all-to-all) to align the dimensions. To ensure consistency and avoid sharding mismatches, mlp_intermediate_sharding should use self.activation_axis_names[1] instead of the hardcoded "activation_length".
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_mlp")) | |
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", self.activation_axis_names[1], "activation_mlp")) |
There was a problem hiding this comment.
Thanks, but I don't think this one is right — the hardcoded activation_length is deliberate, and the suggested change would silently drop context parallelism from the MLP intermediate.
The intermediate is [batch, seq, mlp], and activation_mlp maps to ['tensor', 'tensor_sequence']. activation_norm_length maps to ['tensor_sequence', 'context', 'context_usp_ulysses'] — it shares tensor_sequence with activation_mlp, which is exactly the point: with tensor-sequence parallelism the normed activation shards its sequence dim because its embed dim is replicated, whereas the MLP intermediate already shards tensor_sequence on the mlp dim, so its sequence dim must not. activation_length (['context', 'context_usp_ulysses']) is the correct partner axis.
Resolving both against the rule table shows what the suggestion actually does:
current : P(('data','fsdp','fsdp_transpose','expert'), ('context','context_usp_ulysses'), ('tensor','tensor_sequence'))
suggested: P(('data','fsdp','fsdp_transpose','expert'), None, ('tensor','tensor_sequence'))
The sequence dim goes from context-sharded to replicated, because tensor_sequence is already consumed by activation_mlp and flax then drops the whole rule entry. That is a real regression under context parallelism, not a fix.
This also matches the two models already onboarded to explicit sharding: llama2.py:192 and deepseek.py:83 use the same ("activation_batch", "activation_length", "activation_mlp") triple. Empirically there is no ShardingTypeError — the Gemma parity test (tests/unit/gemma_explicit_sharding_test.py) and the TPU integration test both pass as written, and gemma3-4b was verified end to end with ici_tensor_parallelism=2.
Leaving as is.
|
|
||
| # MLP block. | ||
| mlp_lnx = self.mlp(attn_output, deterministic=deterministic) | ||
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_mlp")) |
There was a problem hiding this comment.
The mlp_intermediate_sharding is hardcoded to use "activation_length" as the sequence/length dimension. However, the layer's activations (and self.activation_axis_names) use "activation_norm_length" (or "prefill_activation_norm_length" in prefill mode).
Under shard_mode=explicit, JAX enforces strict sharding type matching. If the intermediate sharding specifies "activation_length" while the input attn_output uses "activation_norm_length" (or "prefill_activation_norm_length"), JAX will either raise a ShardingTypeError or perform an unnecessary and expensive resharding (all-to-all) to align the dimensions. To ensure consistency and avoid sharding mismatches, mlp_intermediate_sharding should use self.activation_axis_names[1] instead of the hardcoded "activation_length".
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_mlp")) | |
| mlp_intermediate_sharding = create_sharding(self.mesh, ("activation_batch", self.activation_axis_names[1], "activation_mlp")) |
There was a problem hiding this comment.
Thanks, but I don't think this one is right — the hardcoded activation_length is deliberate, and the suggested change would silently drop context parallelism from the MLP intermediate.
The intermediate is [batch, seq, mlp], and activation_mlp maps to ['tensor', 'tensor_sequence']. activation_norm_length maps to ['tensor_sequence', 'context', 'context_usp_ulysses'] — it shares tensor_sequence with activation_mlp, which is exactly the point: with tensor-sequence parallelism the normed activation shards its sequence dim because its embed dim is replicated, whereas the MLP intermediate already shards tensor_sequence on the mlp dim, so its sequence dim must not. activation_length (['context', 'context_usp_ulysses']) is the correct partner axis.
Resolving both against the rule table shows what the suggestion actually does:
current : P(('data','fsdp','fsdp_transpose','expert'), ('context','context_usp_ulysses'), ('tensor','tensor_sequence'))
suggested: P(('data','fsdp','fsdp_transpose','expert'), None, ('tensor','tensor_sequence'))
The sequence dim goes from context-sharded to replicated, because tensor_sequence is already consumed by activation_mlp and flax then drops the whole rule entry. That is a real regression under context parallelism, not a fix.
This also matches the two models already onboarded to explicit sharding: llama2.py:192 and deepseek.py:83 use the same ("activation_batch", "activation_length", "activation_mlp") triple. Empirically there is no ShardingTypeError — the Gemma parity test (tests/unit/gemma_explicit_sharding_test.py) and the TPU integration test both pass as written, and gemma3-4b was verified end to end with ici_tensor_parallelism=2.
Leaving as is.
| if grad_shardings is params_shardings: | ||
| nnx.update(state.model, curr_params) | ||
| # Otherwise (ZeRO-1 + explicit sharding) leave state.model in the ZeRO-1 layout | ||
| # it was donated in. curr_params above is all-gathered over "data" for the | ||
| # forward/backward pass only; writing it back would hand the optimizer params | ||
| # sharded P('fsdp') alongside gradients and moments sharded P(('data', 'fsdp')), | ||
| # which explicit sharding rejects instead of silently reconciling. |
There was a problem hiding this comment.
When shard_optimizer_over_data (ZeRO-1) and shard_mode=explicit are active, grad_shardings is not params_shardings is True. Skipping the update to state.model with the all-gathered curr_params on line 530 correctly prevents a layout mismatch with the optimizer.
However, if parameter_memory_host_offload is also enabled, state.model is unconditionally updated on line 511 with curr_params (which is in the pre-Zero-1 layout params_shardings):
nnx.update(state.model, curr_params) # ensure state.model has device params for optimizer updateThis update overwrites the ZeRO-1 layout of state.model with the pre-Zero-1 layout. Since the update on line 530 is skipped, state.model is left in the pre-Zero-1 layout, while the optimizer moments and gradients are in the ZeRO-1 layout. This will cause a sharding mismatch under explicit sharding during the optimizer update. Consider conditionally updating state.model on line 511 only if grad_shardings is params_shardings (or handle the ZeRO-1 layout update correctly).
There was a problem hiding this comment.
Good catch — the combination is genuinely broken, though the mechanism turned out to be one step further down. Fixed in f58857b.
Reproduced first:
python -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
decoder_block=gemma gradient_accumulation_steps=4 shard_optimizer_over_data=True \
shard_mode=explicit parameter_memory_host_offload=True param_scan_axis=0 \
ici_data_parallelism=2 ici_fsdp_parallelism=-1 ...
jax._src.core.ShardingTypeError: add got incompatible shardings for broadcasting:
('fsdp', None), (('data', 'fsdp'), None)
But dumping the shardings at apply_gradients showed it is the gradients, not state.model, that arrive in the wrong layout — every leaf was in the pre-ZeRO-1 layout despite grad_shardings being correct (P(('data','fsdp'), None) for logits_dense.kernel, memory kind device). The culprit is the host-offload device_put that runs after the ZeRO-1 relayout:
if config.parameter_memory_host_offload:
raw_grads = jax.device_put(raw_grads, max_utils.with_memory_kind(params_shardings, "device"))Only the memory kind should change there, but params_shardings carries the pre-ZeRO-1 layout too, so the device_put undoes the relayout the GA path had just performed. Targeting grad_shardings fixes it, and is a no-op elsewhere since grad_shardings is params_shardings whenever ZeRO-1 + explicit sharding is not active.
On the specific line you flagged: the nnx.update(state.model, curr_params) write-back is fine as written. With only the device_put fixed, gemma runs 3 clean steps with host offload + ZeRO-1 + GA + explicit, and the losses match auto to 3 decimals (10.879 / 10.869 / 10.863–4). Guarding that write-back as well was tried and is not needed — state.model is re-read from the state's own layout at the optimizer update.
One unrelated pre-existing issue surfaced while testing: parameter_memory_host_offload=True with shard_optimizer_over_data=False fails with memory_space of all inputs passed to add must be the same. Got one operand with type: float32<host>[256] and another operand with type: float32[256] — identically under shard_mode=auto, so it is not caused by this PR and is out of scope here.
Added gemma-host-offload as a fourth subtest of test_tpu_gemma_zero1_gradient_accumulation_explicit; all 4 subtests pass on TPU v7-8 in 33 s.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Add shard_mode=explicit support to the Gemma 1/2/3 decoder blocks and make
ZeRO-1 + gradient accumulation + explicit sharding work end to end.
Under auto sharding a missing annotation costs nothing because GSPMD infers a
layout; under explicit sharding JAX type-checks every operation, so each Gemma
layer has to name its activations and hand its sub-modules an out_sharding
(and intermediate_sharding for the MLP) instead of relying on inference.
ZeRO-1 additionally shards the optimizer moments over the "data" axis. Auto
reconciles the resulting mismatch silently; explicit refuses to add a gradient
sharded P('fsdp') to a moment sharded P(('data','fsdp')). get_grad_shardings()
returns the layout gradients must reach the optimizer in, and both the
gradient-accumulation path (via the new out_grad_shardings argument) and the
single-microbatch path now target it. This was broken for every model, not
just Gemma.
Also name the QK/V-norm output layout in attentions.py, so gemma3 with tensor
parallelism gathers the norm scale instead of failing to compile, and keep
unreduced/reduced annotations when truncating a PartitionSpec.
All of it is a no-op under the default shard_mode=auto.
Tested on TPU v7-8: losses match auto for gemma-2b, gemma2-2b, gemma2-9b,
gemma3-4b and gemma3-12b, with and without tensor parallelism, ZeRO-1 and
gradient accumulation. ZeRO-1 without gradient accumulation is 7.5% faster
than auto (0.331 s vs 0.358 s per step on gemma-2b).
Under parameter_memory_host_offload the gradients are moved back to device
memory before the optimizer update. That device_put targeted params_shardings,
which also relayouted them out of the ZeRO-1 layout the gradient-accumulation
path had just put them in, so explicit sharding rejected the moment update with
add got incompatible shardings for broadcasting: ('fsdp', None), (('data', 'fsdp'), None)
Target grad_shardings instead, which equals params_shardings whenever ZeRO-1 +
explicit sharding is not active. Covered by a new host-offload subtest.
…convention gemma_explicit_sharding_test.py was named after the feature rather than the code under test. Fold it into tests/unit/gemma_layers_test.py, alongside the existing gemma3_layers_test.py / gemma4_layers_test.py, so future Gemma layer tests have an obvious home.
0b0dca9 to
01aeda2
Compare
2 participants
Onboard explicit sharding to the Gemma models
Description
Adds
shard_mode=explicitsupport to the Gemma 1 / 2 / 3 decoder blocks, and makes the combination ZeRO-1 + gradient accumulation + explicit sharding work end to end (previously it raised aShardingTypeErrorfor every model, including the already-onboarded llama2 and deepseek).Under
shard_mode=auto, a missing sharding annotation costs nothing — GSPMD infers a layout. Undershard_mode=explicit, JAX type-checks the sharding of every operation, so the same omission is a hard error and every intermediate has to be named. This PR follows the pattern established for llama2 in #2470 and #2783:nn.with_logical_constraint(...)→maybe_shard_with_logical(..., shard_mode=...)RMSNorm,Attention,MlpBlock) are handed an explicitout_sharding(andintermediate_shardingfor the MLP) instead of relying on inference.All of it is a no-op under
shard_mode=auto:maybe_shard_with_logicalfalls back tojax.lax.with_sharding_constraint, andRMSNorm.__call__dropsout_shardingwhen the mode is not explicit.Key changes
out_sharding/intermediate_shardingdown to their sub-modules;gemma,gemma2andgemma3are added to the explicit-sharding allowlist inconfigs/types.py.shard_optimizer_over_datathe optimizer moments are additionally sharded over thedataaxis.autoreconciles the mismatch silently;explicitrefuses to add a gradient shardedP('fsdp')to a moment shardedP(('data','fsdp')). A newsharding.get_grad_shardings()returns the layout the gradients must reach the optimizer in, and both the gradient-accumulation path and the single-microbatch path now target it.gemma3-4bwithici_tensor_parallelism=2failed to compile because the QK-norm scale carries the generic("norm",) -> ["tensor"]rule while the heads dimension is already tensor-sharded. Naming the normed output layout fixes it.Detailed changes by file
src/maxtext/models/gemma.pymaybe_shard_with_logicalfor every activation;out_shardingon both norms, attention and the MLP;intermediate_shardingon the MLP;shard_modethreaded intoRMSNorm.src/maxtext/models/gemma2.pysrc/maxtext/models/gemma3.pyGemma3DecoderLayerandGemma3ScannableBlock.src/maxtext/configs/types.pygemma,gemma2,gemma3added to the explicit-sharding allowlist.src/maxtext/layers/attentions.pyauto.src/maxtext/utils/sharding.pyget_grad_shardings();_extract_param_only()hoisted to module level so it can be shared;truncate_out_sharding()now preservesunreduced/reducedannotations instead of raising on them.src/maxtext/utils/gradient_accumulation.pyout_grad_shardingsargument (defaults toparams_shardings): resolving theunreducedannotation against it emits the cross-replica combine straight into the optimizer's layout — a single reduce-scatter overdatarather than an all-reduce plus a slice.src/maxtext/trainers/pre_train/train.pygrad_shardingsonce and passes it to GA; reshards gradients on the single-microbatch path; under ZeRO-1 + explicit, no longer writes the all-gathered params back intostate.model(which would hand the optimizer params and moments in disagreeing layouts). Theparameter_memory_host_offloaddevice_puttargetsgrad_shardingsrather thanparams_shardings, so moving gradients back to device memory no longer relayouts them out of the ZeRO-1 layout.Tests
New:
tests/unit/gemma_layers_test.py— runs one forward pass per Gemma layer family in both shard modes on an 8-device CPU mesh and asserts the activations match. Verified to fail on all four layer families when the model changes are reverted.tests/unit/sharding_nnx_test.py::TestGetGradShardings— the five branches ofget_grad_shardings, including that the returned tree matches the structure ofnnx.split(model, nnx.Param, ...).tests/unit/sharding_nnx_test.py—truncate_out_shardingkeepsunreduced/reduced.tests/unit/gradient_accumulation_nnx_test.py::TestGradientAccumulationOutGradShardings—out_grad_shardingsdefaults toparams_shardings, overrides the returned layout, and does not change the gradient values.tests/integration/train_tests.py::test_tpu_gemma_zero1_gradient_accumulation_explicit— gemma / gemma2 / gemma3 trained for 3 steps with ZeRO-1 + GA + explicit sharding on TPU, plus agemma-host-offloadsubtest coveringparameter_memory_host_offload=Trueon the same combination.Existing
tests/unit/sharding_nnx_test.py,tests/unit/gradient_accumulation_nnx_test.pyandtests/unit/sharding_test.pypass unchanged.End-to-end runs (TPU v7-8, 4 Ironwood chips / 8 devices)
Losses under
explicitmatchautoto 3 decimals over 20 steps ongemma-2b, and in every smoke run acrossgemma-2b,gemma2-2b,gemma2-9b,gemma3-4b,gemma3-12b— with and without tensor parallelism, ZeRO-1 and gradient accumulation.Step time,
gemma-2b(2.506 B params),per_device_batch_size=1,max_target_length=2048,ici_data_parallelism=2,ici_fsdp_parallelism=4:pdb=4)Peak memory after parameter init is identical in both modes (5.87 GB / 94.75 GB per device).
With
parameter_memory_host_offload=Trueon top of ZeRO-1 + GA,explicitandautoagree to 3 decimals as well (10.879 / 10.869 / 10.863–4 over the first three steps).Known gap. With gradient accumulation,
explicitis 1–3 % slower thanauto. The HLO shows explicit emitting the nine intendeddata-axis reduce-scatters for the scanned parameters, plus a duplicateddata-axis reduction of the ~262 MB tied-embedding gradient (one plain all-reduce and oneall-reduce-scatterlowered as pad + all-reduce + slice). The gap scales with vocabulary size, lives in the shared embedding/logits code rather than in any Gemma layer, and reproduces identically onllama2-7b(1.887 s auto vs 1.912 s explicit with ZeRO-1 + GA=4), which was onboarded before this PR. It is therefore tracked as a follow-up rather than fixed here.Unrelated pre-existing issue.
parameter_memory_host_offload=Truecombined withshard_optimizer_over_data=Falsefails withmemory_space of all inputs passed to add must be the same— identically undershard_mode=auto, so it is not caused by this PR and is left alone.Platform note.
ici_fsdp_parallelism=1(data-parallel only) with ZeRO-1 + GA segfaults the TPU7x compiler for llama2 and Gemma alike, with and without this change — that is b/517509898, which the existing llama2 test already skips for. The new Gemma integration test usesici_data_parallelism=2with FSDP filling the remaining devices, which runs cleanly.Checklist
shard_mode=autoverifiedpyinkandpylintclean on the changed filesshard_mode=auto