Add/fix ModelParallel layout map for qwen3_5_moe#2822
Conversation
Adds Qwen3_5MoeBackbone.get_layout_map so the model can be sharded across a device mesh with keras.distribution.ModelParallel, fixes the QKV/reverse-embedding sharding-axis bug shared by 13 of 16 models in this series, and adds layout rules for the previously-unmapped linear-attention (GatedDeltaNet) projections and the shared-expert gate. Adds a CI-safe synthetic mesh-shape sweep and a Kaggle-gated live-preset sweep alongside the existing test_distribution tier.
![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 implements model parallel distribution support for the Qwen3.5-MoE backbone. Specifically, it adds the get_layout_map static method to Qwen3_5MoeBackbone to retrieve the sharding specifications for the model's weights, including embeddings, attention projections, MoE experts, and routers. Additionally, comprehensive unit tests are introduced in qwen3_5_moe_backbone_test.py to validate the layout map under various mesh shapes and live presets. There are no review comments, so no feedback is provided.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request implements model parallel distribution support for the Qwen3.5-MoE backbone by adding the get_layout_map static method, along with extensive unit tests covering distribution, mesh shapes, and live presets. The reviewer recommends adding a comprehensive usage example to the get_layout_map docstring to comply with the repository's style guide for public APIs.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| """Get a `keras.distribution.LayoutMap` for model parallel distribution. | ||
|
|
||
| The returned `LayoutMap` contains the sharding spec for the | ||
| Qwen3_5Moe backbone's text-decoder full-attention and | ||
| linear-attention (GatedDeltaNet) layers, the MoE expert | ||
| bank/router/shared-expert, and embeddings. A few tiny linear-attention | ||
| weights (in_proj_a/in_proj_b, conv1d_kernel, dt_bias, A_log), the | ||
| shared-expert gate, and the vision encoder/interleave weights are | ||
| intentionally left replicated -- see the comments below and the | ||
| Limitations in the PR description. | ||
|
|
||
| Args: | ||
| device_mesh: keras.distribution.DeviceMesh. The device mesh | ||
| instance for distribution. | ||
| model_parallel_dim_name: str. The axis name of the device mesh, | ||
| where the weights should be partitioned on. | ||
| data_parallel_dim_name: str. The axis name of the device mesh, | ||
| where the data should be partitioned on. | ||
|
|
||
| Returns: | ||
| `keras.distribution.LayoutMap` that contains the sharding spec | ||
| for all the model weights. | ||
| """ |
There was a problem hiding this comment.
The get_layout_map method is a public API and should include a comprehensive usage example in its docstring to fully adhere to the repository style guide.
Here is an example of how you can update the docstring to include a usage example:
@staticmethod
def get_layout_map(
device_mesh,
model_parallel_dim_name="model",
data_parallel_dim_name="batch",
):
"""Get a `keras.distribution.LayoutMap` for model parallel distribution.
The returned `LayoutMap` contains the sharding spec for the
Qwen3_5Moe backbone's text-decoder full-attention and
linear-attention (GatedDeltaNet) layers, the MoE expert
bank/router/shared-expert, and embeddings. A few tiny linear-attention
weights (in_proj_a/in_proj_b, conv1d_kernel, dt_bias, A_log), the
shared-expert gate, and the vision encoder/interleave weights are
intentionally left replicated -- see the comments below and the
Limitations in the PR description.
Args:
device_mesh: keras.distribution.DeviceMesh. The device mesh
instance for distribution.
model_parallel_dim_name: str. The axis name of the device mesh,
where the weights should be partitioned on.
data_parallel_dim_name: str. The axis name of the device mesh,
where the data should be partitioned on.
Returns:
`keras.distribution.LayoutMap` that contains the sharding spec
for all the model weights.
Example:
device_mesh = keras.distribution.DeviceMesh(
shape=(1, 2),
axis_names=("batch", "model"),
devices=keras.distribution.list_devices(),
)
layout_map = Qwen3_5MoeBackbone.get_layout_map(device_mesh)
"""References
- Use Google-style docstrings for all public classes, methods, and functions. Include comprehensive examples showing usage patterns. (link)
- Use the full live-preset config (minus dtype) to build test models instead of a dims-only allowlist, so rope scaling, tie_word_embeddings, dropout, sliding_window_size, mrope_section, router_aux_loss_coefficient and the vision_encoder sub-object are exercised instead of silently dropped. Deserialize vision_encoder back into a layer instance to match from_config's handling. - Make the Tier-3 memory-budget cap tunable via the KERAS_HUB_DISTRIBUTION_TEST_MEM_BUDGET env var instead of a hardcoded 300MB that no real preset can ever fit under, so CI/bigger machines can opt into full-scale verification. - Expand the memory-estimate formula to include untied embeddings (token_embedding/reverse_embeddings is a second real vocab*hidden weight), GQA-scaled attention QKVO, and the shared expert's three matrices, on top of the existing routed expert-bank term. - Add the Gated-DeltaNet fused in_proj_qkv divisibility guard (matching qwen3_5's pattern) to both the Tier-2 and Tier-3 mesh sweeps. - Move json/get_file/CONFIG_FILE imports to module level for consistency with sibling test files. - Reword personal-dev-machine comments/skip-messages to be environment-neutral.
|
@gemini-code-assist review |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
The mesh-shape sweep and live-preset test tiers built single-layer full_attention-only models, so the GatedDeltaNet linear-attention projection rules (in_proj_qkv/in_proj_z/out_proj) and their fused-width divisibility guards were never exercised outside the fixed 2-device test_distribution mesh. Switch both tiers to a two-layer hybrid stack (one linear_attention layer, one full_attention layer), matching sibling qwen3_5's validated pattern, add the linear_attn sharding specs to the shared expected-shardings dict, and account for the linear-attention block in the live-preset memory-budget estimate. Also drop the dangling docstring reference to a nonexistent "Limitations in the PR description" section.
test_distribution and test_layout_map_mesh_shapes were not marked, so the dedicated multi_device CI job (added in keras-team#2809, filtering on -m "multi_device and not kaggle_key_required and not extra_large") would never collect any real test from this model, staying green while doing nothing. test_layout_map_live_presets already carries the marker (plus kaggle_key_required + extra_large) and is left as-is, correctly excluded from this job by design. For test_layout_map_mesh_shapes, the marker is placed as the innermost decorator (directly above the function, below @parameterized.named_parameters(...)), not outermost: absl.testing.parameterized expands each parameter case into a new function via functools.wraps, which only copies attributes already present on the wrapped function's __dict__ at expansion time. Placing pytest.mark above named_parameters sets pytestmark on the intermediate _ParameterizedTestIter object, which is discarded once the metaclass expands per-case methods, silently losing the marker on every generated case. Verified empirically via --collect-only: outer placement collected 0 of 12 parameterized cases under -m multi_device; inner placement collects all 12.
1 participant
Closes #2838
Part of #2807. Depends on #2809 (
TestCase.run_distribution_test) for CI to pass -- this PR's own diff is independent (based on master, not stacked), so it stays small and reviewable regardless of merge order.Summary
Adds
Qwen3_5MoeBackbone.get_layout_mapsoQwen3_5MoeBackbone(a hybrid full-attention/linear-attention MoE model) can be sharded across a device mesh withkeras.distribution.ModelParallel, following the same Megatron-style column/row-parallel convention used elsewhere in this series (e.g. Gemma, Qwen3):query/key/valuekernels here are(hidden_dim, num_heads, head_dim)(contractinghidden_dimfirst), not Gemma's(num_heads, hidden_dim, head_dim).hidden_dimshards on the data axis and heads shard on the model axis -- the correct orientation for Megatron-style tensor parallelism given this weight layout.token_embedding/reverse_embeddings(the(hidden, vocab)output projection) gets its own transposed spec rather than reusing the input embedding's tuple.in_proj_qkvandin_proj_z(column-parallel) andout_proj(row-parallel) have explicit layout rules. The tinyin_proj_a/in_proj_b(project tonum_v_heads, a small axis), the depthwiseconv1d_kernel, and the 1-Ddt_bias/A_logare intentionally left replicated (documented in the code and covered by the test'sallow_replicatedlist).(hidden_dim, 1)) is intentionally left replicated -- its output dim is 1, so there's nothing to shard on the model axis.num_key_value_headsis independent of and typically much smaller thannum_query_heads; sharding it would raise anIndivisibleErrorwhenever the model-mesh dimension doesn't divide it evenly), whilehidden_dimstill shards on the data axis.expert_feedforward_gate_dense/expert_feedforward_output_dense), router, and shared-expert dense layers use column/row-parallel sharding.Limitations
The vision encoder and
interleave_embeddingsweights (used when avision_encoderis attached for multimodal inputs) are not covered by any layout rule and stay fully replicated.get_layout_maponly shards the text-decoder stack (attention/linear-attention/MoE/embeddings); sharding the vision tower is out of scope for this PR.The model-parallel mesh axis can't usefully exceed a preset's
num_query_heads-- that's the axis actually being tensor-parallelized here -- see the design doc's "hard limit" section for the full reasoning. For this model's real presets:qwen3_5_moe_35b_a3bqwen3_5_moe_35b_a3b_baseBoth presets use
num_key_value_heads=2, so any mesh past model-axis=16 fails on the query/attention-output kernels specifically, not onhidden_dimor the MoE expert bank -- those stay divisible well past this range for both presets.Tests
qwen3_5_moe_backbone_test.pyadds test tiers mirroring the pattern already reviewed for Gemma (gemma_backbone_test.py):test_distribution): exact end-to-end sharding-spec + full-weight-coverage + forward/backward/optimizer-step + numerical-parity/convergence check viaTestCase.run_distribution_testfrom Add shared ModelParallel distribution-test helper #2809, on a real 2-device mesh (pinned to exactly 2, notlen(devices), so the default config'snum_key_value_heads=2is exercised deterministically). Usesis_moe=Truefor the looser MoE parity tolerance and asserts both the full-attention and linear-attention (GatedDeltaNet) sublayer specs, since the default 4-layer test config includes both layer types.test_layout_map_mesh_shapes): CI-safe sweep overCAPPED_MESH_SHAPES((2,4),(1,8),(4,4),(2,2,2),(1,1,8),(2,2,4)-- capped at 16 total simulated devices) crossed with two synthetic width classes (QWEN3_5_MOE_35B_A3B_DIMS,QWEN3_5_MOE_SMALL_DIMS) that preserve the real preset's query:kv head ratio (8:1 GQA) and linear key:value head ratio (1:2), scaled down ~20-30x from the realQwen3.5-35B-A3B-class config. Each sweep case builds a two-layer hybrid stack -- onelinear_attentionlayer and onefull_attentionlayer -- so both token-mixer rule sets are asserted in every (width, mesh) combo, not just at the fixed 2-device mesh in Tier 1 (layout rules are per-decoder-block regexes, so depth beyond one-of-each doesn't affect spec matching), and asserts full-weight-coverage via a shared helper. Skips (not fails) any(width, mesh)combo wherenum_query_headsor the linear-attention fusedin_proj_qkvwidth doesn't divide the mesh's model axis -- an inherent tensor-parallelism limit, not a bug.test_layout_map_live_presets, opt-in viakaggle_key_required+multi_device+extra_largemarkers): fetches every real preset'sconfig.json(metadata only, no weights), dedupes by divisibility-relevant dims, forces the same two-layer hybrid stack as Tier 2 onto each width class, and for each unique width class estimates its bf16 two-decoder-block + expert-bank build footprint before attempting a build, skipping (viaself.subTest+ a logged reason) any combo that would exceed a 300MB local safety threshold.Verification
Real local test results (this branch's own base,
master, doesn't yet include #2809'srun_distribution_testhelper, so these were run against a temporary checkout with #2809's helper commit applied on top of this PR's diff -- the diff itself is unchanged and unstacked):Run with
KERAS_BACKEND=jax XLA_FLAGS=--xla_force_host_platform_device_count=16 python -m pytest keras_hub/src/models/qwen3_5_moe/qwen3_5_moe_backbone_test.py -v -n 2, capped at 16 simulated CPU devices per the mesh-shape matrix. Both skips are the expected default-config skips (test_saved_modelneeds--run_large;test_layout_map_live_presetsneeds--run_extra_large+ Kaggle credentials and is intentionally not exercised in this pass by design -- see the module comment aboveCAPPED_MESH_SHAPES).ruff check/ruff format --check: clean on both changed files.