Reuse profiler helpers in mFSDP v2 symmetric memory tests

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -128,20 +128,10 @@ def _mb(num_bytes: int) -> str:
 # work to its enclosing collective or matmul operation.
 _ALL_GATHER_OP_NAME_SUBSTRING = "allgather"
 _REDUCE_SCATTER_OP_NAME_SUBSTRING = "reduce_scatter"
+_ALLREDUCE_OP_NAME_SUBSTRING = "allreduce"
 _GEMM_OP_NAME_SUBSTRING = "aten::mm"
 
 
-def _nccl_events(cuda_events, *name_fragments):
-    """CUDA NCCL events whose name contains any of ``name_fragments`` (case-insensitive)."""
-    return [
-        event
-        for event in cuda_events
-        if "nccl" in event.name.lower()
-        and event.activity_type == "kernel"
-        and any(fragment in event.name.lower() for fragment in name_fragments)
-    ]
-
-
 @pytest.mark.parametrize("num_microbatches", [1, 3])
 def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatches):
     """Minimal per-module FSDP training should match single-rank SGD."""
@@ -279,11 +269,11 @@ def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup):
 
     ``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 NCCL events over a multi-microbatch step, the DP-inner reduce-scatter
+    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 event counts only, not numerics.
+    on kernel counts only, not numerics.
     """
     world_size = distributed_setup.world_size
     device = distributed_setup.device
@@ -326,17 +316,16 @@ def train_one_step() -> None:
         train_one_step()
         torch.cuda.synchronize(device)
 
-    cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"]
-    reduce_scatter_events = _nccl_events(cuda_events, "reducescatter", "reduce_scatter")
-    all_reduce_events = _nccl_events(cuda_events, "allreduce")
+    reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING)
+    allreduce_kernels = collect_linked_kernels(prof, _ALLREDUCE_OP_NAME_SUBSTRING)
     # One DP-outer all-reduce per parameter group -- each child layer plus the
     # root unit's bias -- fired only on the last microbatch. Plain DP fires none.
-    assert len(all_reduce_events) == num_children + 1, [event.name for event in cuda_events]
+    assert len(allreduce_kernels) == num_children + 1, [event.name for event in prof.events()]
     # DP-inner reduce-scatter runs every microbatch; the DP-outer all-reduce runs
     # only on the last, so the counts differ by exactly the microbatch factor.
-    assert len(reduce_scatter_events) == len(all_reduce_events) * num_microbatches, (
-        f"Expected reduce-scatter ({len(reduce_scatter_events)}) to be {num_microbatches}x "
-        f"the DP-outer all-reduce count ({len(all_reduce_events)})."
+    assert len(reduce_scatter_kernels) == len(allreduce_kernels) * num_microbatches, (
+        f"Expected reduce-scatter ({len(reduce_scatter_kernels)}) to be {num_microbatches}x "
+        f"the DP-outer all-reduce count ({len(allreduce_kernels)})."
     )
 
 
@@ -570,7 +559,7 @@ def train_one_iteration() -> None:
     assert len(gemm_kernels) >= 3 * num_children, (
         f"Expected at least {3 * num_children} kernels linked to GEMMs, got "
         f"{len(gemm_kernels)}: "
-        f"{[event.name for event in gemm_kernels]}"
+        f"{[kernel.name for kernel in gemm_kernels]}"
     )
 
     allgather_kernels = collect_linked_kernels(prof, _ALL_GATHER_OP_NAME_SUBSTRING)
@@ -582,16 +571,16 @@ def train_one_iteration() -> None:
     expected_allgather_kernel_count = 0 if use_symm_mem else 2 * num_sharded_modules
     assert len(allgather_kernels) == expected_allgather_kernel_count, (
         f"Expected {expected_allgather_kernel_count} all-gather kernels, got "
-        f"{len(allgather_kernels)}: {[event.name for event in allgather_kernels]}"
+        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 "
-        f"{len(reduce_scatter_kernels)}: {[event.name for event in reduce_scatter_kernels]}"
+        f"{len(reduce_scatter_kernels)}: {[kernel.name for kernel in reduce_scatter_kernels]}"
     )
 
-    allgather_streams = {event.device_resource_id for event in allgather_kernels}
-    reduce_scatter_streams = {event.device_resource_id for event in reduce_scatter_kernels}
-    gemm_streams = {event.device_resource_id for event in gemm_kernels}
+    allgather_streams = {kernel.device_resource_id for kernel in allgather_kernels}
+    reduce_scatter_streams = {kernel.device_resource_id for kernel in reduce_scatter_kernels}
+    gemm_streams = {kernel.device_resource_id for kernel in gemm_kernels}
     if allgather_kernels:
         assert len(allgather_streams) == 1
     assert len(reduce_scatter_streams) == 1
@@ -600,11 +589,11 @@ def train_one_iteration() -> None:
     assert reduce_scatter_streams.isdisjoint(gemm_streams)
 
     allgather_overlap_count = sum(
-        any(events_overlap(event, gemm) for gemm in gemm_kernels) for event in allgather_kernels
+        any(events_overlap(kernel, gemm) for gemm in gemm_kernels) for kernel in allgather_kernels
     )
     reduce_scatter_overlap_count = sum(
-        any(events_overlap(event, gemm) for gemm in gemm_kernels)
-        for event in reduce_scatter_kernels
+        any(events_overlap(kernel, gemm) for gemm in gemm_kernels)
+        for kernel in reduce_scatter_kernels
     )
     expected_allgather_overlap = 2 * (num_children - 1)
     expected_reduce_scatter_overlap = num_children - 1

tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py

@@ -6,7 +6,6 @@
 import torch
 import torch.distributed as dist
 from torch import nn
-from torch.autograd import DeviceType
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.profiler import ProfilerActivity, profile
 
@@ -16,13 +15,16 @@
     Placements,
     fully_shard,
 )
+from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import collect_linked_kernels
 
 # Each sharded Linear's collective must be large enough that NCCL selects its
 # symmetric-memory (ncclSymk*) kernels over ring. Sub-KB collectives fall back to
 # ring on some platforms (e.g. CI with NCCL_NVLS_ENABLE=0), which would make the
 # symmetric-kernel assertions below fail; 1024-wide layers (a few-MiB bf16 weight)
 # reliably engage the symmetric kernels.
 _HIDDEN = 1024
+_ALL_GATHER_OP_NAME_SUBSTRING = "allgather"
+_REDUCE_SCATTER_OP_NAME_SUBSTRING = "reduce_scatter"
 
 
 class TinyModel(nn.Module):
@@ -43,18 +45,6 @@ def _flat_placements() -> Placements:
     return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()])
 
 
-def _kernels(prof: torch.profiler.profile) -> list[str]:
-    return [event.name for event in prof.events() if event.device_type == DeviceType.CUDA]
-
-
-def _is_symmetric_kernel(kernel: str) -> bool:
-    return "ncclSymk" in kernel
-
-
-def _count_symmetric_kernels(kernels: list[str], subname: str) -> int:
-    return sum(1 for kernel in kernels if _is_symmetric_kernel(kernel) and subname in kernel)
-
-
 @pytest.mark.parametrize("num_microbatches", [1, 3])
 def test_fully_shard_symmetric_memory_matches_default_and_profiles_nccl(
     distributed_setup, num_microbatches
@@ -108,11 +98,11 @@ def train(use_symm_mem: bool) -> list[torch.Tensor]:
 
         return losses
 
-    with profile(activities=[ProfilerActivity.CUDA]) as prof_without_symm_mem:
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof_without_symm_mem:
         losses_without_symm_mem = train(use_symm_mem=False)
         torch.cuda.synchronize()
 
-    with profile(activities=[ProfilerActivity.CUDA]) as prof_with_symm_mem:
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof_with_symm_mem:
         losses_with_symm_mem = train(use_symm_mem=True)
         torch.cuda.synchronize()
 
@@ -122,31 +112,44 @@ def train(use_symm_mem: bool) -> list[torch.Tensor]:
         msg="Symmetric-memory FSDP losses did not match default FSDP losses.",
     )
 
-    kernels_without_symm_mem = _kernels(prof_without_symm_mem)
-    assert _count_symmetric_kernels(kernels_without_symm_mem, "AllGather") == 0
-    assert _count_symmetric_kernels(kernels_without_symm_mem, "ReduceScatter") == 0
+    allgather_kernels_without_symm_mem = collect_linked_kernels(
+        prof_without_symm_mem, _ALL_GATHER_OP_NAME_SUBSTRING
+    )
+    reduce_scatter_kernels_without_symm_mem = collect_linked_kernels(
+        prof_without_symm_mem, _REDUCE_SCATTER_OP_NAME_SUBSTRING
+    )
+    assert all("ncclSymk" not in kernel.name for kernel in allgather_kernels_without_symm_mem)
+    assert all("ncclSymk" not in kernel.name for kernel in reduce_scatter_kernels_without_symm_mem)
 
-    kernels_with_symm_mem = _kernels(prof_with_symm_mem)
+    allgather_kernels_with_symm_mem = collect_linked_kernels(
+        prof_with_symm_mem, _ALL_GATHER_OP_NAME_SUBSTRING
+    )
+    reduce_scatter_kernels_with_symm_mem = collect_linked_kernels(
+        prof_with_symm_mem, _REDUCE_SCATTER_OP_NAME_SUBSTRING
+    )
     # 2 sharded modules (fc1, fc2), one reduce-scatter each per microbatch step.
     expected_reduce_scatter_kernel_count = num_training_steps * num_microbatches * 2
-    nccl_kernels_with_symm_mem = [
-        kernel for kernel in kernels_with_symm_mem if "nccl" in kernel.lower()
-    ]
-    assert (
-        _count_symmetric_kernels(kernels_with_symm_mem, "ReduceScatter")
-        == expected_reduce_scatter_kernel_count
-    ), (
+    assert len(reduce_scatter_kernels_with_symm_mem) == expected_reduce_scatter_kernel_count, (
         "Unexpected NCCL symmetric-memory reduce-scatter kernel count. "
-        f"Observed NCCL kernels: {nccl_kernels_with_symm_mem[:20]}"
+        f"Observed reduce-scatter kernels: "
+        f"{[kernel.name for kernel in reduce_scatter_kernels_with_symm_mem[:20]]}"
+    )
+    assert all("ncclSymk" in kernel.name for kernel in reduce_scatter_kernels_with_symm_mem), (
+        "Expected all symmetric-memory reduce-scatter kernels to be ncclSymk kernels. "
+        f"Observed reduce-scatter kernels: "
+        f"{[kernel.name for kernel in reduce_scatter_kernels_with_symm_mem[:20]]}"
     )
 
-    expected_all_gather_kernel_count = 2 * expected_reduce_scatter_kernel_count
-    assert (
-        _count_symmetric_kernels(kernels_with_symm_mem, "AllGather")
-        == expected_all_gather_kernel_count
-    ), (
+    expected_allgather_kernel_count = 2 * expected_reduce_scatter_kernel_count
+    assert len(allgather_kernels_with_symm_mem) == expected_allgather_kernel_count, (
         "Unexpected NCCL symmetric-memory all-gather kernel count. "
-        f"Observed NCCL kernels: {nccl_kernels_with_symm_mem[:20]}"
+        f"Observed all-gather kernels: "
+        f"{[kernel.name for kernel in allgather_kernels_with_symm_mem[:20]]}"
+    )
+    assert all("ncclSymk" in kernel.name for kernel in allgather_kernels_with_symm_mem), (
+        "Expected all symmetric-memory all-gather kernels to be ncclSymk kernels. "
+        f"Observed all-gather kernels: "
+        f"{[kernel.name for kernel in allgather_kernels_with_symm_mem[:20]]}"
     )
 
 
@@ -201,29 +204,32 @@ def test_fully_shard_zero_cta_moves_all_gather_to_copy_engine(distributed_setup)
     x = torch.randn(2, _HIDDEN, device=device, dtype=torch.bfloat16)
     target = torch.randn(2, _HIDDEN, device=device, dtype=torch.bfloat16)
 
-    with profile(activities=[ProfilerActivity.CUDA]) as prof:
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
         for _ in range(num_training_steps):
             optimizer.zero_grad()
             torch.nn.functional.mse_loss(model(x), target).backward()
             optimizer.step()
         torch.cuda.synchronize()
 
-    kernels = _kernels(prof)
-    nccl_kernels = [kernel for kernel in kernels if "nccl" in kernel.lower()]
+    allgather_kernels = collect_linked_kernels(prof, _ALL_GATHER_OP_NAME_SUBSTRING)
+    reduce_scatter_kernels = collect_linked_kernels(prof, _REDUCE_SCATTER_OP_NAME_SUBSTRING)
     # Zero-CTA moves the all-gather to the copy engine: no symmetric-memory all-gather kernel.
-    assert _count_symmetric_kernels(kernels, "AllGather") == 0, (
+    assert not allgather_kernels, (
         f"Expected no symmetric-memory all-gather kernel under zero-CTA. "
-        f"Observed NCCL kernels: {nccl_kernels[:20]}"
+        f"Observed all-gather kernels: {[kernel.name for kernel in allgather_kernels[:20]]}"
     )
     # The reduce-scatter's reduction cannot run on the copy engine, so it stays a
     # symmetric-memory kernel (an SM-launched NVLS multicast reduce): one per sharded
     # module (fc1, fc2) per training step.
     expected_reduce_scatter_kernel_count = num_training_steps * 2
-    assert (
-        _count_symmetric_kernels(kernels, "ReduceScatter") == expected_reduce_scatter_kernel_count
-    ), (
+    assert len(reduce_scatter_kernels) == expected_reduce_scatter_kernel_count, (
         f"Expected {expected_reduce_scatter_kernel_count} symmetric-memory reduce-scatter "
-        f"kernels under zero-CTA. Observed NCCL kernels: {nccl_kernels[:20]}"
+        f"kernels under zero-CTA. Observed reduce-scatter kernels: "
+        f"{[kernel.name for kernel in reduce_scatter_kernels[:20]]}"
+    )
+    assert all("ncclSymk" in kernel.name for kernel in reduce_scatter_kernels), (
+        "Expected all zero-CTA reduce-scatter kernels to be ncclSymk kernels. "
+        f"Observed reduce-scatter kernels: {[kernel.name for kernel in reduce_scatter_kernels[:20]]}"
     )
 
     # Release the dedicated communicator (leaks only on a test failure above, which is fine).