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..7c6dc8ef5ae 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_context, 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_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..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
@@ -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,38 @@
 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.
+    """
+    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}.")
+
+    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 +83,15 @@ 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.")
+    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()
@@ -58,6 +101,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 +134,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 +142,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 +161,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..a1aaa6fc4e8 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,66 @@ 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"
     # 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.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):
             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.")
+
+        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 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:
+        """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,8 +122,9 @@ 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
+    _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
@@ -94,13 +133,15 @@ 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._is_root = False
         self._name = None
         self._unshard_event = None
         owned_parameters = _collect_owned_parameters(self)
@@ -124,60 +165,11 @@ 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:
-        """Return the initialized runtime context."""
-        assert self._context is not None
+        """Return the FSDP context."""
         return self._context
 
     @property
@@ -189,8 +181,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._is_root
 
     def _register_hooks(self) -> None:
         module = cast(nn.Module, self)
@@ -227,10 +219,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()
 
@@ -348,7 +340,7 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str:
 
 
 def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None:
-    """Collect FsdpModules in static backward prefetch order."""
+    """Collect one root's static backward prefetch order."""
     if isinstance(module, FsdpModule):
         order.append(module)
 
@@ -356,6 +348,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_annotation.py b/tests/unit_tests/distributed/mfsdp_v2/test_annotation.py
index f65e0b85724..e5624bfb501 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()
 
@@ -197,7 +202,8 @@ def test_tied_child_parameters_complete_backward_once_per_cycle(distributed_setu
     _setup_nvtx_recording(monkeypatch, events)
     model = TiedLM()
     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())
 
     token_ids = torch.arange(8, device=distributed_setup.device).reshape(2, 4)
     for _ in range(2):
diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_context.py b/tests/unit_tests/distributed/mfsdp_v2/test_context.py
index 102b2fde332..f0bd1ce5fbe 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)
+    model = NestedModel()
 
-    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,81 @@ 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 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]
+
+
+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_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(2)]).to(device)
+
+    with fully_shard_context(device=device):
+        fully_shard(model[0], mesh=mesh, placements=_flat_placements())
+        outer_context = model[0].context
+        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[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
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..3a248848dc7 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)
 
@@ -232,11 +234,10 @@ 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.
-    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 +253,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())
@@ -281,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
@@ -303,9 +301,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 +317,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 +352,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
@@ -367,7 +367,8 @@ def test_tied_child_parameters_allocate_one_physical_weight(distributed_setup):
     """Tied registrations should allocate one DBuffer entry and optimizer parameter."""
     model = TiedLM()
     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())
 
     (parameter_group,) = model.parameter_groups
     (parameter,) = parameter_group.fsdp_parameters
@@ -391,9 +392,9 @@ 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)
 
     x = torch.randn(2, dim, device=device, dtype=dtype)
     with torch.no_grad():
@@ -437,9 +438,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 +477,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)
@@ -545,24 +548,18 @@ 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)
-    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_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,
+            )
 
     x = torch.randn(4096, dim, device=device, dtype=dtype, requires_grad=True)
 
@@ -592,17 +589,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]}"
     )
 
@@ -648,9 +644,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 +671,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 +692,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 +717,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 +753,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 +798,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 +823,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())
 
-    for layer in model:
-        assert layer.context.is_last_microbatch
+    with microbatch(context, is_last=False):
+        assert not context.is_last_microbatch
+
+    assert context.is_last_microbatch
 
 
 def test_cpu_initialized_parameters_shard_to_mesh_device(distributed_setup):
@@ -854,7 +855,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)
@@ -877,7 +879,8 @@ def test_meta_parameters_shard_to_mesh_device(distributed_setup):
         nn.Linear(4, 4, bias=False, device="meta", dtype=torch.bfloat16),
     )
 
-    fully_shard(model, mesh=mesh, placements=_flat_placements())
+    with fully_shard_context(device=device):
+        fully_shard(model, mesh=mesh, placements=_flat_placements())
 
     with torch.no_grad():
         model[0].weight.fill_(2.0)
@@ -901,7 +904,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 +957,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)