Add construction-scoped MFSDP fully shard options

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

@@ -15,7 +15,7 @@
 """Experimental Megatron-FSDP implementation."""
 
 from .dbuffer import DBuffer
-from .fully_shard import fully_shard, microbatch
+from .fully_shard import fully_shard, fully_shard_options, microbatch
 from .optimizer import fully_shard_optimizer
 from .placement import Flat, Partial, Placement, Placements, Replicate
 
@@ -27,6 +27,7 @@
     "Placements",
     "Replicate",
     "fully_shard",
+    "fully_shard_options",
     "fully_shard_optimizer",
     "microbatch",
 ]

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

@@ -16,6 +16,8 @@
 
 from collections.abc import Iterator
 from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
 
 from torch import nn
 from torch.distributed import DeviceMesh
@@ -25,12 +27,37 @@
 from .placement import Placements
 
 
+@dataclass(frozen=True)
+class _FullyShardOptions:
+    """Construction options inherited by fully_shard calls in the current context."""
+
+    use_symmetric_memory: bool = False
+
+
+_FULLY_SHARD_OPTIONS = ContextVar("megatron_fsdp_fully_shard_options", default=_FullyShardOptions())
+
+
+@contextmanager
+def fully_shard_options(*, use_symmetric_memory: bool = False) -> Iterator[None]:
+    """Configure options for ``fully_shard`` calls in this scope.
+
+    Args:
+        use_symmetric_memory: Allocate communication staging buffers from PyTorch's
+            NCCL symmetric-memory pool.
+    """
+    options = _FullyShardOptions(use_symmetric_memory=use_symmetric_memory)
+    token = _FULLY_SHARD_OPTIONS.set(options)
+    try:
+        yield
+    finally:
+        _FULLY_SHARD_OPTIONS.reset(token)
+
+
 def fully_shard(
     module: nn.Module,
     mesh: DeviceMesh,
     placements: Placements,
     mixed_precision_policy: MixedPrecisionPolicy | None = None,
-    use_symm_mem: bool = False,
 ) -> None:
     """Apply FSDP to a module in place.
 
@@ -43,8 +70,6 @@ def fully_shard(
         placements: Parameter, gradient, and optimizer placements.
         mixed_precision_policy: Optional precision policy. Defaults to FP32 main weights
             and parameter-dtype main gradients.
-        use_symm_mem: Allocate all-gather and reduce-scatter staging buffers from
-            PyTorch's NCCL symmetric-memory pool.
     """
     if isinstance(module, FsdpModule):
         raise ValueError("This module is already managed by FSDP.")
@@ -59,7 +84,7 @@ def fully_shard(
             mesh=mesh,
             placements=placements,
             mixed_precision_policy=mixed_precision_policy,
-            use_symm_mem=use_symm_mem,
+            use_symmetric_memory=_FULLY_SHARD_OPTIONS.get().use_symmetric_memory,
         )
     except Exception:
         module.__class__ = original_cls

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

@@ -97,7 +97,7 @@ def __init__(
         mesh: DeviceMesh,
         placements: Placements,
         mixed_precision_policy: MixedPrecisionPolicy,
-        use_symm_mem: bool = False,
+        use_symmetric_memory: bool = False,
     ) -> None:
         """Initialize FSDP runtime state on an already-constructed module."""
         self._context = None
@@ -115,7 +115,7 @@ def __init__(
                 mesh=mesh,
                 placements=placements,
                 mixed_precision_policy=mixed_precision_policy,
-                use_symm_mem=use_symm_mem,
+                use_symmetric_memory=use_symmetric_memory,
             )
             for group_parameters in _group_parameters(owned_parameters)
         ]

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

@@ -57,7 +57,7 @@ def __init__(
         mesh: DeviceMesh,
         placements: Placements,
         mixed_precision_policy: MixedPrecisionPolicy,
-        use_symm_mem: bool = False,
+        use_symmetric_memory: bool = False,
     ) -> None:
         """Create persistent sharded buffers for a group of parameters.
 
@@ -67,7 +67,7 @@ def __init__(
             mesh: Device mesh used for all DBuffer storage in this version.
             placements: Parameter, gradient, and optimizer placements.
             mixed_precision_policy: Precision policy for main weights and gradients.
-            use_symm_mem: Allocate communication staging buffers from PyTorch's
+            use_symmetric_memory: Allocate communication staging buffers from PyTorch's
                 NCCL symmetric-memory pool.
         """
         if not parameters:
@@ -105,7 +105,7 @@ def __init__(
             placements=main_weight_placements,
         )
 
-        if use_symm_mem:
+        if use_symmetric_memory:
             # PyTorch caches this in C++ and returns early when the backend is already NCCL.
             symm_mem.set_backend("NCCL")
             self._symm_mem_pool = symm_mem.get_mem_pool(self.main_weight.device)

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -19,6 +19,7 @@
     Replicate,
     fully_shard,
     fully_shard_optimizer,
+    fully_shard_options,
     microbatch,
 )
 from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy
@@ -472,8 +473,8 @@ def test_root_backward_returns_to_resting_memory(distributed_setup):
     )
 
 
-@pytest.mark.parametrize("use_symm_mem", [False, True], ids=["default", "symmetric_memory"])
-def test_overlaps_communication_and_compute(distributed_setup, use_symm_mem):
+@pytest.mark.parametrize("use_symmetric_memory", [False, True], ids=["default", "symmetric_memory"])
+def test_overlaps_communication_and_compute(distributed_setup, use_symmetric_memory):
     """Forward and backward communication should overlap GEMM compute."""
     world_size = distributed_setup.world_size
     device = distributed_setup.device
@@ -498,7 +499,7 @@ def test_overlaps_communication_and_compute(distributed_setup, use_symm_mem):
     if not dist.is_initialized():
         dist.init_process_group(backend="nccl")
 
-    if use_symm_mem:
+    if use_symmetric_memory:
         # Dedicated communicator with NCCL's zero-CTA policy. cta_policy is a
         # per-communicator property, so scoping it to this group leaves the rest of the
         # bucket on default-CTA symmetric-memory kernels (test_symmetric_memory.py asserts
@@ -520,21 +521,10 @@ def test_overlaps_communication_and_compute(distributed_setup, use_symm_mem):
     model = MultiChildModel(dim=dim, num_children=num_children).to(dtype=dtype)
     placements = _flat_placements()
     policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype)
-    for layer in model.layers:
-        fully_shard(
-            layer,
-            mesh=mesh,
-            placements=placements,
-            mixed_precision_policy=policy,
-            use_symm_mem=use_symm_mem,
-        )
-    fully_shard(
-        model,
-        mesh=mesh,
-        placements=placements,
-        mixed_precision_policy=policy,
-        use_symm_mem=use_symm_mem,
-    )
+    with fully_shard_options(use_symmetric_memory=use_symmetric_memory):
+        for layer in model.layers:
+            fully_shard(layer, mesh=mesh, placements=placements, mixed_precision_policy=policy)
+        fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy)
 
     x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True)
 
@@ -568,7 +558,7 @@ def train_one_iteration() -> None:
     # forward and a backward all-gather and one reduce-scatter. Zero-CTA moves the
     # all-gather to copy-engine memcpys, so it should not emit all-gather kernels.
     num_sharded_modules = num_children + 1
-    expected_allgather_kernel_count = 0 if use_symm_mem else 2 * num_sharded_modules
+    expected_allgather_kernel_count = 0 if use_symmetric_memory else 2 * num_sharded_modules
     assert len(allgather_kernels) == expected_allgather_kernel_count, (
         f"Expected {expected_allgather_kernel_count} all-gather kernels, got "
         f"{len(allgather_kernels)}: {[kernel.name for kernel in allgather_kernels]}"
@@ -597,7 +587,7 @@ def train_one_iteration() -> None:
     )
     expected_allgather_overlap = 2 * (num_children - 1)
     expected_reduce_scatter_overlap = num_children - 1
-    if not use_symm_mem:
+    if not use_symmetric_memory:
         assert allgather_overlap_count >= expected_allgather_overlap, (
             f"Expected at least {expected_allgather_overlap} all-gathers to "
             f"overlap compute, got {allgather_overlap_count}/{len(allgather_kernels)}."

tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py

@@ -14,6 +14,7 @@
     Flat,
     Placements,
     fully_shard,
+    fully_shard_options,
 )
 from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import collect_linked_kernels
 
@@ -58,24 +59,23 @@ def test_fully_shard_symmetric_memory_matches_default_and_profiles_nccl(
     mesh = init_device_mesh(device.type, (world_size,))
     num_training_steps = 5
 
-    def train(use_symm_mem: bool) -> list[torch.Tensor]:
+    def train(use_symmetric_memory: bool) -> list[torch.Tensor]:
         torch.manual_seed(1234)
         model = TinyModel().to(device=device, dtype=torch.bfloat16)
         mixed_precision_policy = MixedPrecisionPolicy(main_params_dtype=torch.float32)
-        fully_shard(
-            model.fc1,
-            mesh=mesh,
-            placements=_flat_placements(),
-            mixed_precision_policy=mixed_precision_policy,
-            use_symm_mem=use_symm_mem,
-        )
-        fully_shard(
-            model.fc2,
-            mesh=mesh,
-            placements=_flat_placements(),
-            mixed_precision_policy=mixed_precision_policy,
-            use_symm_mem=use_symm_mem,
-        )
+        with fully_shard_options(use_symmetric_memory=use_symmetric_memory):
+            fully_shard(
+                model.fc1,
+                mesh=mesh,
+                placements=_flat_placements(),
+                mixed_precision_policy=mixed_precision_policy,
+            )
+            fully_shard(
+                model.fc2,
+                mesh=mesh,
+                placements=_flat_placements(),
+                mixed_precision_policy=mixed_precision_policy,
+            )
         optimizer = torch.optim.SGD(model.parameters(), lr=0.05, foreach=False)
 
         micro_batch_size = 2
@@ -99,11 +99,11 @@ def train(use_symm_mem: bool) -> list[torch.Tensor]:
         return losses
 
     with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof_without_symm_mem:
-        losses_without_symm_mem = train(use_symm_mem=False)
+        losses_without_symm_mem = train(use_symmetric_memory=False)
         torch.cuda.synchronize()
 
     with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof_with_symm_mem:
-        losses_with_symm_mem = train(use_symm_mem=True)
+        losses_with_symm_mem = train(use_symmetric_memory=True)
         torch.cuda.synchronize()
 
     torch.testing.assert_close(
@@ -186,20 +186,19 @@ def test_fully_shard_zero_cta_moves_all_gather_to_copy_engine(distributed_setup)
     num_training_steps = 5
     model = TinyModel().to(device=device, dtype=torch.bfloat16)
     mixed_precision_policy = MixedPrecisionPolicy(main_params_dtype=torch.float32)
-    fully_shard(
-        model.fc1,
-        mesh=mesh,
-        placements=_flat_placements(),
-        mixed_precision_policy=mixed_precision_policy,
-        use_symm_mem=True,
-    )
-    fully_shard(
-        model.fc2,
-        mesh=mesh,
-        placements=_flat_placements(),
-        mixed_precision_policy=mixed_precision_policy,
-        use_symm_mem=True,
-    )
+    with fully_shard_options(use_symmetric_memory=True):
+        fully_shard(
+            model.fc1,
+            mesh=mesh,
+            placements=_flat_placements(),
+            mixed_precision_policy=mixed_precision_policy,
+        )
+        fully_shard(
+            model.fc2,
+            mesh=mesh,
+            placements=_flat_placements(),
+            mixed_precision_policy=mixed_precision_policy,
+        )
     optimizer = torch.optim.SGD(model.parameters(), lr=0.05, foreach=False)
     x = torch.randn(2, _HIDDEN, device=device, dtype=torch.bfloat16)
     target = torch.randn(2, _HIDDEN, device=device, dtype=torch.bfloat16)