Add FSDP NVTX annotations

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

@@ -15,7 +15,7 @@
 """Module mixin for the minimal Megatron-FSDP path."""
 
 from collections.abc import Callable
-from typing import cast
+from typing import Literal, cast
 
 import torch
 from torch import nn
@@ -48,6 +48,9 @@ def __init__(self, root_module: "FsdpModule") -> None:
 class FsdpModule:
     """Mixin attached to modules managed by the minimal FSDP path."""
 
+    # Name relative to the root FSDP module from named_modules().
+    # Root uses "" and None means uninitialized.
+    _name: str | None
     _parameter_groups: tuple[FsdpParameterGroup, ...]
     _context: FsdpContext | None
     _ready_grad_parameters: set[nn.Parameter]
@@ -62,6 +65,7 @@ def __init__(
     ) -> None:
         """Initialize FSDP runtime state on an already-constructed module."""
         self._context = None
+        self._name = None
         owned_parameters = _collect_owned_parameters(self)
         axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes)
         assert axis_indices == tuple(
@@ -111,7 +115,7 @@ def _lazy_init_context(self) -> None:
             return
 
         context = FsdpContext(root_module=self)
-        for submodule in cast(nn.Module, self).modules():
+        for submodule_name, submodule in cast(nn.Module, self).named_modules():
             if not isinstance(submodule, FsdpModule):
                 continue
             if submodule._context is not None:
@@ -120,13 +124,22 @@ def _lazy_init_context(self) -> None:
                     "Run forward through the root FSDP module first."
                 )
             submodule._context = context
+            submodule._name = submodule_name
 
     @property
     def context(self) -> FsdpContext:
         """Return the initialized runtime context."""
         assert self._context is not None
         return self._context
 
+    @property
+    def name(self) -> str:
+        """Return this FSDP unit's name."""
+        name = self._name
+        if name is None:
+            raise RuntimeError("FSDP module name has not been initialized.")
+        return name
+
     def is_root(self) -> bool:
         """Return whether this module is the outermost FSDP unit in its context."""
         return self.context.root_module is self
@@ -157,6 +170,7 @@ def grad_hook(_parameter: nn.Parameter) -> None:
     def pre_forward(self) -> None:
         """Prepare full parameters for forward compute."""
         self._lazy_init_context()
+        torch.cuda.nvtx.range_push(self._nvtx_label("forward"))
         self._ready_grad_parameters.clear()
         for group in self._parameter_groups:
             group.sync_model_weight_from_main_weight()
@@ -166,9 +180,11 @@ def post_forward(self) -> None:
         """Return parameters to their sharded resting state after forward compute."""
         for group in self._parameter_groups:
             group.reshard_parameters()
+        torch.cuda.nvtx.range_pop()
 
     def pre_backward(self) -> None:
         """Prepare full parameters for backward compute."""
+        torch.cuda.nvtx.range_push(self._nvtx_label("backward"))
         for group in self._parameter_groups:
             group.unshard_parameters()
 
@@ -179,11 +195,16 @@ def post_backward(self) -> None:
                 group.reduce_gradients()
             group.reshard_parameters()
         self._ready_grad_parameters.clear()
+        torch.cuda.nvtx.range_pop()
 
     def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]:
         """Return parameter groups owned by this FSDP unit."""
         return self._parameter_groups
 
+    def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str:
+        name = self.name if self.name else "<root>"
+        return f"MFSDP {name} {phase}"
+
 
 def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int:
     if isinstance(axis, int):

tests/unit_tests/distributed/mfsdp_v1/test_annotation.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Unit tests for experimental Megatron-FSDP annotations."""
+
+import re
+from typing import Literal, NamedTuple
+
+import pytest
+import torch
+from torch import nn
+from torch.distributed.device_mesh import init_device_mesh
+
+from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
+    Flat,
+    Placements,
+    fully_shard,
+)
+
+_NVTX_LABEL_PATTERN = re.compile(r"MFSDP (.+) (forward|backward)")
+
+
+class NvtxEvent(NamedTuple):
+    kind: Literal["push", "pop"]
+    name: str
+    phase: str
+
+
+class NestedLinearModel(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.bias = nn.Parameter(torch.ones(dim))
+        self.layers = nn.ModuleList([nn.Linear(dim, dim, bias=False) for _ in range(2)])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.bias
+        for layer in self.layers:
+            x = torch.relu(layer(x))
+        return x
+
+
+def _flat_placements() -> Placements:
+    return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()])
+
+
+def _setup_nvtx_recording(monkeypatch: pytest.MonkeyPatch, events: list[NvtxEvent]) -> None:
+    label_stack: list[tuple[str, str]] = []
+
+    def parse_nvtx_label(label: str) -> tuple[str, str]:
+        match = _NVTX_LABEL_PATTERN.fullmatch(label)
+        assert match is not None
+        return match.groups()
+
+    def record_push(label: str) -> None:
+        name, phase = parse_nvtx_label(label)
+        label_stack.append((name, phase))
+        events.append(NvtxEvent("push", name, phase))
+
+    def record_pop() -> None:
+        name, phase = label_stack.pop()
+        events.append(NvtxEvent("pop", name, phase))
+
+    monkeypatch.setattr(torch.cuda.nvtx, "range_push", record_push)
+    monkeypatch.setattr(torch.cuda.nvtx, "range_pop", record_pop)
+
+
+def _get_distributed_setup(request: pytest.FixtureRequest):
+    try:
+        return request.getfixturevalue("distributed_setup")
+    except pytest.FixtureLookupError:
+        pytest.skip("distributed_setup fixture is only available in the Megatron-FSDP test bucket")
+
+
+def test_fsdp_sibling_roots_emit_root_nvtx_ranges_after_training_step(request, monkeypatch):
+    """Independent FSDP roots should each emit root-labeled NVTX ranges."""
+    distributed_setup = _get_distributed_setup(request)
+    events: list[NvtxEvent] = []
+    _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())
+
+    model(torch.ones(2, 4, device=distributed_setup.device)).sum().backward()
+
+    assert [(event.kind, event.name, event.phase) for event in events] == [
+        ("push", "<root>", "forward"),
+        ("pop", "<root>", "forward"),
+        ("push", "<root>", "forward"),
+        ("pop", "<root>", "forward"),
+        ("push", "<root>", "backward"),
+        ("pop", "<root>", "backward"),
+        ("push", "<root>", "backward"),
+        ("pop", "<root>", "backward"),
+    ]
+
+
+def test_fsdp_training_hooks_emit_stacked_nvtx_ranges(request, monkeypatch):
+    """Nested training hooks should emit concise NVTX ranges."""
+    distributed_setup = _get_distributed_setup(request)
+    events: list[NvtxEvent] = []
+    _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())
+
+    model(torch.ones(2, 4, device=distributed_setup.device)).sum().backward()
+
+    assert [(event.kind, event.name, event.phase) for event in events] == [
+        ("push", "<root>", "forward"),
+        ("push", "layers.0", "forward"),
+        ("pop", "layers.0", "forward"),
+        ("push", "layers.1", "forward"),
+        ("pop", "layers.1", "forward"),
+        ("pop", "<root>", "forward"),
+        ("push", "<root>", "backward"),
+        ("push", "layers.1", "backward"),
+        ("pop", "layers.1", "backward"),
+        ("push", "layers.0", "backward"),
+        ("pop", "layers.0", "backward"),
+        ("pop", "<root>", "backward"),
+    ]