Add ModelParallel layout map for Qwen3#2811
Conversation
Adds `Qwen3Backbone.get_layout_map` to shard backbone weights across a device mesh for `keras.distribution.ModelParallel`, following the Megatron-style column/row-parallel convention used elsewhere in this series (e.g. Gemma). QKV kernels are `(hidden, num_heads, head_dim)`, so `hidden` (the contracting dim) is sharded on the data-parallel axis and heads on the model-parallel axis; key/value heads are left replicated on the model axis since GQA's num_key_value_heads is often too small to divide evenly. `token_embedding/reverse_embeddings` is `(hidden, vocab)`, the transpose of the input embedding, so it gets its own `(data_dim, model_dim)` spec rather than reusing the input embedding's tuple. Test coverage adds a Tier-1 `test_distribution` (untied embeddings, exact sharding spec check via the new `run_distribution_test` helper), a Tier-2 CI-safe sweep over capped mesh shapes with memory-scaled synthetic dims preserving real query:kv head ratios across three width classes, and a Tier-3 opt-in sweep against live preset configs with a memory-budget guard that skips any width class whose estimated build size would be unsafe to construct locally.
![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 introduces the get_layout_map static method to Qwen3Backbone to support model parallel distribution, along with comprehensive unit tests verifying the layout map across various mesh shapes and live presets. The feedback recommends using a context manager with an explicit UTF-8 encoding when opening preset configuration files in the tests to prevent resource leaks and platform-specific encoding issues.
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.
| for preset in Qwen3Backbone.presets: | ||
| try: | ||
| path = get_file(preset, CONFIG_FILE) | ||
| cfg = json.load(open(path))["config"] |
There was a problem hiding this comment.
Opening a file using open(path) without specifying an explicit encoding can lead to UnicodeDecodeError on platforms where the default encoding is not UTF-8 (such as Windows). Additionally, not using a context manager (with statement) leaves the file descriptor open until garbage collection, which can cause resource leaks when running many tests in a loop.
Using with open(path, "r", encoding="utf-8") as f: resolves both issues.
| cfg = json.load(open(path))["config"] | |
| with open(path, "r", encoding="utf-8") as f: | |
| cfg = json.load(f)["config"] |
References
- Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. (link)
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request adds a static method get_layout_map to Qwen3Backbone to support model parallel distribution, along with comprehensive unit tests to verify the layout map across various mesh shapes and live presets. The reviewer recommended using a with statement and specifying encoding="utf-8" when opening the preset configuration file in the tests to prevent resource leaks and ensure platform-agnostic behavior.
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.
| path = get_file(preset, CONFIG_FILE) | ||
| cfg = json.load(open(path))["config"] |
There was a problem hiding this comment.
Opening a file using open() without a with statement or explicitly closing it can lead to unclosed file descriptors, causing resource leaks. Additionally, it is a best practice to specify encoding="utf-8" when reading JSON files to ensure platform-agnostic behavior (especially on Windows platforms). Use a with statement to safely manage the file lifecycle.
| path = get_file(preset, CONFIG_FILE) | |
| cfg = json.load(open(path))["config"] | |
| path = get_file(preset, CONFIG_FILE) | |
| with open(path, "r", encoding="utf-8") as f: | |
| cfg = json.load(f)["config"] |
Use a context manager when opening the preset config file and pass explicit encoding, per Gemini review on keras-team#2811.
…ng memory estimate in qwen3 distribution test - test_layout_map_live_presets now passes the full fetched preset config (minus dtype) to Qwen3Backbone instead of a dims-only allowlist, so the live-preset build actually exercises the preset's real architecture flags (rope scaling, tie_word_embeddings, etc.) instead of silently dropping them. - The Tier-3 memory-budget cap is now overridable via the KERAS_HUB_DISTRIBUTION_TEST_MEM_BUDGET env var (defaults to the previous 300MB, so local behavior is unchanged) instead of a hardcoded ceiling no real preset could ever fit under. - The memory estimate formula now includes GQA-scaled attention QKVO params and doubles the embedding term when tie_word_embeddings=False, matching the model's actual untied reverse_embeddings weight. - Reworded comments/skip messages to be environment-neutral instead of referencing a specific personal dev machine.
|
@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! |
test_distribution and test_layout_map_mesh_shapes construct real DeviceMesh/ModelParallel setups but were missing the multi_device pytest marker, so the CI job that filters on "multi_device and not kaggle_key_required and not extra_large" was silently collecting zero tests from this file. Note: for test_layout_map_mesh_shapes, the marker must be placed below @parameterized.named_parameters (directly on the function), not above it -- absl.testing.parameterized's TestGeneratorMetaclass expands the decorated method into per-case methods without copying pytestmark from the outer _ParameterizedTestIter, so a mark placed above the parameterized decorator is silently dropped from every generated test case.
get_layout_map now takes an optional num_query_heads argument. Real Qwen3 presets have as few as 8 query heads, and larger model-parallel mesh axes (32-way, 64-way, ...) don't evenly divide every preset's head count -- sharding that axis unconditionally raises IndivisibleError. When num_query_heads is given and doesn't divide the model-axis mesh size, the query/attention_output kernels' head axis now falls back to fully replicated instead of sharded, mirroring the existing key/value fallback. Omitting num_query_heads preserves prior (always-sharded) behavior. Adds test_layout_map_query_heads_fallback, a regression test building a real (1, 8) mesh with num_query_heads=4 to prove the fallback triggers instead of raising IndivisibleError.
This reverts commit c17c820.
1 participant
Closes #2831
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
Qwen3Backbone.get_layout_mapsoQwen3Backbonecan be sharded across a device mesh withkeras.distribution.ModelParallel, following the same Megatron-style column/row-parallel convention already used elsewhere in this series (e.g. Gemma).Sharding rules:
token_embedding/embeddings:(model_dim, data_dim)-- vocab-parallel input embedding.token_embedding/reverse_embeddings:(data_dim, model_dim). This is the transpose of the input embedding ((hidden, vocab), not(vocab, hidden)), so it needs its own spec rather than reusing the input embedding's tuple -- only exists as a real weight whentie_word_embeddings=False.self_attention.*query.kernel:(data_dim, model_dim, None). QKV kernels are(hidden, num_heads, head_dim)(see thebqm,muh->bquheinsum inQwen3Attention.build), sohidden(the contracting dim) shards on the data axis and heads shard on the model axis -- this matches theattention_outputrule below, which already shards heads onmodel_dim.self_attention.*(key|value).kernel:(data_dim, None, None). Key/value heads are sized bynum_key_value_heads(GQA), which is independent of and typically much smaller thannum_query_heads, so sharding that axis on the model-parallel dim would raise anIndivisibleErrorwhenever the model-mesh dimension doesn't divide it evenly;hiddenstill shards on the data axis for memory savings.attention_output.kernel:(model_dim, None, data_dim).feedforward_intermediate_dense.kernel/feedforward_gate_dense.kernel:(data_dim, model_dim).feedforward_output_dense.kernel:(model_dim, data_dim).Query-head-count divisibility fallback
get_layout_mapalso takes an optionalnum_query_headsargument. Real Qwen3 presets range from 8 to 64 query heads, and a model-parallel mesh axis (16-way, 32-way, 64-way, ...) does not evenly divide every one of those -- sharding the query/attention_outputhead axis unconditionally would raiseIndivisibleErrorfor those combinations. Whennum_query_headsis passed and does not divide the model-parallel axis size,query.kernelandattention_output.kernelfall back to leaving their head axis fully replicated instead of sharded, the same fallback already applied to the key/value kernels above. Whennum_query_headsis omitted (the default), the head axis is always sharded on the model-parallel dim, preserving prior behavior.Tests
qwen3_backbone_test.pyadds three tiers, mirroring the pattern already reviewed and merged for Gemma (gemma_backbone_test.py):test_distribution): exact end-to-end sharding-spec + coverage + forward/backward/optimizer-step + numerical-parity check via the newTestCase.run_distribution_testhelper from Add shared ModelParallel distribution-test helper #2809, on a real(1, 2)mesh, withtie_word_embeddings=Falsesoreverse_embeddingsis a real weight and its spec is actually exercised.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 three synthetic width classes (QWEN3_0_6B_DIMS,QWEN3_4B_DIMS,QWEN3_32B_DIMS) that preserve each real preset's exact query:kv head ratio (2:1, 4:1, 8:1) while scaling every other dimension down ~20-30x from the real preset so the sweep is cheap to build repeatedly, asserting the same spec-coverage helper (_assert_qwen3_shardings_and_coverage) plus an "every rank>=2 weight is mapped" completeness check. Skips (not fails) any(width, mesh)combo wherenum_query_headsdoesn't divide the mesh's model axis, sinceget_layout_mapisn't givennum_query_headsin this sweep and so cannot apply the fallback -- an inherent tensor-parallelism limit under that call convention, not a bug.test_layout_map_query_heads_fallback: dedicated regression test for the fallback above -- builds a real(1, 8)device mesh withnum_query_heads=4(which does not divide 8), passesnum_query_headsthrough toget_layout_map, and asserts the query/attention_outputkernels build successfully with a fully replicated head axis instead of raisingIndivisibleError.test_layout_map_live_presets, opt-in viakaggle_key_required+multi_device+extra_largemarkers): fetches every real preset'sconfig.json(cheap metadata only, never weights), dedupes by divisibility-relevant dims, and for each unique width class estimates its bf16 build footprint ((vocab*hidden + 3*hidden*intermediate) * 2 * 3) before attempting a build, skipping (viaself.subTest+ a logged reason) any combo that would exceed a 300MB local safety threshold rather than risking an OOM on a real multi-GB checkpoint.Verification
Ran locally against a checkout with #2809's
run_distribution_testhelper patched intoTestCase:All Tier 1, Tier 2, and the new fallback regression test pass. Skips are the mesh shapes needing more than 8 simulated devices,
test_saved_model(@pytest.mark.large, unrelated to distribution), and aTestCase.test_sessioncollection artifact; the 1 deselection is Tier 3 (kaggle_key_required), which needs real Kaggle credentials and is opt-in by design.Locally,
ruff checkandruff format --checkboth pass clean on the two changed files.Limitations
The model-parallel mesh axis can't usefully exceed a preset's
num_query_heads-- that's the axis actually being tensor-parallelized here -- so mesh sizes beyond that just fall back to a fully replicated head axis instead of giving any further sharding benefit; see the design doc's "hard limit" section for why.Real numbers for Qwen3 presets:
qwen3_0.6b_enqwen3_1.7b_enqwen3_4b_enqwen3_8b_enqwen3_14b_enqwen3_32b_enqwen3_embedding_0.6b_enqwen3_embedding_4b_enqwen3_embedding_8b_enharrier_embedding_oss_0.6bNote the
qwen3_embedding_*presets are much tighter than theirnum_query_headsalone would suggest (16-32) -- their safe ceiling is driven bytoken_embedding/embeddings, whose vocab dim doesn't divide evenly even at a 2-way model axis. That's a separate, vocab-size-driven failure mode, not the query-head ceiling described above, and it's worth checking independently if you're sharding one of the embedding presets.