[refactor] Common combined-1F1B schedule-plan base (1/4 of #4798)#4941
Open
Connor-XY wants to merge 6 commits into
Open
[refactor] Common combined-1F1B schedule-plan base (1/4 of #4798)#4941Connor-XY wants to merge 6 commits into
Connor-XY wants to merge 6 commits into
Conversation
wujingyue
reviewed
May 28, 2026
wujingyue
left a comment
Contributor
There was a problem hiding this comment.
I'll review after expert reviews are done.
Contributor
|
But thanks for breaking down your PR and thanks for the refactor! |
0784750 to
d65a251
Compare
Contributor
Author
|
/claude review |
d65a251 to
d163ed8
Compare
Contributor
Author
|
/claude review |
d163ed8 to
22539e6
Compare
Contributor
Author
|
/claude review |
22539e6 to
c1c3585
Compare
Contributor
Author
|
/ok to test c1c3585 |
b32b394 to
123d858
Compare
Move model-agnostic schedule-plan helpers out of gpt/fine_grained_callables.py
into megatron/core/models/common/ so non-GPT models (HybridStack) can build the
same combined-1F1B / EP-overlap schedule plans. No behavior change for GPT/MTP.
- Add megatron/core/models/common/{utils.py, fine_grained_callables.py,
model_chunk_schedule_plan.py} with the shared abstractions.
- Reduce megatron/core/models/gpt/fine_grained_callables.py to GPT-specific
pieces; reuse the new common base classes.
- Switch GraphableMegatronModule.init_backward_dw_wrapper to import
_BackwardDWWrapper from the new common module.
- Relax combined_forward_backward_step's GPTModel-only assert to a duck-type
on build_schedule_plan so any model implementing it can participate.
Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…unk) Carries over upstream commit f6ea23b from NVIDIA#4798: when a hybrid layer pattern places MTP in a post_process VPP chunk that holds no main HybridStack layers (e.g. trailing pipe before the MTP separator), the EP-overlap schedule never invokes ``_maybe_apply_final_norm`` on the main path, so the unnormalized hidden_states feed straight into the LM head and lm_loss diverges by ~10x. Fix in ``submodule_mtp_pre_dispatch_forward``: run the main decoder's ``final_norm`` just before ``torch.chunk``, gated on ``len(model.decoder.layers) == 0`` and ``isinstance(model, HybridModel)``. The HybridModel import is deferred inside the function so this file does not gain a module-level dependency on the hybrid package — PR 1 remains independently importable. Part 1/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
123d858 to
eb4c325
Compare
Contributor
Author
|
/ok to test eb4c325 |
Contributor
Author
|
/ok to test 21484bc |
Phlip79
approved these changes
Jun 11, 2026
santhnm2
approved these changes
Jun 11, 2026
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
5 participants
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?
Part 1 of 4 splitting #4798 by @Wohox and @Connor-XY into smaller PRs to reduce the number of CODEOWNERS reviewer groups per PR. Original changes by @Wohox and @Connor-XY; this PR is purely the model-agnostic refactor with no behavior change for GPT/MTP.
Summary
Move model-agnostic combined-1F1B schedule-plan helpers out of
gpt/fine_grained_callables.pyintomegatron/core/models/common/so non-GPT models (HybridStack, in part 2) can build the same schedule plans:megatron/core/models/common/utils.py(houses_BackwardDWWrapper, previously ingpt/fine_grained_callables.py),common/fine_grained_callables.py, andcommon/model_chunk_schedule_plan.pywith the shared abstractions.gpt/fine_grained_callables.pyto GPT-specific pieces; reuse the new common base classes.GraphableMegatronModule.init_backward_dw_wrapperto import_BackwardDWWrapperfromcommon.utils.combined_forward_backward_step'sGPTModel-only assert to a duck-type onbuild_schedule_planso any model implementing it can participate in EP overlap.final_normin MTP pre-dispatch when a VPP chunk has no main HybridStack layers (theHybridModelreference is a deferred import inside the function, so this file remains importable without the hybrid package).Why this slice
Touches 5 reviewer groups:
core-adlr,core-nemo,gpt,transformer,pipeline-parallelism. The hybrid- / FSDP- / training-specific reviewers are not needed here.Dependencies
None — this PR is mergeable on its own. Parts 2/3/4 depend on this PR.
Validation
The full integrated change was validated in #4798 (unit tests + DeepSeek-V3 deterministic Transformer/Hybrid baseline + EP-overlap smoke tests). This slice is a refactor with no behavior change for the GPT/MTP path; behavior was confirmed bitwise-identical by @Wohox on EOS 8-node for the carried-over f6ea23b fix (lm_loss bitwise-identical at iter 1 with A2A_OVERLAP=1 vs 0).
Issue tracking
Linked issue: part of #4798.
Pre-checks
🤖 Generated with Claude Code