[DEV] fix(megatron-fsdp): reduce padding for grouped expert weights

[xuwchen]

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

main PR: #4979

What does this PR do ?

MFSDP computes a chunk size factor (CSF) for each bucket as shape[1:].numel(), which flattens all dimensions except the first one.

For per-expert 2D expert weights:

• linear_fc1: (2 * moe_ffn_hidden_size, hidden_size)
• linear_fc2: (hidden_size, moe_ffn_hidden_size)

shape[1:].numel() is just the last dimension, so the CSF stays small.

For grouped 3D expert weights:

• linear_fc1: (num_local_experts, 2 * moe_ffn_hidden_size, hidden_size)
• linear_fc2: (num_local_experts, hidden_size, moe_ffn_hidden_size)

shape[1:].numel() becomes the full per-expert matrix size. This can make the CSF much larger than the equivalent per-expert 2D layout. When multiple expert weights share a bucket, MFSDP uses divisibility/LCM logic to choose a common CSF. The oversized CSF can force the bucket size to be padded to a much larger alignment unit, increasing AllGather traffic. In the reported configuration this matched the observed ~33% communication increase, corresponding to ~25% bucket padding.

This PR routes grouped expert weights with heterogeneous CSFs into separate buckets via a new _should_split_from_grouped_expert_bucket helper, while keeping the 2D / non-expert paths unchanged.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[wujingyue]

Could you add a unit test? These algorithmic changes are easy to regress without test coverage guarding them. The unit test can be something like https://github.com/NVIDIA/Megatron-LM/pull/4835/changes#diff-f1aaffa52ab9eab0c69292e44285ebb8c9387f1c3ebd62191d82920b5817cccbR110 and doesn't have to run FSDP end-to-end.

[wujingyue]

LGTM otherwise

[wujingyue]

Suggested change

	def _get_csf_base(group: ParameterGroup, param: torch.nn.Parameter) -> int:
	def _get_chunk_size_factor_base(group: ParameterGroup, param: torch.nn.Parameter) -> int:

Also, you only need an is_expert_param boolean instead of the entire ParameterGroup.

[xuwchen]

Addressed in #5013 (the continuation PR). Variable naming uses the full param_chunk_size_factor. _get_csf_base is also removed and replaced by _should_split_from_grouped_expert_bucket, which takes is_expert_param: bool.

[xuwchen]

Continued in #5013. The original PR couldn't be reopened due to a force-push lock. The review feedback has been addressed in the new PR.