[main] add thd sequence packing dispatcher support for main by HaochenYuan · Pull Request #5008 · NVIDIA/Megatron-LM

HaochenYuan · 2026-05-27T08:08:57Z

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

What does this PR do ?

Corresponding main branch PR for dev branch PR 4816 & PR 5048.

Previous dev PR 2924 added support for thd format in training, but for MoE dispatcher, only all2all type is supported. This PR adds the deepep & hybridep backend.

While currently the thd e2e support PR 3386 is not merged, this PR will remain draft state until 3386 proceeds.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

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"

When your PR is ready, click Ready for Review.
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.

github-actions · 2026-05-27T08:09:10Z

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:

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

See the contribution guide for more details.

copy-pr-bot · 2026-05-27T08:09:12Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Signed-off-by: HaochenYuan <haocheny@nvidia.com>

jiemingz · 2026-07-20T20:25:17Z

    moe_permute_fusion_into_hybridep: bool = False
    """Fuse token rearrangement ops during token dispatching for HybridEP."""

+    moe_hybridep_pad_variable_tokens: bool = False


Is this feature only intended to be used for sequence packing? If so it feels awkward that if a user were to enable a sequence packed dataset that produced uneven sequence lengths, they would also need to know about this knob (which doesnt mention anything about sequence packing)

I think we generally want to avoid creating new configs when we could otherwise reasonably infer them from elsewhere. Could we set this path by doing something like detecting the presence of packed_seq_params by plumbing it down the moe_layer?

Yes we expect users to turn this toggle on specifically for HybridEP + THD and we will guide our customers when should we turn this on. On the other hand, having packed sequence doesn't always mean we needs to pad if we apply offline packing and the dataset was created without the need to pad.

However, the accuracy impact of doing offline packing & memory overhead concerns of always padding to max seq-len are often key factors for us to introduce this specific toggle to do minimum amount of padding while also being able to use HybridEP.

can we rename this and update the description to reflect this reasoning? maybe moe_hybridep_pad_uneven_sequence_packed_inputs ?

I agree that documentation might need to be updated, but moe_hybridep_pad_variable_tokens might be a better name because this is a general enough feature to handle dynamic sized inputs across EP ranks, while technically sequence packed inputs is orthogonal, because we can have the following two cases:

if I choose not to do seq-packing, but processing single document (accuracy concerns about offline seqpacking might lead to this), I will need this toggle to do minimum amount of padding.

if I choose to do seq-packing, there are cases where my data distribution is good enough that it's actually acceptable to pad to MBS * seqlen because the utilization rate is very high, so that we don't really need to pad right before HybridEP, for example Nemotron RL seems to this case.
Overall, I avoided to bring in seq-packing & thd into the naming in the first place.

then maybe: moe_hybridep_pad_uneven_dispatch_inputs ? to disambiguate we're padding the input tokens which vary across ranks

I am okay with this one

Signed-off-by: HaochenYuan <haocheny@nvidia.com>

svcnvidia-nemo-ci · 2026-07-23T23:01:58Z