[michal2409]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Fix the FC1 out_features of FusedSharedExpertMLP for tensor parallel (TP) > 1. Then TE can fuse the shared expert MLP.

The problem

_make_fused_grouped_swiglu_ops makes the TE GroupedLinear op for FC1 with
out_features=fc1_weight.size(0) * tp_world_size. But fc1_weight is the local TP shard
(op.weight0 = fc1_weight), and te.pytorch.ops.GroupedLinear has no TP mode. So the op must use the local size.

With TP > 1, TE does not fuse GroupedLinear + ScaledSwiGLU + GroupedLinear, because
validate_grouped_mlp_dims needs fc1.out_features == 2 * fc2.in_features.
TE then runs the unfused ops. The unfused ScaledSwiGLU needs one scale for each row,
but the shared expert gives one unit scale. The forward pass stops with:

transformer_engine/pytorch/csrc/extensions/activation.cpp:367 in function scaled_activation_compute:
Assertion failed: scales_tensor.numel() == rows. scaled activation expects one scale per input row

This problem does not occur with TP=1, because then the two sizes are equal.
We found it with DeepSeek-V3 on GB200: TP=2 with sequence parallel, MXFP8,
use_grouped_gemm_for_shared_expert=True, and moe_shared_expert_glu_interleave_size=32.

Changes

• Use out_features=fc1_weight.size(0) for the FC1 op.
• Remove the get_pg_size import, because it has no use now.
• Add test_make_fused_grouped_swiglu_ops_uses_local_shard_sizes_with_tp.
  The existing tests set the TP size to 1, so they did not find the problem.

Tests

• tests/unit_tests/transformer/moe/test_shared_experts.py on one GB200 node (4 GPUs): 23 tests pass.
  test_shared_expert_forward_backward[2-4-alltoall] needs 8 GPUs (TP2 x EP4), so it cannot run there.
• The new test on main without this fix fails with assert 16 == 8.
• black, isort, ruff, and pylint (versions from pyproject.toml) pass on the two changed files.
• End-to-end test with the same one-line change on 26.09-alpha.rc2:
  • Model: a DeepSeek-V3 proxy with 19 layers and 256 experts.
  • System: 16 GB200 nodes, TP2 PP4 VP4 EP16, MXFP8, full-iteration CUDA graph.
  • Without the fix, warmup step 0 stops with the assertion above.
  • With the fix, 10 training steps and the evaluation complete.
  • The eval loss (9.698) agrees with a run that uses the non-fused shared expert (9.696).

Issue tracking

Linked issue: none. Related: #8024 makes the fused path use the TP group of the module for its collectives.
The two changes do not overlap. A TP > 1 run needs both if tp_group is not the global TP group.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• Determinism test for a changed GPU kernel: not applicable, because this PR changes no kernel.
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[gautham-kollu]

/ok to test bc56a09

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40678856079

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40683345995