From 22964567981587d0c68a24d8a062af3e9687b39d Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 29 Jul 2026 14:26:38 +0000 Subject: [PATCH] Add construction-scoped fully shard options Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/__init__.py | 3 +- .../megatron_fsdp/experimental/fully_shard.py | 33 ++++++++-- .../src/megatron_fsdp/experimental/module.py | 4 +- .../experimental/parameter_group.py | 6 +- .../distributed/mfsdp_v2/test_fully_shard.py | 30 +++------ .../mfsdp_v2/test_symmetric_memory.py | 61 +++++++++---------- 6 files changed, 76 insertions(+), 61 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index bae27be831c..897375be7a3 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/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", ] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py index 0f53c34f359..da3b50aa660 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py +++ b/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 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..2794b9e2e75 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/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) ] 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..80bb9bf651f 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 @@ -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) 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..d7b029db084 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/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)}." diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py b/tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py index 87207548140..6c0cfd55c57 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py +++ b/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)