Skip to content

[codex] Add experimental FSDP meta-parameter support - #2

Closed
wujingyue wants to merge 1 commit into
fsdp/minimalfrom
meta
Closed

[codex] Add experimental FSDP meta-parameter support#2
wujingyue wants to merge 1 commit into
fsdp/minimalfrom
meta

Conversation

@wujingyue

Copy link
Copy Markdown
Owner

This is intentionally separated from the minimal FSDP cleanup in the current stacked branch so the meta-parameter support can be kept, revised, or dropped independently.

What changed:

  • Restores experimental fully_shard handling for direct meta parameters by materializing with to_empty(...) and invoking reset_parameters() / _reset_parameters() before ParameterGroup construction.
  • Restores the distributed test that verifies meta parameters are replaced by initialized sharded DTensors.
  • Keeps the CPU-initialized real-parameter coverage added in the base branch.
  • Does not touch non-experimental FSDP files.

Validation:

  • CUDA_VISIBLE_DEVICES= uv run --no-sync python -m torch.distributed.run --nproc-per-node 8 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py (8 passed, 1 skipped)

Notes:

  • The plain requested uv run ... command could not run in this session because uv attempted to modify /opt/venv and hit a permission error.
  • Running with CUDA visible could not complete on this host because it has 2 GPUs while the test command launches 8 local ranks.

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
@wujingyue

Copy link
Copy Markdown
Owner Author

Moved to upstream PR NVIDIA#5369.

@wujingyue wujingyue closed this Jun 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant