From 3a61d62d3bab57a826e77491012ed670895ba990 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 8 Jul 2026 17:43:38 +0000 Subject: [PATCH 1/2] Allow parameterless FSDP root modules Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/module.py | 11 ++-- .../distributed/mfsdp_v2/test_fully_shard.py | 55 +++++++++++++++---- 2 files changed, 51 insertions(+), 15 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 3c97fda3242..5e6bfa4ce3a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -147,8 +147,13 @@ def _lazy_init_context(self) -> None: if self._context is not None: return - context = FsdpContext(device=self._parameter_groups[0].main_weight.device, root_module=self) - for submodule_name, submodule in cast(nn.Module, self).named_modules(): + module = cast(nn.Module, self) + first_parameter = next(module.parameters(), None) + if first_parameter is None: + raise RuntimeError("FSDP root module requires at least one parameter in its subtree.") + + context = FsdpContext(device=first_parameter.device, root_module=self) + for submodule_name, submodule in module.named_modules(): if not isinstance(submodule, FsdpModule): continue if submodule._context is not None: @@ -306,8 +311,6 @@ def visit(submodule: nn.Module, submodule_fqn: str) -> None: visit(child_module, child_fqn) visit(root_module, "") - if not parameters: - raise ValueError("fully_shard requires at least one unowned parameter.") return parameters 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 229cd0bff4b..497cb510ca5 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -386,6 +386,30 @@ def train_one_iteration() -> None: ) +def test_parameterless_parent_with_child_units_trains(distributed_setup): + """A parent with no unowned parameters should still root trainable child FSDP units.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + + mesh = init_device_mesh(device.type, (world_size,)) + torch.manual_seed(5678) + model = nn.Sequential(nn.Linear(4, 4, bias=False), nn.Linear(4, 2, bias=False)).to(device) + + fully_shard(model[0], mesh=mesh, placements=_flat_placements()) + fully_shard(model[1], mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + assert model.parameter_groups == () + + optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + x = torch.randn(3, 4, device=device) + + optimizer.zero_grad(set_to_none=True) + loss = model(x).sum() + loss.backward() + optimizer.step() + + def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): """A non-trainable parameter group should not allocate persistent main gradients.""" world_size = distributed_setup.world_size @@ -485,24 +509,33 @@ def test_microbatch_scopes_child_contexts(distributed_setup): def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): - """CPU-initialized parameters should be sharded with their real values.""" + """A CPU model should support sharding a child before moving the full model to CUDA.""" 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 = nn.Linear(4, 4, bias=False) + model = nn.Sequential( + nn.Linear(4, 4, bias=False), + nn.Linear(4, 4, bias=False), + ) with torch.no_grad(): - model.weight.fill_(3.0) - expected_weight = model.weight.detach().to(device) + model[0].weight.fill_(2.0) + model[1].weight.fill_(3.0) + x = torch.ones(1, 4) + expected_output = model(x).to(device) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + # Shard the second layer's parameters onto the mesh device; the unwrapped + # first layer's parameters remain on CPU until model.to(device) below. + fully_shard(model[1], mesh=mesh, placements=_flat_placements()) - (group,) = model.parameter_groups - full_weight = group.model_weight.allgather(0).get_local_tensor(0) - assert full_weight.device.type == device.type - torch.testing.assert_close(full_weight, expected_weight) + assert model[0].weight.device.type == "cpu" + assert isinstance(model[1].weight, DTensor) + assert model[1].weight.device == device + + model.to(device) + + output = model(x.to(device)) + torch.testing.assert_close(output, expected_output) def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): From 384414f1d667a6e4281ff931613b9075c7a0e2ec Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 13 Jul 2026 17:21:15 +0000 Subject: [PATCH 2/2] Fix FSDP empty root test formatting Signed-off-by: Jingyue Wu --- tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 497cb510ca5..faa59919d3a 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -514,10 +514,7 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): device = distributed_setup.device mesh = init_device_mesh(device.type, (world_size,)) - model = nn.Sequential( - nn.Linear(4, 4, bias=False), - nn.Linear(4, 4, bias=False), - ) + model = nn.Sequential(nn.Linear(4, 4, bias=False), nn.Linear(4, 4, bias=False)) with torch.no_grad(): model[0].weight.fill_(2.0) model[1].weight.fill_(3.0)