From 7adea24e6f42f87f616238468ca75668638418a1 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 1 Aug 2026 06:49:59 +0000 Subject: [PATCH 1/8] Add eager MFSDP construction context Signed-off-by: Jingyue Wu --- .../distributed/fsdp/mcore_fsdp_adapter.py | 30 ++-- .../megatron_fsdp/experimental/__init__.py | 5 +- .../megatron_fsdp/experimental/fully_shard.py | 60 +++++-- .../src/megatron_fsdp/experimental/module.py | 166 +++++++++++------- .../distributed/mfsdp_v2/test_annotation.py | 23 ++- .../distributed/mfsdp_v2/test_context.py | 105 +++++++++-- .../distributed/mfsdp_v2/test_cuda_graph.py | 8 +- .../distributed/mfsdp_v2/test_fully_shard.py | 147 +++++++++------- .../distributed/mfsdp_v2/test_optimizer.py | 31 ++-- .../mfsdp_v2/test_symmetric_memory.py | 59 ++++--- 10 files changed, 401 insertions(+), 233 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index d50cb220a24..a8393fcf147 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -57,6 +57,7 @@ Flat, Placements, fully_shard, + fully_shard_context, ) HAVE_MEGATRON_FSDP = True @@ -563,19 +564,22 @@ def __init__( placements = Placements( dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()] ) - for submodule in reversed(list(module.modules())): - if submodule is module: - # The root is always sharded after selected child units so it is not - # wrapped twice when its type also appears in fsdp_unit_modules. - continue - if any(isinstance(submodule, module_type) for module_type in fsdp_unit_modules): - fully_shard( - submodule, - mesh=mesh, - placements=placements, - mixed_precision_policy=self.mp_policy, - ) - fully_shard(module, mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy) + with fully_shard_context(device=device): + for submodule in reversed(list(module.modules())): + if submodule is module: + # The root is always sharded after selected child units so it is not + # wrapped twice when its type also appears in fsdp_unit_modules. + continue + if any(isinstance(submodule, module_type) for module_type in fsdp_unit_modules): + fully_shard( + submodule, + mesh=mesh, + placements=placements, + mixed_precision_policy=self.mp_policy, + ) + fully_shard( + module, mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy + ) super().__init__(config=config, module=module) @staticmethod 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..01f798dc639 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -15,18 +15,21 @@ """Experimental Megatron-FSDP implementation.""" from .dbuffer import DBuffer -from .fully_shard import fully_shard, microbatch +from .fully_shard import fully_shard, fully_shard_context, microbatch +from .module import FsdpContext from .optimizer import fully_shard_optimizer from .placement import Flat, Partial, Placement, Placements, Replicate __all__ = [ "DBuffer", "Flat", + "FsdpContext", "Partial", "Placement", "Placements", "Replicate", "fully_shard", + "fully_shard_context", "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 076a5ee227e..de680d066c7 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 @@ -17,7 +17,9 @@ import dataclasses from collections.abc import Iterator from contextlib import contextmanager +from contextvars import ContextVar +import torch from torch import nn from torch.distributed import DeviceMesh @@ -25,6 +27,35 @@ from .module import FsdpContext, FsdpModule from .placement import MeshAxis, Placements +_FSDP_CONTEXT = ContextVar[FsdpContext | None]("megatron_fsdp_context", default=None) + + +@contextmanager +def fully_shard_context(device: torch.device | None = None) -> Iterator[FsdpContext]: + """Construct FSDP modules that share runtime streams and prefetch orders. + + Independent roots are ordered by their root-level ``fully_shard`` calls. + Construction must finish before any of the registered modules run forward. + + Args: + device: CUDA device on which to create communication streams. Defaults to + the current CUDA device. + """ + device = device or torch.device("cuda", torch.cuda.current_device()) + if device.type != "cuda": + raise ValueError(f"fully_shard_context requires a CUDA device, got {device}.") + + context = FsdpContext(device=device) + token = _FSDP_CONTEXT.set(context) + try: + yield context + except Exception: + raise + else: + context.finalize() + finally: + _FSDP_CONTEXT.reset(token) + def fully_shard( module: nn.Module, @@ -49,6 +80,9 @@ def fully_shard( """ if isinstance(module, FsdpModule): raise ValueError("This module is already managed by FSDP.") + context = _FSDP_CONTEXT.get() + if context is None: + raise RuntimeError("fully_shard must run inside fully_shard_context.") placements = _normalize_placements(mesh, placements) mixed_precision_policy = mixed_precision_policy or MixedPrecisionPolicy() @@ -58,6 +92,7 @@ def fully_shard( assert isinstance(module, FsdpModule) FsdpModule.__init__( module, + context=context, mesh=mesh, placements=placements, mixed_precision_policy=mixed_precision_policy, @@ -90,7 +125,7 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: @contextmanager -def microbatch(module: nn.Module, is_last: bool) -> Iterator[None]: +def microbatch(context: FsdpContext, is_last: bool) -> Iterator[None]: """Mark an FSDP microbatch as the last accumulation microbatch. At present, this is only needed for HSDP/HFSDP gradient accumulation, so @@ -98,20 +133,17 @@ def microbatch(module: nn.Module, is_last: bool) -> Iterator[None]: parallelism finalizes gradients on every backward and does not need it. Args: - module: Module tree whose FSDP roots should use this microbatch state. + context: FSDP context whose roots should use this microbatch state. is_last: Whether forwards in this scope are for the last microbatch. """ - contexts: list[FsdpContext] = [] - _collect_fsdp_contexts(module, contexts) - previous_states = [(context, context.is_last_microbatch) for context in contexts] - for context in contexts: - context.is_last_microbatch = is_last + context.ensure_finalized() + previous_state = context.is_last_microbatch + context.is_last_microbatch = is_last try: yield finally: - for context, is_last_microbatch in previous_states: - context.is_last_microbatch = is_last_microbatch + context.is_last_microbatch = previous_state def _attach_mixin(module: nn.Module) -> None: @@ -120,13 +152,3 @@ def _attach_mixin(module: nn.Module) -> None: module_cls = module.__class__ fsdp_cls = type(f"ExperimentalFsdp{module_cls.__name__}", (FsdpModule, module_cls), {}) module.__class__ = fsdp_cls - - -def _collect_fsdp_contexts(module: nn.Module, contexts: list[FsdpContext]) -> None: - if isinstance(module, FsdpModule): - module._lazy_init_context() - contexts.append(module.context) - return - - for child in module.children(): - _collect_fsdp_contexts(child, contexts) 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 476c1cb6bc7..fd01994124f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -28,7 +28,7 @@ class FsdpContext: - """Runtime stream and prefetch state shared by one FSDP subtree.""" + """Runtime stream and prefetch state shared by FSDP roots constructed together.""" allgather_stream: torch.cuda.Stream reduce_scatter_stream: torch.cuda.Stream @@ -36,28 +36,72 @@ class FsdpContext: # unnecessary because it can be detected when ``model_weight``, after syncing # from ``main_weight``, has placements different from ``Placements.optimizer``. is_last_microbatch: bool - root_module: "FsdpModule" + root_modules: tuple["FsdpModule", ...] # Static orders used to drive all-gather prefetch. We may want to switch to # capturing runtime order if static module order proves too fragile. Each # FsdpModule tracks its own materialized state via ``FsdpModule._unshard_event``. forward_order: IndexedOrder["FsdpModule"] backward_order: IndexedOrder["FsdpModule"] - def __init__(self, device: torch.device, root_module: "FsdpModule") -> None: - """Create rank-local runtime state for a root FSDP subtree. + def __init__(self, device: torch.device) -> None: + """Create rank-local runtime state for FSDP modules on ``device``. Args: device: Device on which this context schedules communication. - root_module: Outermost module that owns this context. """ - self.root_module = root_module + self.root_modules = () self.is_last_microbatch = True self.forward_order = IndexedOrder() self.backward_order = IndexedOrder() + self._registered_modules: list[FsdpModule] = [] + self._is_finalized = False with torch.cuda.device(device): self.allgather_stream = torch.cuda.Stream() self.reduce_scatter_stream = torch.cuda.Stream() + def register_module(self, module: "FsdpModule") -> None: + """Register a module constructed in this context.""" + if self._is_finalized: + raise RuntimeError("Cannot register an FSDP module after its context is finalized.") + self._registered_modules.append(module) + + def finalize(self) -> None: + """Finalize roots, names, and cross-root prefetch orders.""" + if self._is_finalized: + raise RuntimeError("FSDP context is already finalized.") + + visited: set[FsdpModule] = set() + root_modules: list[FsdpModule] = [] + for module in reversed(self._registered_modules): + if module in visited: + continue + root_modules.append(module) + _collect_context_modules(cast(nn.Module, module), self, visited) + root_modules.reverse() + self.root_modules = tuple(root_modules) + + seen: set[FsdpModule] = set() + for root_module in self.root_modules: + _collect_forward_order( + cast(nn.Module, root_module), self, self.forward_order, seen, name="" + ) + + forward_modules = seen.copy() + seen.clear() + for root_module in reversed(self.root_modules): + _collect_backward_order(cast(nn.Module, root_module), self, self.backward_order, seen) + + if seen != forward_modules: + raise RuntimeError("FSDP context produced inconsistent prefetch orders.") + self._is_finalized = True + + def ensure_finalized(self) -> None: + """Raise if construction has not completed for this context.""" + if not self._is_finalized: + raise RuntimeError( + "FSDP context is not finalized. Exit fully_shard_context before running forward." + ) + def current_stream(self) -> torch.cuda.Stream: """Current stream on this context's device.""" return torch.cuda.current_stream(self.allgather_stream.device) @@ -84,7 +128,7 @@ class FsdpModule: # Root uses "" and None means uninitialized. _name: str | None _parameter_groups: tuple[FsdpParameterGroup, ...] - _context: FsdpContext | None + _context: FsdpContext _num_ready_grad_parameters: int _num_trainable_parameters: int # Event recorded after this FsdpModule's full parameters are materialized. @@ -94,13 +138,14 @@ class FsdpModule: def __init__( self, + context: FsdpContext, mesh: DeviceMesh, placements: Placements, mixed_precision_policy: MixedPrecisionPolicy, use_symm_mem: bool = False, ) -> None: """Initialize FSDP runtime state on an already-constructed module.""" - self._context = None + self._context = context self._name = None self._unshard_event = None owned_parameters = _collect_owned_parameters(self) @@ -124,55 +169,7 @@ def __init__( len(group.fsdp_parameters) for group in self._parameter_groups if group.requires_grad ) self._register_hooks() - - def _lazy_init_context(self) -> None: - """Initialize one shared runtime context for this FSDP root subtree. - - MFSDP v2 requires users to apply ``fully_shard`` bottom-up, so child FSDP - modules are constructed before their eventual root module is constructed. - This method resolves the root lazily on the first forward through the - outermost FSDP module and shares that one context with every FSDP - descendant. - - Alternatives considered: - - Eagerly initialize contexts during ``fully_shard``. When a parent is - sharded, we could create a new root context and reassign it to all - descendant FSDP modules. This creates transient child contexts that are - never used if the parent is later sharded, and each parent shard must - walk its descendants again, making nested sharding quadratic. - - Store an ``is_root`` field on each FSDP module. ``fully_shard`` could - mark newly sharded modules as roots and clear that flag on descendant - FSDP modules when a parent is sharded. This avoids creating unused - contexts but moves root tracking onto every FSDP module, adding - per-module state that must stay consistent with the final sharded - module hierarchy. - """ - if self._context is not None: - return - - root_module = cast(nn.Module, self) - first_parameter = next(root_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) - # named_modules() yields FsdpModules in registration order, which is the static - # forward execution order used to prefetch the next FsdpModule's all-gather. - for submodule_name, submodule in root_module.named_modules(): - if not isinstance(submodule, FsdpModule): - continue - if submodule._context is not None: - raise RuntimeError( - "FSDP context is already initialized for a descendant module. " - "Run forward through the root FSDP module first." - ) - submodule._context = context - submodule._name = submodule_name - context.forward_order.append(submodule) - - # Backward starts from the root pre-backward hook before visiting child - # subtrees in reverse module order. - _collect_backward_order(root_module, context.backward_order) + context.register_module(self) @property def context(self) -> FsdpContext: @@ -189,8 +186,8 @@ def name(self) -> str: return name def is_root(self) -> bool: - """Return whether this module is the outermost FsdpModule in its context.""" - return self.context.root_module is self + """Return whether this module is an outermost FsdpModule in its context.""" + return self in self.context.root_modules def _register_hooks(self) -> None: module = cast(nn.Module, self) @@ -227,10 +224,10 @@ def pre_forward(self) -> None: While this FsdpModule computes, we issue the next FsdpModule's all-gather on the comm stream, so ``AG_{i+1}`` is launched before ``F_i`` finishes. """ - self._lazy_init_context() + context = self.context + context.ensure_finalized() torch.cuda.nvtx.range_push(self._nvtx_label("forward")) self._num_ready_grad_parameters = 0 - context = self.context allgather_stream = context.allgather_stream current_stream = context.current_stream() @@ -347,13 +344,54 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str: return f"MFSDP {name} {phase}" -def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None: - """Collect FsdpModules in static backward prefetch order.""" +def _collect_context_modules( + module: nn.Module, context: FsdpContext, modules: set["FsdpModule"] +) -> None: + """Collect same-context FSDP modules reachable from ``module``.""" + if isinstance(module, FsdpModule): + if module._context is not context or module in modules: + return + modules.add(module) + + for child in module.children(): + _collect_context_modules(child, context, modules) + + +def _collect_forward_order( + module: nn.Module, + context: FsdpContext, + order: IndexedOrder["FsdpModule"], + seen: set["FsdpModule"], + name: str, +) -> None: + """Collect one root's static forward prefetch order.""" if isinstance(module, FsdpModule): + if module._context is not context or module in seen: + return + seen.add(module) + module._name = name + order.append(module) + + for child_name, child in module.named_children(): + child_fqn = f"{name}.{child_name}" if name else child_name + _collect_forward_order(child, context, order, seen, child_fqn) + + +def _collect_backward_order( + module: nn.Module, + context: FsdpContext, + order: IndexedOrder["FsdpModule"], + seen: set["FsdpModule"], +) -> None: + """Collect one root's static backward prefetch order.""" + if isinstance(module, FsdpModule): + if module._context is not context or module in seen: + return + seen.add(module) order.append(module) for child in reversed(list(module.children())): - _collect_backward_order(child, order) + _collect_backward_order(child, context, order, seen) def _collect_owned_parameters(root_module: nn.Module) -> dict[str, nn.Parameter]: diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_annotation.py b/tests/unit_tests/distributed/mfsdp_v2/test_annotation.py index f65e0b85724..a635accb5d9 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_annotation.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_annotation.py @@ -14,6 +14,7 @@ Flat, Placements, fully_shard, + fully_shard_context, ) _NVTX_LABEL_PATTERN = re.compile(r"MFSDP (.+) (forward|backward)") @@ -95,8 +96,9 @@ def test_fsdp_sibling_roots_emit_root_nvtx_ranges_after_training_step( _setup_nvtx_recording(monkeypatch, events) model = NestedLinearModel(dim=4).to(distributed_setup.device) mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) - fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) - fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=distributed_setup.device): + fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) + fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) model(torch.ones(2, 4, device=distributed_setup.device)).sum().backward() @@ -118,9 +120,10 @@ def test_fsdp_training_hooks_emit_stacked_nvtx_ranges(distributed_setup, monkeyp _setup_nvtx_recording(monkeypatch, events) model = NestedLinearModel(dim=4).to(distributed_setup.device) mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) - fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) - fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=distributed_setup.device): + fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) + fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) model(torch.ones(2, 4, device=distributed_setup.device)).sum().backward() @@ -148,7 +151,8 @@ def test_fsdp_frozen_parameters_emit_balanced_backward_nvtx_range(distributed_se for parameter in model.parameters(): parameter.requires_grad_(False) mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=distributed_setup.device): + fully_shard(model, mesh=mesh, placements=_flat_placements()) x = torch.ones(2, 4, device=distributed_setup.device, requires_grad=True) model(x).sum().backward() @@ -171,9 +175,10 @@ def test_fsdp_frozen_child_without_grad_inputs_skips_backward_nvtx_range( for parameter in model.layers[0].parameters(): parameter.requires_grad_(False) mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) - fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) - fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=distributed_setup.device): + fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) + fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) model(torch.ones(2, 4, device=distributed_setup.device)).sum().backward() diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_context.py b/tests/unit_tests/distributed/mfsdp_v2/test_context.py index 102b2fde332..b8f349a2bad 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_context.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_context.py @@ -2,6 +2,7 @@ """Unit tests for experimental Megatron-FSDP runtime contexts.""" +import pytest import torch from torch import nn from torch.distributed.device_mesh import init_device_mesh @@ -10,6 +11,7 @@ Flat, Placements, fully_shard, + fully_shard_context, ) @@ -74,14 +76,17 @@ def _flat_placements() -> Placements: def test_child_then_parent_share_one_context(distributed_setup): - """A parent FsdpModule should lazily create one context for its subtree.""" + """Modules constructed together should eagerly share one context.""" device = distributed_setup.device mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) model = NestedModel().to(device) - fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device) as context: + fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) + assert model.context is context + assert model.inner.context is context with torch.no_grad(): model(torch.ones(2, 4, device=device)) @@ -91,16 +96,17 @@ def test_child_then_parent_share_one_context(distributed_setup): assert not model.inner.is_root() -def test_two_child_subtrees_then_parent_collapse_to_one_context(distributed_setup): - """Sharding a parent should lazily assign one context across child subtrees.""" +def test_two_child_subtrees_then_parent_share_one_context(distributed_setup): + """One construction scope should assign one context across child subtrees.""" device = distributed_setup.device mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) model = MultiChildModel(dim=4, num_children=2).to(device) - fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) - fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) + fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) with torch.no_grad(): model(torch.ones(2, 4, device=device)) @@ -109,22 +115,26 @@ def test_two_child_subtrees_then_parent_collapse_to_one_context(distributed_setu assert model.layers[1].context is model.context -def test_sibling_roots_without_parent_keep_separate_contexts(distributed_setup): - """Independent FSDP roots should not share runtime scheduling state.""" +def test_sibling_roots_share_context_and_cross_root_orders(distributed_setup): + """Independent roots should share streams and follow construction order.""" device = distributed_setup.device mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) model = MultiChildModel(dim=4, num_children=2).to(device) - fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) - fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.layers[0], mesh=mesh, placements=_flat_placements()) + fully_shard(model.layers[1], mesh=mesh, placements=_flat_placements()) with torch.no_grad(): model(torch.ones(2, 4, device=device)) - assert model.layers[0].context is not model.layers[1].context + context = model.layers[0].context + assert model.layers[1].context is context assert model.layers[0].is_root() assert model.layers[1].is_root() + assert list(context.forward_order) == [model.layers[0], model.layers[1]] + assert list(context.backward_order) == [model.layers[1], model.layers[0]] def test_nested_prefetch_orders_use_dfs(distributed_setup): @@ -134,10 +144,11 @@ def test_nested_prefetch_orders_use_dfs(distributed_setup): mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) model = NestedSiblingModel(dim=4).to(device) - fully_shard(model.left.inner, mesh=mesh, placements=_flat_placements()) - fully_shard(model.left, mesh=mesh, placements=_flat_placements()) - fully_shard(model.right, mesh=mesh, placements=_flat_placements()) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.left.inner, mesh=mesh, placements=_flat_placements()) + fully_shard(model.left, mesh=mesh, placements=_flat_placements()) + fully_shard(model.right, mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) with torch.no_grad(): model(torch.ones(2, 4, device=device)) @@ -145,3 +156,63 @@ def test_nested_prefetch_orders_use_dfs(distributed_setup): context = model.context assert list(context.forward_order) == [model, model.left, model.left.inner, model.right] assert list(context.backward_order) == [model, model.right, model.left, model.left.inner] + + +def test_nested_and_sibling_roots_use_cross_root_orders(distributed_setup): + """Context orders should concatenate nested roots at construction boundaries.""" + device = distributed_setup.device + mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) + model = NestedSiblingModel(dim=4).to(device) + + with fully_shard_context(device=device): + fully_shard(model.left.inner, mesh=mesh, placements=_flat_placements()) + fully_shard(model.left, mesh=mesh, placements=_flat_placements()) + fully_shard(model.right, mesh=mesh, placements=_flat_placements()) + + context = model.left.context + assert context.root_modules == (model.left, model.right) + assert list(context.forward_order) == [model.left, model.left.inner, model.right] + assert list(context.backward_order) == [model.right, model.left, model.left.inner] + + +def test_fully_shard_requires_context(distributed_setup): + """fully_shard should reject construction without an active context.""" + device = distributed_setup.device + mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) + model = nn.Linear(4, 4, bias=False).to(device) + + with pytest.raises(RuntimeError, match="inside fully_shard_context"): + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + +def test_forward_requires_finalized_context(distributed_setup): + """Forward should be unavailable until construction scope exit.""" + device = distributed_setup.device + mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) + model = nn.Linear(4, 4, bias=False).to(device) + x = torch.ones(2, 4, device=device) + + with fully_shard_context(device=device): + fully_shard(model, mesh=mesh, placements=_flat_placements()) + with pytest.raises(RuntimeError, match="Exit fully_shard_context"): + model(x) + + model(x) + + +def test_nested_context_restores_outer_context(distributed_setup): + """A nested construction scope should restore the enclosing active context.""" + device = distributed_setup.device + mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) + model = nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(3)]).to(device) + + with fully_shard_context(device=device): + fully_shard(model[0], mesh=mesh, placements=_flat_placements()) + outer_context = model[0].context + with fully_shard_context(device=device): + fully_shard(model[1], mesh=mesh, placements=_flat_placements()) + fully_shard(model[2], mesh=mesh, placements=_flat_placements()) + + assert model[0].context is outer_context + assert model[2].context is outer_context + assert model[1].context is not outer_context diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py b/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py index 08920bcea98..86d764f72f7 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_cuda_graph.py @@ -12,6 +12,7 @@ Flat, Placements, fully_shard, + fully_shard_context, ) logger = logging.getLogger(__name__) @@ -51,9 +52,10 @@ def test_captures_full_iteration(distributed_setup): static_target = torch.zeros_like(static_input) placements = _flat_placements() - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=placements) - fully_shard(model, mesh=mesh, placements=placements) + with fully_shard_context(device=device): + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=placements) + fully_shard(model, mesh=mesh, placements=placements) optimizer = torch.optim.SGD(model.parameters(), lr=0.25, foreach=False) 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 e820abe515d..448facc751c 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -18,6 +18,7 @@ Placements, Replicate, fully_shard, + fully_shard_context, fully_shard_optimizer, microbatch, ) @@ -161,8 +162,9 @@ def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatch model = TinyModel().to(device) model.load_state_dict(baseline.state_dict()) - fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) - fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) @@ -234,9 +236,10 @@ def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_ # Shard the child layers, then the model, so the children share a root context # and reduce through the overlap path instead of as independent roots. - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) - fully_shard(model, mesh=mesh, placements=_hsdp_placements()) + with fully_shard_context(device=device) as context: + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) + fully_shard(model, mesh=mesh, placements=_hsdp_placements()) baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) @@ -252,7 +255,7 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]: for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): is_last = microbatch_index == num_microbatches - 1 - with microbatch(model, is_last=is_last): + with microbatch(context, is_last=is_last): loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) (loss / num_microbatches).backward() losses.append(loss.detach()) @@ -303,9 +306,10 @@ def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup): dim = 8 num_children = 2 model = MultiChildModel(dim=dim, num_children=num_children).to(device) - for layer in model.layers: - fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) - fully_shard(model, mesh=mesh, placements=_hsdp_placements()) + with fully_shard_context(device=device) as context: + for layer in model.layers: + fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) + fully_shard(model, mesh=mesh, placements=_hsdp_placements()) optimizer = torch.optim.SGD(model.parameters(), lr=0.05) num_microbatches = 3 @@ -318,7 +322,7 @@ def train_one_step() -> None: optimizer.zero_grad(set_to_none=True) for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches): is_last = microbatch_index == num_microbatches - 1 - with microbatch(model, is_last=is_last): + with microbatch(context, is_last=is_last): loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target) (loss / num_microbatches).backward() optimizer.step() @@ -353,8 +357,9 @@ def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup): mesh = init_device_mesh(device.type, (world_size,)) model = NestedModel().to(device) - fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) (inner_group,) = model.inner.parameter_groups (outer_group,) = model.parameter_groups @@ -391,9 +396,10 @@ def test_forward_peak_memory_bounds_in_flight_child_all_gathers(distributed_setu model = MultiChildModel(dim=dim, num_children=4).to(dtype=dtype, device=device) 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) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + with fully_shard_context(device=device): + 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(2, dim, device=device, dtype=dtype) with torch.no_grad(): @@ -437,9 +443,10 @@ def test_root_forward_returns_to_resting_memory(distributed_setup): model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) 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) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + with fully_shard_context(device=device): + 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(2, dim, device=device, dtype=dtype) torch.cuda.synchronize(device) @@ -475,9 +482,10 @@ def test_root_backward_returns_to_resting_memory(distributed_setup): model = MultiChildModel(dim=dim, num_children=2).to(dtype=dtype, device=device) 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) - fully_shard(model, mesh=mesh, placements=placements, mixed_precision_policy=policy) + with fully_shard_context(device=device): + 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(2, dim, device=device, dtype=dtype, requires_grad=True) output = model(x) @@ -548,21 +556,22 @@ 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: + with fully_shard_context(device=device): + for layer in model.layers: + fully_shard( + layer, + mesh=mesh, + placements=placements, + mixed_precision_policy=policy, + use_symm_mem=use_symm_mem, + ) fully_shard( - layer, + model, 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, - ) x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True) @@ -648,9 +657,10 @@ def test_parameterless_parent_with_child_modules_trains(distributed_setup): 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()) + with fully_shard_context(device=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 == () @@ -674,7 +684,8 @@ def test_frozen_parameter_group_does_not_allocate_main_grad(distributed_setup): model = nn.Linear(4, 4, bias=False).to(device) model.weight.requires_grad_(False) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model, mesh=mesh, placements=_flat_placements()) (group,) = model.parameter_groups assert not group.requires_grad @@ -694,7 +705,8 @@ def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_se with torch.no_grad(): model.weight.fill_(1.0) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model, mesh=mesh, placements=_flat_placements()) x = torch.full((1, 1), float(rank + 1), device=device) model(x).sum().backward() @@ -718,12 +730,13 @@ def test_next_forward_uses_optimizer_updated_weights(distributed_setup): with torch.no_grad(): model.weight.fill_(1.0) - fully_shard( - model, - mesh=mesh, - placements=_flat_placements(), - mixed_precision_policy=MixedPrecisionPolicy(main_params_dtype=torch.float32), - ) + with fully_shard_context(device=device): + fully_shard( + model, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=MixedPrecisionPolicy(main_params_dtype=torch.float32), + ) # SGD's foreach/fused CUDA paths require matching parameter and gradient dtypes. # Use the scalar path to exercise FP32 main weights with default BF16 main grads. optimizer = torch.optim.SGD(model.parameters(), lr=0.25, foreach=False) @@ -753,8 +766,9 @@ def test_optimizer_post_step_syncs_once_per_parameter_group(distributed_setup, m mesh = init_device_mesh(device.type, (world_size,)) model = TinyModel().to(device=device, dtype=torch.bfloat16) - fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) - fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) parameter_groups = (*model.fc1.parameter_groups, *model.fc2.parameter_groups) sync_counts = {parameter_group: 0 for parameter_group in parameter_groups} @@ -797,8 +811,9 @@ def test_fully_shard_adam_mixed_precision_losses_match_baseline(distributed_setu baseline = TinyModel().to(device=device, dtype=torch.bfloat16) model = TinyModel().to(device=device, dtype=torch.bfloat16) model.load_state_dict(baseline.state_dict()) - fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) - fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) baseline_optimizer = torch.optim.Adam(baseline.parameters(), lr=0.01) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) @@ -821,22 +836,21 @@ def test_fully_shard_adam_mixed_precision_losses_match_baseline(distributed_setu optimizer.step() -def test_microbatch_scopes_child_contexts(distributed_setup): - """microbatch() should scope FSDP child contexts under an unwrapped parent.""" +def test_microbatch_scopes_context(distributed_setup): + """microbatch() should scope state on the supplied FSDP context.""" world_size = distributed_setup.world_size device = distributed_setup.device mesh = init_device_mesh(device.type, (world_size,)) model = nn.Sequential(nn.Linear(1, 1, bias=False), nn.Linear(1, 1, bias=False)).to(device) - for layer in model: - fully_shard(layer, mesh=mesh, placements=_flat_placements()) - - with microbatch(model, is_last=False): + with fully_shard_context(device=device) as context: for layer in model: - assert not layer.context.is_last_microbatch + fully_shard(layer, mesh=mesh, placements=_flat_placements()) + + with microbatch(context, is_last=False): + assert not context.is_last_microbatch - for layer in model: - assert layer.context.is_last_microbatch + assert context.is_last_microbatch def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): @@ -854,7 +868,8 @@ def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup): # 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()) + with fully_shard_context(device=device): + fully_shard(model[1], mesh=mesh, placements=_flat_placements()) assert model[0].weight.device.type == "cpu" assert isinstance(model[1].weight, DTensor) @@ -901,7 +916,8 @@ def test_non_leaf_parameter_view_survives_storage_resize(distributed_setup): mesh = init_device_mesh(device.type, (world_size,)) model = NonLeafViewModel().to(device) - fully_shard(model, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model, mesh=mesh, placements=_flat_placements()) group = model.parameter_groups[0] x = torch.randn(8, device=device, requires_grad=True) @@ -953,15 +969,16 @@ def train_steps(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Ten torch.manual_seed(4321) model = nn.Sequential(*[nn.Linear(dim, dim, dtype=dtype) for _ in range(layers)]).to(device) - for layer in model: - fully_shard( - layer, - mesh=mesh, - placements=_flat_placements(), - mixed_precision_policy=MixedPrecisionPolicy( - main_params_dtype=dtype, main_grads_dtype=dtype - ), - ) + with fully_shard_context(device=device): + for layer in model: + fully_shard( + layer, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=MixedPrecisionPolicy( + main_params_dtype=dtype, main_grads_dtype=dtype + ), + ) optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) torch.cuda.empty_cache() diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py index 4ba4d1be271..c3427073993 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -12,6 +12,7 @@ Flat, Placements, fully_shard, + fully_shard_context, fully_shard_optimizer, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy @@ -42,8 +43,9 @@ def test_adam_without_adapter_raises_precision_error(distributed_setup): mesh = init_device_mesh(device.type, (world_size,)) torch.manual_seed(2026) model = TinyModel().to(device=device, dtype=torch.bfloat16) - fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) - fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + with fully_shard_context(device=device): + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) @@ -68,18 +70,19 @@ def test_fused_adam_adapter_accepts_mismatched_grads(distributed_setup): mixed_precision_policy = MixedPrecisionPolicy( main_params_dtype=torch.float32, main_grads_dtype=torch.bfloat16 ) - 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, - ) + with fully_shard_context(device=device): + 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 = FusedAdam(model.parameters(), lr=0.01) fully_shard_optimizer(optimizer, precision_aware=True) 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..6ae3f05a95b 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_context, ) from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import collect_linked_kernels @@ -62,20 +63,21 @@ def train(use_symm_mem: 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_context(device=device): + 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, + ) optimizer = torch.optim.SGD(model.parameters(), lr=0.05, foreach=False) micro_batch_size = 2 @@ -186,20 +188,21 @@ 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_context(device=device): + 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, + ) 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) From 09895897e6722846d81d531802aa8e8388d318a2 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 1 Aug 2026 14:54:15 +0000 Subject: [PATCH 2/8] Refine MFSDP context ownership Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 3 +++ .../src/megatron_fsdp/experimental/module.py | 13 ++++++------ .../distributed/mfsdp_v2/test_context.py | 20 ++++++++++--------- 3 files changed, 21 insertions(+), 15 deletions(-) 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 de680d066c7..0445735eb9c 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 @@ -41,6 +41,9 @@ def fully_shard_context(device: torch.device | None = None) -> Iterator[FsdpCont device: CUDA device on which to create communication streams. Defaults to the current CUDA device. """ + if _FSDP_CONTEXT.get() is not None: + raise RuntimeError("fully_shard_context does not support nesting.") + device = device or torch.device("cuda", torch.cuda.current_device()) if device.type != "cuda": raise ValueError(f"fully_shard_context requires a CUDA device, got {device}.") 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 fd01994124f..22434bff979 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -36,7 +36,6 @@ class FsdpContext: # unnecessary because it can be detected when ``model_weight``, after syncing # from ``main_weight``, has placements different from ``Placements.optimizer``. is_last_microbatch: bool - root_modules: tuple["FsdpModule", ...] # Static orders used to drive all-gather prefetch. We may want to switch to # capturing runtime order if static module order proves too fragile. Each # FsdpModule tracks its own materialized state via ``FsdpModule._unshard_event``. @@ -49,7 +48,6 @@ def __init__(self, device: torch.device) -> None: Args: device: Device on which this context schedules communication. """ - self.root_modules = () self.is_last_microbatch = True self.forward_order = IndexedOrder() self.backward_order = IndexedOrder() @@ -78,17 +76,18 @@ def finalize(self) -> None: root_modules.append(module) _collect_context_modules(cast(nn.Module, module), self, visited) root_modules.reverse() - self.root_modules = tuple(root_modules) + for root_module in root_modules: + root_module._is_root = True seen: set[FsdpModule] = set() - for root_module in self.root_modules: + for root_module in root_modules: _collect_forward_order( cast(nn.Module, root_module), self, self.forward_order, seen, name="" ) forward_modules = seen.copy() seen.clear() - for root_module in reversed(self.root_modules): + for root_module in reversed(root_modules): _collect_backward_order(cast(nn.Module, root_module), self, self.backward_order, seen) if seen != forward_modules: @@ -130,6 +129,7 @@ class FsdpModule: _parameter_groups: tuple[FsdpParameterGroup, ...] _context: FsdpContext _num_ready_grad_parameters: int + _is_root: bool _num_trainable_parameters: int # Event recorded after this FsdpModule's full parameters are materialized. # ``None`` lets pre_forward enqueue an all-gather unless an earlier FsdpModule @@ -146,6 +146,7 @@ def __init__( ) -> None: """Initialize FSDP runtime state on an already-constructed module.""" self._context = context + self._is_root = False self._name = None self._unshard_event = None owned_parameters = _collect_owned_parameters(self) @@ -187,7 +188,7 @@ def name(self) -> str: def is_root(self) -> bool: """Return whether this module is an outermost FsdpModule in its context.""" - return self in self.context.root_modules + return self._is_root def _register_hooks(self) -> None: module = cast(nn.Module, self) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_context.py b/tests/unit_tests/distributed/mfsdp_v2/test_context.py index b8f349a2bad..8dc07879102 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_context.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_context.py @@ -170,7 +170,9 @@ def test_nested_and_sibling_roots_use_cross_root_orders(distributed_setup): fully_shard(model.right, mesh=mesh, placements=_flat_placements()) context = model.left.context - assert context.root_modules == (model.left, model.right) + assert model.left.is_root() + assert model.right.is_root() + assert not model.left.inner.is_root() assert list(context.forward_order) == [model.left, model.left.inner, model.right] assert list(context.backward_order) == [model.right, model.left, model.left.inner] @@ -200,19 +202,19 @@ def test_forward_requires_finalized_context(distributed_setup): model(x) -def test_nested_context_restores_outer_context(distributed_setup): - """A nested construction scope should restore the enclosing active context.""" +def test_fully_shard_context_rejects_nesting(distributed_setup): + """A construction scope should reject an ambiguous nested context.""" device = distributed_setup.device mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) - model = nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(3)]).to(device) + model = nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(2)]).to(device) with fully_shard_context(device=device): fully_shard(model[0], mesh=mesh, placements=_flat_placements()) outer_context = model[0].context - with fully_shard_context(device=device): - fully_shard(model[1], mesh=mesh, placements=_flat_placements()) - fully_shard(model[2], mesh=mesh, placements=_flat_placements()) + with pytest.raises(RuntimeError, match="does not support nesting"): + with fully_shard_context(device=device): + pass + fully_shard(model[1], mesh=mesh, placements=_flat_placements()) assert model[0].context is outer_context - assert model[2].context is outer_context - assert model[1].context is not outer_context + assert model[1].context is outer_context From 4f102fd117e72592961e8c4c2a2ab0c562b63fd0 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 1 Aug 2026 15:28:33 +0000 Subject: [PATCH 3/8] Simplify MFSDP context finalization Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/module.py | 62 +++++-------------- 1 file changed, 14 insertions(+), 48 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 22434bff979..87671f0c2bb 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -74,24 +74,21 @@ def finalize(self) -> None: if module in visited: continue root_modules.append(module) - _collect_context_modules(cast(nn.Module, module), self, visited) + _collect_fsdp_modules(cast(nn.Module, module), visited) root_modules.reverse() for root_module in root_modules: root_module._is_root = True - seen: set[FsdpModule] = set() for root_module in root_modules: - _collect_forward_order( - cast(nn.Module, root_module), self, self.forward_order, seen, name="" - ) + for name, module in cast(nn.Module, root_module).named_modules(): + if not isinstance(module, FsdpModule): + continue + module._name = name + self.forward_order.append(module) - forward_modules = seen.copy() - seen.clear() for root_module in reversed(root_modules): - _collect_backward_order(cast(nn.Module, root_module), self, self.backward_order, seen) + _collect_backward_order(cast(nn.Module, root_module), self.backward_order) - if seen != forward_modules: - raise RuntimeError("FSDP context produced inconsistent prefetch orders.") self._is_finalized = True def ensure_finalized(self) -> None: @@ -174,8 +171,7 @@ def __init__( @property def context(self) -> FsdpContext: - """Return the initialized runtime context.""" - assert self._context is not None + """Return the FSDP context.""" return self._context @property @@ -345,54 +341,24 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str: return f"MFSDP {name} {phase}" -def _collect_context_modules( - module: nn.Module, context: FsdpContext, modules: set["FsdpModule"] -) -> None: - """Collect same-context FSDP modules reachable from ``module``.""" +def _collect_fsdp_modules(module: nn.Module, modules: set["FsdpModule"]) -> None: + """Collect FSDP modules reachable from ``module``.""" if isinstance(module, FsdpModule): - if module._context is not context or module in modules: + if module in modules: return modules.add(module) for child in module.children(): - _collect_context_modules(child, context, modules) - - -def _collect_forward_order( - module: nn.Module, - context: FsdpContext, - order: IndexedOrder["FsdpModule"], - seen: set["FsdpModule"], - name: str, -) -> None: - """Collect one root's static forward prefetch order.""" - if isinstance(module, FsdpModule): - if module._context is not context or module in seen: - return - seen.add(module) - module._name = name - order.append(module) + _collect_fsdp_modules(child, modules) - for child_name, child in module.named_children(): - child_fqn = f"{name}.{child_name}" if name else child_name - _collect_forward_order(child, context, order, seen, child_fqn) - -def _collect_backward_order( - module: nn.Module, - context: FsdpContext, - order: IndexedOrder["FsdpModule"], - seen: set["FsdpModule"], -) -> None: +def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None: """Collect one root's static backward prefetch order.""" if isinstance(module, FsdpModule): - if module._context is not context or module in seen: - return - seen.add(module) order.append(module) for child in reversed(list(module.children())): - _collect_backward_order(child, context, order, seen) + _collect_backward_order(child, order) def _collect_owned_parameters(root_module: nn.Module) -> dict[str, nn.Parameter]: From ac48e297d20d069a1c6eda409425f5de23e2f032 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 1 Aug 2026 15:37:36 +0000 Subject: [PATCH 4/8] Clarify MFSDP subtree collection helper Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/module.py | 6 +++--- 1 file changed, 3 insertions(+), 3 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 87671f0c2bb..5500ac3daa1 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -74,7 +74,7 @@ def finalize(self) -> None: if module in visited: continue root_modules.append(module) - _collect_fsdp_modules(cast(nn.Module, module), visited) + _collect_fsdp_modules_under(cast(nn.Module, module), visited) root_modules.reverse() for root_module in root_modules: root_module._is_root = True @@ -341,7 +341,7 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str: return f"MFSDP {name} {phase}" -def _collect_fsdp_modules(module: nn.Module, modules: set["FsdpModule"]) -> None: +def _collect_fsdp_modules_under(module: nn.Module, modules: set["FsdpModule"]) -> None: """Collect FSDP modules reachable from ``module``.""" if isinstance(module, FsdpModule): if module in modules: @@ -349,7 +349,7 @@ def _collect_fsdp_modules(module: nn.Module, modules: set["FsdpModule"]) -> None modules.add(module) for child in module.children(): - _collect_fsdp_modules(child, modules) + _collect_fsdp_modules_under(child, modules) def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None: From a582bd901ca05913616f999ae93ed0fe9b797d2f Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 1 Aug 2026 15:46:00 +0000 Subject: [PATCH 5/8] Simplify MFSDP overlap tests Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 36 ++++++------------- 1 file changed, 11 insertions(+), 25 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 448facc751c..d548d2ea485 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -234,8 +234,6 @@ def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_ model = MultiChildModel(dim=dim, num_children=2).to(device) model.load_state_dict(baseline.state_dict()) - # Shard the child layers, then the model, so the children share a root context - # and reduce through the overlap path instead of as independent roots. with fully_shard_context(device=device) as context: for layer in model.layers: fully_shard(layer, mesh=mesh, placements=_hsdp_placements()) @@ -284,13 +282,10 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]: def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup): """HSDP reduce-scatters DP-inner every microbatch but all-reduces DP-outer once. - ``fully_shard(model)`` makes the child units share a root context so their - reductions run through the overlap path rather than as independent roots. - Counting linked NCCL kernels over a multi-microbatch step, the DP-inner reduce-scatter - fires once per microbatch per group while the DP-outer all-reduce that - finalizes main_grad fires only on the last microbatch, so the reduce-scatter - count is exactly ``num_microbatches`` times the all-reduce count. This asserts - on kernel counts only, not numerics. + Counting linked NCCL kernels over a multi-microbatch step, the DP-inner + reduce-scatter fires once per microbatch per group while the DP-outer + all-reduce that finalizes main_grad fires only on the last microbatch. This + asserts on kernel counts only, not numerics. """ world_size = distributed_setup.world_size device = distributed_setup.device @@ -399,7 +394,6 @@ def test_forward_peak_memory_bounds_in_flight_child_all_gathers(distributed_setu with fully_shard_context(device=device): 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(2, dim, device=device, dtype=dtype) with torch.no_grad(): @@ -553,7 +547,7 @@ def test_overlaps_communication_and_compute(distributed_setup, use_symm_mem): dp_group = dist.new_group(backend="nccl") mesh = DeviceMesh.from_group(dp_group, device.type) - model = MultiChildModel(dim=dim, num_children=num_children).to(dtype=dtype) + model = MultiChildModel(dim=dim, num_children=num_children).to(device=device, dtype=dtype) placements = _flat_placements() policy = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype) with fully_shard_context(device=device): @@ -565,13 +559,6 @@ def test_overlaps_communication_and_compute(distributed_setup, use_symm_mem): 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, - ) x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True) @@ -601,17 +588,16 @@ def train_one_iteration() -> None: allgather_kernels = collect_linked_kernels(prof, _ALL_GATHER_OP_NAME_SUBSTRING) reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING) - # The num_children child layers plus the root are each a sharded module; each does a - # 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 + # Each child layer does a 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. + expected_allgather_kernel_count = 0 if use_symm_mem else 2 * num_children 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]}" ) - assert len(reduce_scatter_kernels) == num_sharded_modules, ( - f"Expected {num_sharded_modules} reduce-scatter kernels, got " + assert len(reduce_scatter_kernels) == num_children, ( + f"Expected {num_children} reduce-scatter kernels, got " f"{len(reduce_scatter_kernels)}: {[kernel.name for kernel in reduce_scatter_kernels]}" ) From e223003bee13bba8267d02d560ea977f560d2b80 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 1 Aug 2026 15:58:25 +0000 Subject: [PATCH 6/8] Keep FsdpContext internal Signed-off-by: Jingyue Wu --- .../distributed/fsdp/src/megatron_fsdp/experimental/__init__.py | 2 -- 1 file changed, 2 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 01f798dc639..7c6dc8ef5ae 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -16,14 +16,12 @@ from .dbuffer import DBuffer from .fully_shard import fully_shard, fully_shard_context, microbatch -from .module import FsdpContext from .optimizer import fully_shard_optimizer from .placement import Flat, Partial, Placement, Placements, Replicate __all__ = [ "DBuffer", "Flat", - "FsdpContext", "Partial", "Placement", "Placements", From bcf2bdb84548d3728582b9a3aaecfffd45e8a771 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 5 Aug 2026 00:05:29 +0000 Subject: [PATCH 7/8] Guard MFSDP context construction Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/fully_shard.py | 6 +++ .../src/megatron_fsdp/experimental/module.py | 46 ++++++++----------- .../distributed/mfsdp_v2/test_context.py | 18 +++++++- 3 files changed, 43 insertions(+), 27 deletions(-) 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 0445735eb9c..135382ab1f1 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 @@ -86,6 +86,12 @@ def fully_shard( context = _FSDP_CONTEXT.get() if context is None: raise RuntimeError("fully_shard must run inside fully_shard_context.") + for submodule in module.modules(): + if isinstance(submodule, FsdpModule) and submodule.context is not context: + raise ValueError( + "Cannot fully_shard a module containing an FSDP child from another " + "fully_shard_context." + ) placements = _normalize_placements(mesh, placements) mixed_precision_policy = mixed_precision_policy or MixedPrecisionPolicy() 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 5500ac3daa1..877792dc960 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -68,26 +68,22 @@ def finalize(self) -> None: if self._is_finalized: raise RuntimeError("FSDP context is already finalized.") - visited: set[FsdpModule] = set() - root_modules: list[FsdpModule] = [] - for module in reversed(self._registered_modules): - if module in visited: - continue - root_modules.append(module) - _collect_fsdp_modules_under(cast(nn.Module, module), visited) - root_modules.reverse() - for root_module in root_modules: - root_module._is_root = True - - for root_module in root_modules: - for name, module in cast(nn.Module, root_module).named_modules(): + children: set[FsdpModule] = set() + for module in self._registered_modules: + _collect_fsdp_children(cast(nn.Module, module), children) + # FsdpModules that are not descendants of any other FsdpModule. + roots = [module for module in self._registered_modules if module not in children] + + for root in roots: + root._is_root = True + for name, module in cast(nn.Module, root).named_modules(): if not isinstance(module, FsdpModule): continue module._name = name self.forward_order.append(module) - for root_module in reversed(root_modules): - _collect_backward_order(cast(nn.Module, root_module), self.backward_order) + for root in reversed(roots): + _collect_backward_order(cast(nn.Module, root), self.backward_order) self._is_finalized = True @@ -341,17 +337,6 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str: return f"MFSDP {name} {phase}" -def _collect_fsdp_modules_under(module: nn.Module, modules: set["FsdpModule"]) -> None: - """Collect FSDP modules reachable from ``module``.""" - if isinstance(module, FsdpModule): - if module in modules: - return - modules.add(module) - - for child in module.children(): - _collect_fsdp_modules_under(child, modules) - - def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None: """Collect one root's static backward prefetch order.""" if isinstance(module, FsdpModule): @@ -361,6 +346,15 @@ def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"] _collect_backward_order(child, order) +def _collect_fsdp_children(module: nn.Module, children: set["FsdpModule"]) -> None: + """Collect the nearest FSDP descendants of ``module``.""" + for child in module.children(): + if isinstance(child, FsdpModule): + children.add(child) + else: + _collect_fsdp_children(child, children) + + def _collect_owned_parameters(root_module: nn.Module) -> dict[str, nn.Parameter]: parameters: dict[str, nn.Parameter] = {} diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_context.py b/tests/unit_tests/distributed/mfsdp_v2/test_context.py index 8dc07879102..f0bd1ce5fbe 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_context.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_context.py @@ -80,7 +80,7 @@ def test_child_then_parent_share_one_context(distributed_setup): device = distributed_setup.device mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) - model = NestedModel().to(device) + model = NestedModel() with fully_shard_context(device=device) as context: fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) @@ -218,3 +218,19 @@ def test_fully_shard_context_rejects_nesting(distributed_setup): assert model[0].context is outer_context assert model[1].context is outer_context + + +def test_fully_shard_rejects_child_from_another_context(distributed_setup): + """A parent cannot join a context different from an FSDP child context.""" + device = distributed_setup.device + mesh = init_device_mesh(device.type, (distributed_setup.world_size,)) + model = NestedModel() + + with fully_shard_context(device=device) as first_context: + fully_shard(model.inner, mesh=mesh, placements=_flat_placements()) + + with fully_shard_context(device=device): + with pytest.raises(ValueError, match="another fully_shard_context"): + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + assert model.inner.context is first_context From 0f315c6c69bf3b489a27d251f1ab11f2e4d08e5e Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 5 Aug 2026 00:14:25 +0000 Subject: [PATCH 8/8] Clear finalized MFSDP construction state Signed-off-by: Jingyue Wu --- .../distributed/fsdp/src/megatron_fsdp/experimental/module.py | 2 ++ 1 file changed, 2 insertions(+) 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 877792dc960..a1aaa6fc4e8 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -51,6 +51,7 @@ def __init__(self, device: torch.device) -> None: self.is_last_microbatch = True self.forward_order = IndexedOrder() self.backward_order = IndexedOrder() + # Construction-only; empty after finalization. self._registered_modules: list[FsdpModule] = [] self._is_finalized = False with torch.cuda.device(device): @@ -85,6 +86,7 @@ def finalize(self) -> None: for root in reversed(roots): _collect_backward_order(cast(nn.Module, root), self.backward_order) + self._registered_modules.clear() self._is_finalized = True def ensure_finalized(self) -> None: