From dc05b8b978f83d3a4c9cbc698cebbe2a4694ce6c Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 12 Jun 2026 19:05:53 +0000 Subject: [PATCH 1/3] Add experimental FSDP meta parameter support Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/module.py | 32 +++++++++++++++++-- .../experimental/parameter_group.py | 5 +++ .../distributed/mfsdp_v2/test_fully_shard.py | 30 +++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index 33e6985d8ed..ac689193be6 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -103,7 +103,7 @@ def __init__( self._context = None self._name = None self._unshard_event = None - owned_parameters = _collect_owned_parameters(self) + owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) assert axis_indices == tuple( range(mesh.ndim) @@ -373,12 +373,40 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: return dim_names.index(axis) -def _collect_owned_parameters(root_module: nn.Module) -> dict[str, nn.Parameter]: +def _mesh_device(mesh: DeviceMesh) -> torch.device: + if mesh.device_type == "cuda": + return torch.device("cuda", torch.cuda.current_device()) + return torch.device(mesh.device_type) + + +def _materialize_and_collect_owned_parameters( + root_module: nn.Module, device: torch.device +) -> dict[str, nn.Parameter]: parameters: dict[str, nn.Parameter] = {} def visit(submodule: nn.Module, submodule_fqn: str) -> None: direct_parameters = list(submodule.named_parameters(recurse=False)) + if any(parameter.is_meta for _, parameter in direct_parameters): + if any(not parameter.is_meta for _, parameter in direct_parameters): + raise ValueError( + f"Module {submodule_fqn!r} mixes meta and non-meta direct parameters. " + "Initialize all direct parameters on meta or none of them." + ) + submodule.to_empty(device=device, recurse=False) + with torch.no_grad(): + if hasattr(submodule, "reset_parameters"): + submodule.reset_parameters() + elif hasattr(submodule, "_reset_parameters"): + submodule._reset_parameters() + else: + raise ValueError( + f"Module {submodule_fqn!r} does not have " + "reset_parameters or _reset_parameters." + ) + # Module.to_empty may replace Parameters, so collect direct parameters again. + direct_parameters = list(submodule.named_parameters(recurse=False)) + for local_parameter_name, parameter in direct_parameters: parameter_fqn = ( f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index fd5eb5d2033..472a43632e5 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -86,6 +86,11 @@ def __init__( self.dtype = first_parameter.dtype self.requires_grad = first_parameter.requires_grad for name, parameter in parameters.items(): + if parameter.is_meta: + raise ValueError( + f"Expected parameter {name!r} to be materialized before " + "ParameterGroup construction." + ) if parameter.dtype != self.dtype: raise ValueError( f"Expected parameter {name!r} to have dtype {self.dtype}, " diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 510690d09ec..c1e0dbdfb03 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -104,6 +104,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return SaveNonLeafWeightView.apply(x, weight_view) +class ConstantMetaModel(nn.Module): + """Model whose meta parameter is initialized by reset_parameters().""" + + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) + + def reset_parameters(self) -> None: + """Initialize the weight to a deterministic value.""" + self.weight.fill_(3.0) + + def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) @@ -838,6 +850,24 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): torch.testing.assert_close(output, expected_output) +def test_meta_parameters_initialize_with_reset_parameters(distributed_setup): + """Meta parameters should be replaced by sharded DTensors and initialized in place.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + + mesh = init_device_mesh(device.type, (world_size,)) + model = ConstantMetaModel() + + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + group = model.parameter_groups[0] + full_weight = group.model_weight.allgather(0).get_local_tensor(0) + assert not full_weight.is_meta + torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.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 From 73984c2a81992bdd4c50f2db036c7bb44d4dd4fd Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 29 Jul 2026 21:16:26 +0000 Subject: [PATCH 2/3] Allow MFSDP v2 to shard meta parameters Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/dbuffer.py | 10 +++-- .../src/megatron_fsdp/experimental/module.py | 32 +--------------- .../experimental/parameter_group.py | 22 ++++++----- .../distributed/mfsdp_v2/test_fully_shard.py | 38 +++++++++---------- 4 files changed, 38 insertions(+), 64 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 8381f9a3a5c..c6344955822 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/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( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py index ac689193be6..33e6985d8ed 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -103,7 +103,7 @@ def __init__( self._context = None self._name = None self._unshard_event = None - owned_parameters = _materialize_and_collect_owned_parameters(self, _mesh_device(mesh)) + owned_parameters = _collect_owned_parameters(self) axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes) assert axis_indices == tuple( range(mesh.ndim) @@ -373,40 +373,12 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: return dim_names.index(axis) -def _mesh_device(mesh: DeviceMesh) -> torch.device: - if mesh.device_type == "cuda": - return torch.device("cuda", torch.cuda.current_device()) - return torch.device(mesh.device_type) - - -def _materialize_and_collect_owned_parameters( - root_module: nn.Module, device: torch.device -) -> dict[str, nn.Parameter]: +def _collect_owned_parameters(root_module: nn.Module) -> dict[str, nn.Parameter]: parameters: dict[str, nn.Parameter] = {} def visit(submodule: nn.Module, submodule_fqn: str) -> None: direct_parameters = list(submodule.named_parameters(recurse=False)) - if any(parameter.is_meta for _, parameter in direct_parameters): - if any(not parameter.is_meta for _, parameter in direct_parameters): - raise ValueError( - f"Module {submodule_fqn!r} mixes meta and non-meta direct parameters. " - "Initialize all direct parameters on meta or none of them." - ) - submodule.to_empty(device=device, recurse=False) - with torch.no_grad(): - if hasattr(submodule, "reset_parameters"): - submodule.reset_parameters() - elif hasattr(submodule, "_reset_parameters"): - submodule._reset_parameters() - else: - raise ValueError( - f"Module {submodule_fqn!r} does not have " - "reset_parameters or _reset_parameters." - ) - # Module.to_empty may replace Parameters, so collect direct parameters again. - direct_parameters = list(submodule.named_parameters(recurse=False)) - for local_parameter_name, parameter in direct_parameters: parameter_fqn = ( f"{submodule_fqn}.{local_parameter_name}" if submodule_fqn else local_parameter_name diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 472a43632e5..e40b9ae50d7 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -86,11 +86,6 @@ def __init__( self.dtype = first_parameter.dtype self.requires_grad = first_parameter.requires_grad for name, parameter in parameters.items(): - if parameter.is_meta: - raise ValueError( - f"Expected parameter {name!r} to be materialized before " - "ParameterGroup construction." - ) if parameter.dtype != self.dtype: raise ValueError( f"Expected parameter {name!r} to have dtype {self.dtype}, " @@ -162,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) @@ -177,8 +183,6 @@ def __init__( self.sharded_parameters = tuple(sharded_parameters) self.unsharded_parameters = tuple(unsharded_parameters) - # Compute weights must be initialized before the first forward; subsequent - # refreshes happen from the FSDP optimizer's post-step hook. self.sync_model_weight_from_main_weight() self._switch_to_sharded_parameters() self._unsharded_model_weight.release_storage() diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index c1e0dbdfb03..7797fdf02f1 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -104,18 +104,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return SaveNonLeafWeightView.apply(x, weight_view) -class ConstantMetaModel(nn.Module): - """Model whose meta parameter is initialized by reset_parameters().""" - - def __init__(self) -> None: - super().__init__() - self.weight = nn.Parameter(torch.empty(4, 4, device="meta")) - - def reset_parameters(self) -> None: - """Initialize the weight to a deterministic value.""" - self.weight.fill_(3.0) - - def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) @@ -850,22 +838,30 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): torch.testing.assert_close(output, expected_output) -def test_meta_parameters_initialize_with_reset_parameters(distributed_setup): - """Meta parameters should be replaced by sharded DTensors and initialized in place.""" +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 - if world_size < 2: - pytest.skip("This test requires at least 2 ranks.") mesh = init_device_mesh(device.type, (world_size,)) - model = ConstantMetaModel() + 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()) - group = model.parameter_groups[0] - full_weight = group.model_weight.allgather(0).get_local_tensor(0) - assert not full_weight.is_meta - torch.testing.assert_close(full_weight, torch.full_like(full_weight, 3.0)) + 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): From 236eb596687762c17a4b4079a272f655f4dc3b6e Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 30 Jul 2026 03:25:08 +0000 Subject: [PATCH 3/3] Document initial MFSDP weight synchronization Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/parameter_group.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index e40b9ae50d7..a4fabfd308f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -183,6 +183,8 @@ def __init__( self.sharded_parameters = tuple(sharded_parameters) self.unsharded_parameters = tuple(unsharded_parameters) + # Compute weights must be initialized before the first forward; subsequent + # refreshes happen from the FSDP optimizer's post-step hook. self.sync_model_weight_from_main_weight() self._switch_to_sharded_parameters() self._unsharded_model_weight.release_storage()