Skip to content

Move FSDP model weight sync to optimizer post-step#5949

Merged
wujingyue merged 1 commit into
NVIDIA:mainfrom
wujingyue:post-step
Jul 23, 2026
Merged

Move FSDP model weight sync to optimizer post-step#5949
wujingyue merged 1 commit into
NVIDIA:mainfrom
wujingyue:post-step

Conversation

@wujingyue

@wujingyue wujingyue commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Move experimental FSDP main-weight to model-weight synchronization from forward hooks to the optimizer post-step hook, so it runs once per optimizer step instead of once per microbatch.
  • Add precision_aware=True for optimizers, such as FusedAdam, that accept mixed-precision gradients without temporary casts.
  • Initialize compute weights at FSDP setup and add multi-step synchronization coverage.

Validation

  • python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py
  • Focused FSDP post-step synchronization, next-forward freshness, and mixed-precision Adam parity tests.
  • Black, isort, Ruff, and git diff --check.

@copy-pr-bot

copy-pr-bot Bot commented Jul 21, 2026

Copy link
Copy Markdown

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: Jingyue Wu <jingyuew@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Jul 23, 2026

@wujingyue wujingyue left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From @shjwudp : for drop-in replacement and potentially CPU offloading, we may want to sync in hooks instead of post optimizer step. Revisit this when there's a clearer use case.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Jul 23, 2026
@wujingyue
wujingyue added this pull request to the merge queue Jul 23, 2026
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30019192213

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 23, 2026
@wujingyue
wujingyue added this pull request to the merge queue Jul 23, 2026
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30029209950

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30030913542

Merged via the queue into NVIDIA:main with commit 15d4b34 Jul 23, 2026
97 of 99 checks passed
@wujingyue
wujingyue deleted the post-step branch July 23, 2026 19:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: low Run tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants