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 2d3cc2a677a..081cc5c260f 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/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 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 13dc1744628..87207548140 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py +++ b/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,6 +15,7 @@ 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 @@ -23,6 +23,8 @@ # 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).