Repository navigation
Conversation
Split from NVIDIA#3791; original changes by @parthmannan. Plumb packed_seq_params through the transformer layer MLP call sites, MoELayer.forward/route, and the TopKRouter so that, under dynamic (hybrid) context parallelism, aux and seq-aux losses are reduced over the per-sequence CP sub-group followed by the TP group instead of the static TP-CP group, with a logging-only scale correcting the sub-group overcounting in the reduced metrics. Extend get_tokens_per_expert_and_token_count to accept a list of process groups and all-reduce sequentially over them. Clamp the hybrid CP scheduler's gpus_needed to a minimum CP size of EP/TP so expert ranks stay in sync, and scale MoE/MTP logged metrics in training_log by the hybrid-CP per-iteration group count when available. Co-authored-by: Parth Mannan <pmannan@nvidia.com> Signed-off-by: ilml <tolong@nvidia.com>
Contributor
Author
|
/ok to test 78d7e90 |
This was referenced Sep 8, 2026
3 of 5 tasks
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
MoE under dynamic CP:
packed_seq_paramsis plumbed throughTransformerLayer→MoELayer→ router soTopKRouterreduces aux/seq-aux losses over the runtime[cp_group, tp_group];get_tokens_per_expert_and_token_countaccepts a list of process groups;gpus_needed()clamps the minimum CP size for expert-rank sync (HACK/TODO preserved from the original);training_logscales MoE/MTP metrics by the hybrid-CP group count with a fallback toget_num_microbatches()so this PR is independently mergeable before the packing PR.Series (split of #3791)
Split of #3791 ("Update Dynamic CP to support sequence packing, MXFP8 and Mamba") against current
main, tracked in the Dynamic Context Parallelism project. Original changes by @parthmannan in #3791; hunks already covered by main (via the merged sequence-packing series and hybrid-CP fixes) were dropped, and the rest re-implemented against main's current structure. All PRs targetmain; merge order:Contribution process