Skip to content

Fix tied parameters - #6164

Merged
wujingyue merged 4 commits into
NVIDIA:mainfrom
wujingyue:codex/repro/mfsdp-tied-weights
Aug 1, 2026
Merged

Fix tied parameters#6164
wujingyue merged 4 commits into
NVIDIA:mainfrom
wujingyue:codex/repro/mfsdp-tied-weights

Conversation

@wujingyue

@wujingyue wujingyue commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

What changed

  • collect MFSDP v2 parameter registrations without deduplicating names, while representing each physical Parameter once
  • retain every owning-module-relative FQN and use the complete tuple when swapping sharded and unsharded parameter replacements
  • register one gradient hook per physical trainable parameter and track backward readiness with a resettable counter
  • add a tied-embedding regression test covering two forward/backward cycles, balanced annotations, DBuffer size, and optimizer-visible parameter count

Why

PyTorch's default parameter enumeration deduplicates tied parameters by object identity. MFSDP v2 instead collected registrations by FQN, so the same physical parameter could receive multiple buffer entries and replacement objects. That broke weight identity and prevented the ready-parameter set from reaching the duplicated trainable-parameter count during backward.

Impact

MFSDP v2 now preserves tied parameter aliases while allocating, sharding, reducing, and exposing only one physical optimizer parameter.

Validation

  • ruff check on all changed Python files
  • Black check on all changed Python files
  • Python bytecode compilation on all changed Python files
  • python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2
    • 53 passed
    • 11 skipped because they require 4-5 ranks

Related issue

Related to #6139.

@copy-pr-bot

copy-pr-bot Bot commented Jul 30, 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 codex/repro/mfsdp-tied-weights branch 3 times, most recently from c135937 to 14fc287 Compare July 30, 2026 23:49
@wujingyue wujingyue changed the title Fix MFSDP v2 tied parameters Fix tied parameters Jul 30, 2026
@wujingyue
wujingyue force-pushed the codex/repro/mfsdp-tied-weights branch 12 times, most recently from d988824 to 1dd1cc2 Compare July 31, 2026 00:40
@wujingyue
wujingyue marked this pull request as ready for review July 31, 2026 00:43
@wujingyue
wujingyue requested review from a team as code owners July 31, 2026 00:43
@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/30681159582

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Aug 1, 2026
@wujingyue
wujingyue force-pushed the codex/repro/mfsdp-tied-weights branch from 6d0e97e to f2c6df9 Compare August 1, 2026 04:55
@copy-pr-bot

copy-pr-bot Bot commented Aug 1, 2026

Copy link
Copy Markdown

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.

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@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/30688416365

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Aug 1, 2026
Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
@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/30720670220

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: medium Run functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Integrate NeMo Automodel Qwen3.5-35B-A3B recipe with MFSDP v2

5 participants