Repository navigation
Shard MuseSpark's experts whole under tensor parallelism - #49787
IlyasMoutawwakil wants to merge 9 commits into
Conversation
should_convert_module matched an entry from the start of the module name with
no boundary after it, so "model.layers.1" also skipped model.layers.10 to 19,
and Qwen3.8-27B-FP8's router entry "mlp.gate" also skipped the quantized
mlp.gate_proj (its FP8 weight then reached a bf16 linear). A match from the
start now has to end on a "." it consumed, or right before a "." or the end
of the name, which keeps regex entries that carry their own boundary such as
"^model\.audio_tower(?:\.|$)". The endswith clause Hub skip-lists rely on
("gate" covering shared_expert_gate) is unchanged.
Supersedes #47990, which anchors with "$" and so stops matching such regex
entries.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… file Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
process_target_pattern unescaped \. before extracting the capturing group, so the reversed rename's group matched any character in place of a dot. Unescape only outside the group. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ision mlp keys The `.block_sparse_moe.` -> `.mlp.` rename reversed on save turned every `.mlp.` key back into `.block_sparse_moe.`, the dense and vision ones too. Name the MoE block's children in the pattern; the reversed group keeps its escaped dots. Run the reverse-mapping test that was skipped, on the composite model only, as for Llava. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
MuseSpark normalizes each expert's output before the ranks combine them, so an expert split across ranks cannot compute its norm. Its TP plan now gives the MoE the expert-parallel layout (router masking, whole experts per rank, the norm replicated with its gradient all-reduced), and its EP plan replicates the router. The model test drops its override of test_moe_parallel_plans_shard_experts: the common test now passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ac9c920 to
72c152d
Compare
daad369 to
5dd6ef5
Compare
vasqu
left a comment
There was a problem hiding this comment.
Imo this doesn't need to be fixed, don't feel like the TP here for moe is important
| "layers.*.mlp.experts": "moe_tp_experts", | ||
| } | ||
| base_model_ep_plan = { | ||
| "layers.*.mlp.gate": "replicate", |
There was a problem hiding this comment.
? why is this suddenly needed?
| # Post-expert normalization requires complete expert outputs before ranks combine them. | ||
| "layers.*.mlp.gate": "ep_router", | ||
| "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", | ||
| "layers.*.mlp.experts.down_proj": "grouped_gemm", | ||
| "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", | ||
| "layers.*.mlp.experts": "moe_tp_experts", |
There was a problem hiding this comment.
do we really need this? Werent we always using colwise/rowwise for TP. Kinda feels like a conflict between what we keep under TP and what under EP (esp after the split). Otherwise we would need to update all moe tps no?
There was a problem hiding this comment.
agree, it's an ugly workaround
|
yeh it's not worth it, will just skip tp quantization checks for it |
5a70522 to
19d9f5c
Compare
…mapping' into fg-split/fix-muse-spark-moe-plan # Conflicts: # src/transformers/conversion_mapping.py # src/transformers/quantizers/quantizers_utils.py # tests/quantization/test_quantizers_utils.py # tests/utils/test_core_model_loading.py
|
[For maintainers] Suggested jobs to run (before merge) run-slow: muse_spark |
What. MuseSpark normalizes each expert's output before the ranks combine it, so an expert split across ranks
cannot compute its norm. The TP plan now gives the MoE the expert-parallel layout: router masking, whole experts per
rank, and the norm replicated with an all-reduced gradient. The EP plan replicates the router. The model test drops
its override of
test_moe_parallel_plans_shard_experts, since the common test now passes.Depends on. the MiniMax-M3-VL fix branch (#49786, stacked only).
How tested.
tests/models/muse_spark.🤖 Generated with Claude Code