Add meta-parameter support

[wujingyue]

What does this PR do?

Adds meta-parameter support to MFSDP v2. fully_shard() can now accept meta parameters, allocate their sharded storage directly on the device mesh, and leave their values for the caller to initialize or checkpoint-load.

This helps with #6139

Why is this useful?

MFSDP v2 already supports models whose weights are initialized on CPU, but that requires the full unsharded model weights to occupy host memory before sharding. For large models, CPU memory can be precious and may become a startup bottleneck. Constructing the model on the meta device retains only parameter metadata such as shapes and dtypes, allowing MFSDP v2 to create the real sharded storage without first materializing a full CPU copy.

Summary

• Allocate real local DBuffer storage on the mesh device from tensor shapes and dtypes.
• Skip copying source values for meta tensors, leaving their sharded ranges uninitialized.
• Materialize each original meta Parameter onto its real unsharded buffer with torch.utils.swap_tensors() because .data cannot rebind a meta TensorImpl to a real device TensorImpl.
• Preserve the existing constructor-time main-weight-to-model-weight synchronization; callers must synchronize again after initializing or loading exposed main weights.
• Do not rely on to_empty() or reset_parameters() / _reset_parameters(); fully_shard() materializes storage without invoking user initialization hooks.
• Cover a full BF16 meta model: shard it, fill its exposed weights, simulate checkpoint post-load synchronization, and verify the forward result.

Initialization semantics

I chose not to rely on reset_parameters() because it is not guaranteed to be available for every nn.Module. Automatically initializing only modules that happen to provide a reset hook would also be confusing: fully_shard() would initialize some modules while leaving others uninitialized. Instead, fully_shard() consistently materializes storage without choosing parameter values, and the caller explicitly initializes or checkpoint-loads them.

Checkpoint loading

The test explicitly synchronizes every parameter group's main weights to model weights after filling them. This simulates the post-load behavior being added by #6024; once that PR lands, its checkpoint loader will perform the synchronization.

Testing

• python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py::test_meta_parameters_shard_to_mesh_device
  • 1 passed per rank
• python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2
  • 52 passed, 11 topology-dependent skips per rank
• Ruff and Black 24 checks on the changed Python files
• git diff --check

[copy-pr-bot]

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.

[copy-pr-bot]

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.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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