Integrate MFSDP v2 with fine-grained 1F1B scheduling by shjwudp · Pull Request #6692 · NVIDIA/Megatron-LM

shjwudp · 2026-08-20T02:02:38Z

Summary

Part 2 of the MFSDP v2 combined-1F1B split. This PR contains the MCore adapter and models/common schedule integration; the generic experimental lifecycle remains in #6484.

infer MFSDP custom scheduling from overlap_moe_expert_parallel_comm
install fine-grained demand-unshard hooks with nearest-FSDP-owner tracking and idempotent setup
keep forward resharding explicit and let the parameter countdown finalize backward reduction
add MFSDP v2 no_sync() handling and wait for the final reduce-scatter before gradient consumers
add focused adapter tests and 20-step MFSDP-v2/DistOpt overlap parity coverage

Dependency

Depends on #6484 and is based on its verified pull-request/6484 mirror. Retarget this PR to main before merging; merging while based on the mirror would land into the bot scratch branch instead of main.

Validation

Black 26.3.0: passed
isort 5.13.2: passed
Ruff 0.9.10, including explicit undefined/unused-symbol checks: passed
git diff --check: passed
pre-split 20-step, 2-GPU MFSDP-v2/DistOpt parity: passed; final LM loss 4.30552197 vs. 4.30558395
GPU rerun on current main: pending cluster MFA availability

copy-pr-bot · 2026-08-20T02:02:41Z

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.

shjwudp · 2026-08-20T02:14:08Z

/ok to test 33627e2

shjwudp · 2026-08-20T05:14:07Z

/ok to test 0a0f601

shjwudp · 2026-08-20T05:19:49Z

/ok to test 1fdb108

Clarify the API contract: the callback runs gradient reduction once all trainable parameters have accumulated their gradients, not via a module-level full-backward post hook (which can fire at the wrong time, e.g. with TE delayed wgrad). Addresses review feedback on NVIDIA#6484. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

pre_backward() runs from the registered backward hook or the 1F1B schedule, neither of which passes a value, so the knob is dead. Restore the unconditional context.register_post_backward_final_callback() call. Addresses NVIDIA#6484 review feedback. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

The range_push was moved after context=self.context during the life-cycle refactor; it is a no-op reorder, so put it back to keep the diff focused. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Rename _post_backward_hook_registered to _grad_reduction_callback_registered to match the renamed register_grad_reduction_callback API, and document that external schedulers calling unshard() directly must perform the root all-gather-stream sync that pre_forward() normally does first. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…rop duplicate-guard - Keep the public method name register_post_backward_hook (the callback also reshards; rename to grad_reduction_callback was misleading). - Drop the _post_backward_hook_registered state so FsdpModule can accept multiple hooks and let the fine-grained scheduler do its own checking. - Comment why register_load_state_dict_pre_hook is registered unconditionally. - Rename _lifecycle_events to _forward_backward_events. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Wrap the unshard/reshard/gradient_reduce nvtx push/pop in a context manager so an early return cannot skip the matching pop. forward/backward span multiple lifecycle methods and keep their explicit push/pop. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Drop tests that assert on private attributes/internal methods (behavior- not-implementation): test_fully_shard_can_disable_execution_hooks, test_register_post_backward_hook_waits_for_all_parameter_gradients, and test_multiple_forwards_before_backwards_reset_gradient_readiness. register_hooks and the fine-grained scheduler get their coverage in the upcoming 1F1B PRs. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…IA#6484 lifecycle) Port the NVIDIA#6692 fine-grained scheduler integration on top of split/pr-a-overlap (a2ba13b), dropping the removed custom_schedule lifecycle and relying on the NVIDIA#6484 MFSDP v2 API (register_post_backward_hook, idempotent unshard()/reshard()). - Add fine_grained_mfsdp_scheduler.py (setup_combined_1f1b_hooks, submodule unshard/reshard hooks, reshard_fsdp_module) with no custom_schedule. - Wire it in model_chunk_schedule_plan.py. - Add end-to-end ND-parallel overlap parity test. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…hedule_plan.py Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…IA#6484 lifecycle) Port the NVIDIA#6692 fine-grained scheduler integration on top of split/pr-a-overlap (a2ba13b), dropping the removed custom_schedule lifecycle and relying on the NVIDIA#6484 MFSDP v2 API (register_post_backward_hook, idempotent unshard()/reshard()). - Add fine_grained_mfsdp_scheduler.py (setup_combined_1f1b_hooks, submodule unshard/reshard hooks, reshard_fsdp_module) with no custom_schedule. - Wire it in model_chunk_schedule_plan.py. - Add end-to-end ND-parallel overlap parity test. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…IA#6484 lifecycle) Port the NVIDIA#6692 fine-grained scheduler integration on top of split/pr-a-overlap (a2ba13b), dropping the removed custom_schedule lifecycle and relying on the - Add fine_grained_mfsdp_scheduler.py (setup_combined_1f1b_hooks, submodule unshard/reshard hooks, reshard_fsdp_module) with no custom_schedule. - Wire it in model_chunk_schedule_plan.py. - Add end-to-end ND-parallel overlap parity test. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

shjwudp mentioned this pull request Aug 20, 2026

Document MFSDP v2 1F1B overlap design #6693

Closed

shjwudp added the Run functional tests label Aug 20, 2026

copy-pr-bot Bot temporarily deployed to public August 20, 2026 02:14 Inactive

copy-pr-bot Bot temporarily deployed to test August 20, 2026 02:15 Inactive

shjwudp mentioned this pull request Aug 20, 2026

Integrate MFSDP v2 with 1F1B overlap #6694

Closed

copy-pr-bot Bot temporarily deployed to public August 20, 2026 02:17 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 02:18 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 02:30 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:14 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:17 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:18 Inactive

shjwudp force-pushed the split/pr6484-fine-grained-hooks branch from 0a0f601 to 1fdb108 Compare August 20, 2026 05:18

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:20 Inactive

copy-pr-bot Bot temporarily deployed to test August 20, 2026 05:20 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:23 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:24 Inactive

copy-pr-bot Bot temporarily deployed to public August 20, 2026 05:36 Inactive

shjwudp changed the title ~~Integrate MFSDP v2 with the MCore 1F1B lifecycle~~ Integrate MFSDP v2 with fine-grained 1F1B scheduling Sep 3, 2026

shjwudp changed the base branch from main to pull-request/6484 September 3, 2026 14:22

shjwudp force-pushed the split/pr6484-fine-grained-hooks branch from e200680 to 51723a8 Compare September 3, 2026 14:34

copy-pr-bot Bot force-pushed the pull-request/6484 branch from 7dcc669 to cff249e Compare September 4, 2026 06:44

shjwudp added 8 commits September 6, 2026 23:39

Restore nvtx push position in pre_backward

d2b3f90

The range_push was moved after context=self.context during the life-cycle refactor; it is a no-op reorder, so put it back to keep the diff focused. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

shjwudp force-pushed the split/pr6484-fine-grained-hooks branch from e5bdfec to 7172f5f Compare September 7, 2026 02:38

shjwudp added 2 commits September 7, 2026 13:22

Register 1f1b fine-grained hooks at FSDP construction stage

7372a95

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

remove duplicate setup_combined_1f1b_hooks register at model_chunk_sc…

17d43d7

…hedule_plan.py Signed-off-by: Jack Chang <jianbinc@nvidia.com>

shjwudp marked this pull request as ready for review September 7, 2026 05:28

shjwudp requested review from a team as code owners September 7, 2026 05:28

github-actions Bot added the complexity: medium label Sep 7, 2026

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

simply unit test

8a3a4e7

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

shjwudp marked this pull request as draft September 7, 2026 12:22

copy-pr-bot Bot deleted the branch NVIDIA:pull-request/6484 September 7, 2026 12:54

copy-pr-bot Bot closed this Sep 7, 2026

shjwudp mentioned this pull request Sep 7, 2026

Integrate MFSDP v2 with fine-grained 1F1B scheduling #7112

Merged

Navigation Menu

Repository navigation

Integrate MFSDP v2 with fine-grained 1F1B scheduling - #6692

Integrate MFSDP v2 with fine-grained 1F1B scheduling#6692
shjwudp wants to merge 11 commits into
NVIDIA:pull-request/6484from
shjwudp:split/pr6484-fine-grained-hooks

shjwudp commented Aug 20, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 20, 2026

Uh oh!

shjwudp commented Aug 20, 2026

Uh oh!

shjwudp commented Aug 20, 2026

Uh oh!

shjwudp commented Aug 20, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

shjwudp commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependency

Validation

Uh oh!

copy-pr-bot Bot commented Aug 20, 2026

Uh oh!

shjwudp commented Aug 20, 2026

Uh oh!

shjwudp commented Aug 20, 2026

Uh oh!

shjwudp commented Aug 20, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

shjwudp commented Aug 20, 2026 •

edited

Loading