Skip to content

Allow parameterless FSDP root modules - #5711

Merged
wujingyue merged 3 commits into
NVIDIA:mainfrom
wujingyue:fsdp-empty-root
Jul 17, 2026
Merged

Allow parameterless FSDP root modules#5711
wujingyue merged 3 commits into
NVIDIA:mainfrom
wujingyue:fsdp-empty-root

Conversation

@wujingyue

@wujingyue wujingyue commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Summary

  • allow experimental FSDP modules to own zero unowned parameters
  • add a training regression for a parameterless root wrapping child FSDP units

Why

FSDP overlap coordination is scoped within a root FSDP module, matching the PyTorch FSDP2 model. That means users and tests will often wrap higher-level container modules to define the root scope, even when all real parameters are already owned by child FSDP units.

The rejected-empty-parameter check was originally added as a sanity check. With root-scoped overlap, that sanity check is too restrictive: an empty root can be a valid container for child FSDP units. Allowing it avoids forcing artificial root-owned parameters just to establish that scope, which keeps unit tests simpler and makes the expected higher-level wrapping pattern work naturally.

Testing

  • uv run --no-sync isort megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
  • uv run --no-sync black --skip-magic-trailing-comma --skip-string-normalization --check --diff megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
  • uv run --no-sync pylint megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
  • git diff --check
  • uv run --no-sync python -m torch.distributed.run --nproc-per-node 1 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py::test_parameterless_parent_with_child_units_trains --tb=short
  • uv run --no-sync python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py::test_parameterless_parent_with_child_units_trains --tb=short

@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.

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 13, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

# Conflicts:
#	megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py
@wujingyue
wujingyue enabled auto-merge July 17, 2026 03:00
@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 17, 2026
@wujingyue
wujingyue added this pull request to the merge queue Jul 17, 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/29555614995

@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 17, 2026
@wujingyue
wujingyue added this pull request to the merge queue Jul 17, 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/29560018176

Merged via the queue into NVIDIA:main with commit 61f3114 Jul 17, 2026
182 checks passed
@wujingyue
wujingyue deleted the fsdp-empty-root branch July 17, 2026 07:26
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