Skip to content

Pair frozen FSDP backward hooks - #5710

Merged
wujingyue merged 3 commits into
NVIDIA:mainfrom
wujingyue:nvtx-frozen-backward
Jul 16, 2026
Merged

Pair frozen FSDP backward hooks#5710
wujingyue merged 3 commits into
NVIDIA:mainfrom
wujingyue:nvtx-frozen-backward

Conversation

@wujingyue

@wujingyue wujingyue commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Rename the FSDP backward-completion counter from training to trainable parameters.
  • Pair frozen-only FSDP units with a module full-backward hook so post_backward() runs after pre_backward().
  • Add a regression test that exposes the frozen-parameter NVTX imbalance from Add FSDP NVTX annotations #5704 review feedback.
  • Move the annotation regression test into the mfsdp_v2 unit-test bucket, where the experimental FSDP v2 fixture lives.

Validation

  • python -m py_compile megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py tests/unit_tests/distributed/mfsdp_v2/test_annotation.py
  • git diff origin/main...HEAD --check

@copy-pr-bot

copy-pr-bot Bot commented Jul 8, 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.

@wujingyue
wujingyue force-pushed the nvtx-frozen-backward branch from 487ddcb to 47b4f3e Compare July 8, 2026 17:02
@wujingyue
wujingyue marked this pull request as ready for review July 8, 2026 17:02
@wujingyue
wujingyue requested review from a team as code owners July 8, 2026 17:02
@wujingyue
wujingyue changed the base branch from pull-request/5704 to main July 8, 2026 17:14
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 14, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@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 15, 2026
@wujingyue
wujingyue added this pull request to the merge queue Jul 15, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to a conflict with the base branch Jul 15, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

# Conflicts:
#	megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py
@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/29468459658

Merged via the queue into NVIDIA:main with commit 38626c2 Jul 16, 2026
181 checks passed
@wujingyue
wujingyue deleted the nvtx-frozen-backward branch July 16, 2026 04:33
chochowski pushed a commit to chochowski/Megatron-LM that referenced this pull request Jul 20, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: mchochowski <mchochowski@nvidia.com>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
svcnvidia-nemo-ci pushed a commit to dimapihtar/Megatron-LM that referenced this pull request Aug 4, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
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