From d41bcda8bcc37683a9378a5494861b69a1f90530 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 8 Jul 2026 00:53:00 +0000 Subject: [PATCH] Add FSDP NVTX annotations Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/module.py | 25 +++- .../distributed/mfsdp_v1/test_annotation.py | 123 ++++++++++++++++++ 2 files changed, 146 insertions(+), 2 deletions(-) create mode 100644 tests/unit_tests/distributed/mfsdp_v1/test_annotation.py 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 9c56a106d85..3f1d24b1517 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/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,6 +124,7 @@ 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: @@ -127,6 +132,14 @@ def context(self) -> FsdpContext: 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 "" + return f"MFSDP {name} {phase}" + def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int: if isinstance(axis, int): diff --git a/tests/unit_tests/distributed/mfsdp_v1/test_annotation.py b/tests/unit_tests/distributed/mfsdp_v1/test_annotation.py new file mode 100644 index 00000000000..9aee7734172 --- /dev/null +++ b/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", "", "forward"), + ("pop", "", "forward"), + ("push", "", "forward"), + ("pop", "", "forward"), + ("push", "", "backward"), + ("pop", "", "backward"), + ("push", "", "backward"), + ("pop", "", "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", "", "forward"), + ("push", "layers.0", "forward"), + ("pop", "layers.0", "forward"), + ("push", "layers.1", "forward"), + ("pop", "layers.1", "forward"), + ("pop", "", "forward"), + ("push", "", "backward"), + ("push", "layers.1", "backward"), + ("pop", "layers.1", "backward"), + ("push", "layers.0", "backward"), + ("pop", "layers.0", "backward"), + ("pop", "", "backward"), + ]