Add meta-parameter support

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py

@@ -201,15 +201,17 @@ def from_local(
     def distribute_tensors(
         cls, tensors: Iterable[torch.Tensor], mesh: DeviceMesh, placements: Iterable[Placement]
     ) -> "DBuffer":
-        """Distribute full local tensor values into a DBuffer.
+        """Distribute full local tensors into a DBuffer.
 
         Args:
-            tensors: Full tensor values available on this rank.
+            tensors: Full tensors available on this rank. Meta tensors contribute
+                shape and dtype metadata but no values.
             mesh: Device mesh whose dimensions correspond to ``placements``.
             placements: Per-mesh-axis DBuffer placements.
 
         Returns:
-            A DBuffer whose local storage matches ``placements``.
+            A DBuffer whose real local storage matches ``placements``. Ranges
+            corresponding to meta tensors are left uninitialized.
         """
         tensors = tuple(tensor.detach().contiguous() for tensor in tensors)
         if not tensors:
@@ -232,7 +234,7 @@ def distribute_tensors(
         # observable through get_local_tensor() and can remain unspecified.
         for index, tensor in enumerate(tensors):
             owned_range = buffer._get_owned_range(index)
-            if owned_range is None:
+            if owned_range is None or tensor.is_meta:
                 continue
 
             source_slice = tensor.view(-1).narrow(

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

@@ -157,8 +157,19 @@ def __init__(
         unsharded_parameters: list[nn.Parameter] = []
         main_grad_dtype = self.main_grad.dtype if self.main_grad is not None else None
         for index, parameter in enumerate(parameters.values()):
-            parameter.data = self._unsharded_model_weight.get_local_tensor(index)
-            parameter.grad = None
+            unsharded_tensor = self._unsharded_model_weight.get_local_tensor(index)
+            if parameter.is_meta:
+                # A meta Parameter cannot set .data to a real tensor because their
+                # TensorImpl types are incompatible, so swap in a materialized Parameter.
+                # This may be problematic if attributes from the original Parameter need
+                # to be copied to the unsharded Parameter.
+                materialized_parameter = nn.Parameter(
+                    unsharded_tensor, requires_grad=parameter.requires_grad
+                )
+                torch.utils.swap_tensors(parameter, materialized_parameter)
+            else:
+                parameter.data = unsharded_tensor
+                parameter.grad = None
             setattr(parameter, _CONTAINING_PARAMETER_GROUP_ATTR, self)
             unsharded_parameters.append(parameter)
 

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -838,6 +838,32 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup):
     torch.testing.assert_close(output, expected_output)
 
 
+def test_meta_parameters_shard_to_mesh_device(distributed_setup):
+    """A sharded meta model should support initialization and forward."""
+    world_size = distributed_setup.world_size
+    device = distributed_setup.device
+
+    mesh = init_device_mesh(device.type, (world_size,))
+    model = nn.Sequential(
+        nn.Linear(4, 4, bias=False, device="meta", dtype=torch.bfloat16),
+        nn.Linear(4, 4, bias=False, device="meta", dtype=torch.bfloat16),
+    )
+
+    fully_shard(model, mesh=mesh, placements=_flat_placements())
+
+    with torch.no_grad():
+        model[0].weight.fill_(2.0)
+        model[1].weight.fill_(3.0)
+    # The exposed parameters update FP32 main weights, while forward uses separate BF16
+    # model weights. This simulates load_checkpoint() until
+    # https://github.com/NVIDIA/Megatron-LM/pull/6024 lands and syncs after loading.
+    for parameter_group in model.parameter_groups:
+        parameter_group.sync_model_weight_from_main_weight()
+
+    output = model(torch.ones(1, 4, device=device, dtype=torch.bfloat16))
+    torch.testing.assert_close(output, torch.full_like(output, 96.0))
+
+
 def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup):
     """A non-leaf parameter view saved for backward should survive full-storage resize."""
     world_size = distributed_setup.world_size