Skip to content

Shard MuseSpark's experts whole under tensor parallelism - #49787

Closed
IlyasMoutawwakil wants to merge 9 commits into
fg-split/fix-minimax-m3-vl-save-mappingfrom
fg-split/fix-muse-spark-moe-plan
Closed

IlyasMoutawwakil wants to merge 9 commits into
fg-split/fix-minimax-m3-vl-save-mappingfrom
fg-split/fix-muse-spark-moe-plan

Conversation

@IlyasMoutawwakil

@IlyasMoutawwakil IlyasMoutawwakil commented Oct 20, 2026 •

Copy link
Copy Markdown
Member

🤗 Transformers CI

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

IlyasMoutawwakil and others added 2 commits October 20, 2026 05:45
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>
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

IlyasMoutawwakil and others added 6 commits October 20, 2026 08:59
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>
@IlyasMoutawwakil
IlyasMoutawwakil force-pushed the fg-split/fix-muse-spark-moe-plan branch from ac9c920 to 72c152d Compare October 20, 2026 16:04
@IlyasMoutawwakil
IlyasMoutawwakil force-pushed the fg-split/fix-minimax-m3-vl-save-mapping branch from daad369 to 5dd6ef5 Compare October 20, 2026 16:04

@vasqu vasqu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

? why is this suddenly needed?

Comment on lines +113 to +118
# 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",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

agree, it's an ugly workaround

@IlyasMoutawwakil

Copy link
Copy Markdown
Member Author

yeh it's not worth it, will just skip tp quantization checks for it

@IlyasMoutawwakil
IlyasMoutawwakil removed this pull request from stack #49788 October 20, 2026 16:25
@IlyasMoutawwakil
IlyasMoutawwakil force-pushed the fg-split/fix-minimax-m3-vl-save-mapping branch 5 times, most recently from 5a70522 to 19d9f5c Compare October 23, 2026 13:35
…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
@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: muse_spark

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants