Qwen3 has no layout map yet, so it can't be sharded with ModelParallel out of the box. This adds Qwen3Backbone.get_layout_map following the same column/row-parallel convention Gemma uses. The QKV kernels here are (hidden, num_heads, head_dim), so hidden shards on the data axis and heads on the model axis; k/v stay replicated because num_key_value_heads under GQA is small and often won't divide the model mesh dim. reverse_embeddings gets its own transposed spec since it's (hidden, vocab), not the input embedding's (vocab, hidden).
Tests come in the usual three tiers (exact spec + coverage, a capped mesh-shape sweep over scaled-down width classes that keep the real query:kv ratios, and an opt-in live-preset sweep with a memory guard). Part of #2807, tracked in #2811.
Qwen3 has no layout map yet, so it can't be sharded with
ModelParallelout of the box. This addsQwen3Backbone.get_layout_mapfollowing the same column/row-parallel convention Gemma uses. The QKV kernels here are(hidden, num_heads, head_dim), so hidden shards on the data axis and heads on the model axis; k/v stay replicated becausenum_key_value_headsunder GQA is small and often won't divide the model mesh dim.reverse_embeddingsgets its own transposed spec since it's(hidden, vocab), not the input embedding's(vocab, hidden).Tests come in the usual three tiers (exact spec + coverage, a capped mesh-shape sweep over scaled-down width classes that keep the real query:kv ratios, and an opt-in live-preset sweep with a memory guard). Part of #2807, tracked in #2811.