test_reduce_grads.py

import os

# Must pop before PyTorch imports
os.environ.pop('NCCL_MIN_NCHANNELS', None)
os.environ.pop('NCCL_MAX_NCHANNELS', None)
os.environ.pop('NCCL_NVLS_ENABLE', None)

import argparse
import torch
import torch.distributed as dist

import deep_ep
from deep_ep.utils.envs import dist_print, get_nvlink_gbs, get_physical_domain_size, init_dist
from deep_ep.utils.testing import bench, parse_num_bytes

# Period of `make_pattern`'s ramp, chosen so adjacent 12 KiB chunks have different contents
NUM_PATTERN_VALUES = (1 << 12) - 1

# Hidden sizes in fp32 elements, including full chunks, partial tails, and buffer reuse
CORRECTNESS_HIDDENS = (1 << 20, 1 << 18, 16392, 6144, 3080, 3072, 3064, 5120, 1536, 384, 24)


def make_pattern(hidden: int, key: int) -> torch.Tensor:
    """A deterministic fp32 blob: a ramp of period `NUM_PATTERN_VALUES` shifted by `key`.

    Every blob in the test carries a distinct key, so a swapped source shows up as a
    mismatch at every element. Adjacent 12 KiB chunks also have different contents, making
    a one-chunk offset error visible.

    The values are small integers on purpose. `lb_reduce_grads` sums redundant expert gradients into
    `expert_grads` with `cp.reduce.async.bulk`, an element-wise atomic add, so the order in
    which several borrowers land is not fixed. Values this size stay exact in fp32, which
    makes the expected value order-independent and lets the check use `torch.equal` rather
    than a tolerance.

    Plain arithmetic rather than an RNG: it must come out bit-identical on every GPU,
    and RNG launch configs depend on device properties.
    """
    idx = torch.arange(hidden, dtype=torch.int64, device='cuda') + key
    return ((idx % NUM_PATTERN_VALUES) - NUM_PATTERN_VALUES // 2).to(torch.float32)


def expert_grad_key(expert_idx: int) -> int:
    """Key of an original expert gradient. Non-negative, while `redundant_expert_grad_key` is negative, so the two
    families of blobs never share a key."""
    return expert_idx


def redundant_expert_grad_key(rank_idx: int, redundant_expert_idx: int, num_redundant_experts: int) -> int:
    """Key of a redundant expert gradient, unique over `(rank_idx, redundant_expert_idx)`."""
    return -1 - (rank_idx * num_redundant_experts + redundant_expert_idx)


def make_redundancy_mapping(num_ranks: int, num_redundant_experts: int, num_experts: int,
                            empty_ratio: float, seed: int, pattern: str = 'random') -> torch.Tensor:
    """The `[num_ranks, num_redundant_experts]` redundancy mapping, identical on every rank of the LSA domain.

    Both indices are domain-local: row `l` is the rank at position `l` of the domain, and an
    entry is an expert id in `[0, num_experts)`, the domain's own expert space.

    A rank never hosts a redundant instance of its local experts, so `random` samples each row from
    the `num_experts - num_local_experts` remote experts and shifts it past its own block.

    `ring` and `spread` are benchmarking patterns rather than realistic ones. Both give every
    rank exactly `num_redundant_experts` redundant gradients to pull, removing the imbalance that makes `random`'s GB/s
    hard to read, and they differ only in how many peers those gradients come from: `ring` puts
    all of them on one peer (rank `r` borrows only from rank `r + 1`), `spread` walks the
    owner across all `num_ranks - 1` peers. Comparing the two separates a per-peer link limit
    from an aggregate NVLink limit at identical bytes and identical balance.
    """
    # NOTES: everything here is pinned to CPU and explicitly typed, because `init_dist`
    # sets the default device to CUDA and the default dtype to bf16 -- a bf16 `torch.rand`
    # would silently change which redundant expert slots come out empty
    num_local_experts = num_experts // num_ranks
    if pattern in ('ring', 'spread'):
        borrower = torch.arange(num_ranks, dtype=torch.int32, device='cpu').unsqueeze(1)
        redundant_expert_indices = torch.arange(num_redundant_experts, dtype=torch.int32, device='cpu').unsqueeze(0)
        # The offset is never 0 mod num_ranks, so a rank never borrows from itself
        offset = 1 if pattern == 'ring' else 1 + redundant_expert_indices % (num_ranks - 1)
        table = (borrower + offset) % num_ranks * num_local_experts + redundant_expert_indices % num_local_experts
    else:
        assert pattern == 'random', f'Unknown pattern {pattern}'
        generator = torch.Generator().manual_seed(seed)
        table = torch.randint(0, num_experts - num_local_experts, (num_ranks, num_redundant_experts),
                              generator=generator, dtype=torch.int32, device='cpu')
        own_block_start = (torch.arange(num_ranks, dtype=torch.int32, device='cpu') * num_local_experts).unsqueeze(1)
        table += (table >= own_block_start).to(torch.int32) * num_local_experts
        empty = torch.rand((num_ranks, num_redundant_experts), generator=generator,
                           dtype=torch.float32, device='cpu') < empty_ratio
        table.masked_fill_(empty, -1)

    owners = torch.div(table, num_local_experts, rounding_mode='floor')
    assert ((table < 0) | (owners != torch.arange(num_ranks, dtype=torch.int32, device='cpu').unsqueeze(1))).all(), \
        'A rank must never hold a redundant instance of its local experts'
    return table.contiguous().cuda()


def expected_expert_grads(redundancy_mapping: torch.Tensor, rank_idx: int, num_ranks: int,
                          num_redundant_experts: int, num_local_experts: int, hidden: int) -> torch.Tensor:
    """What `expert_grads` must hold afterwards: its seed plus the gradients from every matching redundant expert.

    A rank never pulls from its own row.
    """
    expected = torch.stack([make_pattern(hidden, expert_grad_key(rank_idx * num_local_experts + local_expert_idx))
                            for local_expert_idx in range(num_local_experts)])
    for remote_rank_idx in range(num_ranks):
        if remote_rank_idx == rank_idx:
            continue
        for redundant_expert_idx in range(num_redundant_experts):
            expert_idx = int(redundancy_mapping[remote_rank_idx, redundant_expert_idx].item())
            if expert_idx >= 0 and expert_idx // num_local_experts == rank_idx:
                expected[expert_idx % num_local_experts] += \
                    make_pattern(hidden, redundant_expert_grad_key(remote_rank_idx, redundant_expert_idx, num_redundant_experts))
    return expected


@torch.inference_mode()
def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace) -> None:
    rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks)
    num_rdma_ranks, num_nvl_ranks = get_physical_domain_size(group)
    assert num_rdma_ranks == 1, 'Grad reduction pulls over NVLink only, so it is intra-node'

    # Settings
    num_experts = args.num_experts
    assert num_experts % num_ranks == 0
    num_local_experts = num_experts // num_ranks
    num_redundant_experts, num_sms = args.num_redundant_experts, args.num_sms

    # Correctness sweeps a table of hidden sizes; the benchmark runs on its own, much larger
    # one, so that its GB/s is not distorted by launch overhead.
    perf_hidden = args.num_hidden_bytes // torch.float32.itemsize
    hiddens = (*CORRECTNESS_HIDDENS, perf_hidden)

    assert all(h * torch.float32.itemsize % deep_ep.get_num_tma_alignment() == 0 for h in hiddens), \
        'Each expert gradient must be TMA-aligned'
    assert num_experts + num_ranks * num_redundant_experts < NUM_PATTERN_VALUES, \
        'Two blobs would share a key modulo the pattern period'

    # Plan every symmetric redundant expert gradient pool once, including the benchmark pools.
    plan = deep_ep.BufferAllocator()
    redundant_expert_grad_sets = [plan.allocate((num_redundant_experts, h), torch.float32) for h in hiddens]
    buffer = deep_ep.EPBuffer(group, num_bytes=deep_ep.get_num_allocation_alignment(),
                              lb_allocation_plan_or_num_bytes=plan, explicitly_destroy=True)
    num_sms = buffer.lb_get_theoretical_num_sms() if num_sms == 0 else num_sms

    dist_print(f'Config:\n'
               f' > Ranks: {num_rdma_ranks} x {num_nvl_ranks}\n'
               f' > Experts: {num_experts} ({num_local_experts} local), '
               f'redundant experts: {num_redundant_experts}\n'
               f' > Correctness hidden sweep: {list(CORRECTNESS_HIDDENS)} fp32\n'
               f' > Perf hidden: {perf_hidden} fp32 '
               f'({perf_hidden * torch.float32.itemsize / 1e6:.1f} MB per expert)\n'
               f' > #SM: {num_sms}\n',
               once_in_node=True)

    # This rank's original expert gradients, one set per hidden size. They stay
    # local -- the kernel only reduces into them, so they never have to be symmetric.
    expert_grad_sets = [torch.empty((num_local_experts, h), dtype=torch.float32, device='cuda') for h in hiddens]

    *sweep_expert_grad_sets, perf_expert_grads = expert_grad_sets
    *sweep_redundant_expert_grad_sets, perf_redundant_expert_grads = redundant_expert_grad_sets

    def seed_buffers(hidden: int, expert_grads: torch.Tensor, redundant_expert_grads: torch.Tensor) -> None:
        """Refill both sides. `expert_grads` accumulates, so it must be reseeded every round."""
        for local_expert_idx in range(num_local_experts):
            expert_grads[local_expert_idx] = make_pattern(
                hidden, expert_grad_key(rank_idx * num_local_experts + local_expert_idx))
        for redundant_expert_idx in range(num_redundant_experts):
            redundant_expert_grads[redundant_expert_idx] = make_pattern(hidden, redundant_expert_grad_key(rank_idx, redundant_expert_idx, num_redundant_experts))

    # Check correctness over every hidden size.
    dist_print(f'Checking correctness over {len(CORRECTNESS_HIDDENS)} hidden sizes '
               f'x {args.num_rounds} rounds ...', once_in_node=True)
    for hidden, expert_grads, redundant_expert_grads in zip(CORRECTNESS_HIDDENS, sweep_expert_grad_sets, sweep_redundant_expert_grad_sets, strict=True):
        for round_idx in range(args.num_rounds):
            redundancy_mapping = make_redundancy_mapping(num_ranks, num_redundant_experts, num_experts,
                                                   args.empty_ratio, args.seed + round_idx, args.pattern)

            # The seed must be visible to peers before anyone pulls
            seed_buffers(hidden, expert_grads, redundant_expert_grads)
            torch.cuda.synchronize()
            dist.barrier(group)

            buffer.lb_reduce_grads(
                redundant_expert_grads, expert_grads, redundancy_mapping, num_sms=num_sms,
                previous_event=buffer.capture() if round_idx % 2 else None).wait()
            torch.cuda.synchronize()

            expected = expected_expert_grads(redundancy_mapping, rank_idx, num_ranks,
                                          num_redundant_experts, num_local_experts, hidden)
            for local_expert_idx in range(num_local_experts):
                assert torch.equal(expert_grads[local_expert_idx], expected[local_expert_idx]), \
                    f'{hidden=}, {round_idx=}: gradient of local expert ' \
                    f'{local_expert_idx} is wrong ' \
                    f'({int((expert_grads[local_expert_idx] != expected[local_expert_idx]).sum().item())} ' \
                    f'of {hidden} elements)'

            # Never writes the redundant expert gradient pool. A swapped TMA source/destination would land
            # here rather than in the check above, so it is worth testing explicitly.
            dist.barrier(group)
            for redundant_expert_idx in range(num_redundant_experts):
                assert torch.equal(redundant_expert_grads[redundant_expert_idx],
                                   make_pattern(hidden, redundant_expert_grad_key(rank_idx, redundant_expert_idx, num_redundant_experts))), \
                    f'{hidden=}, {round_idx=}: redundant expert gradient {redundant_expert_idx} was modified'

    # Perf: count the redundant experts contributing gradients to each rank.
    # That is the same set `lb_prefetch_weights` pushes, just read from the other end.
    redundancy_mapping = make_redundancy_mapping(num_ranks, num_redundant_experts, num_experts,
                                           args.empty_ratio, args.seed, args.pattern)
    owners = torch.div(redundancy_mapping, num_local_experts, rounding_mode='floor')
    redundant_experts_per_owner = [int(((redundancy_mapping >= 0) & (owners == r)).sum().item()) for r in range(num_ranks)]
    num_bytes_per_redundant_expert = perf_hidden * torch.float32.itemsize
    num_pulled_bytes = redundant_experts_per_owner[rank_idx] * num_bytes_per_redundant_expert
    num_critical_bytes = max(redundant_experts_per_owner) * num_bytes_per_redundant_expert

    # Values are never checked here, so a zero fill is enough -- it only keeps the repeated
    # accumulation away from whatever `torch.empty` left in the pages.
    perf_expert_grads.zero_()
    perf_redundant_expert_grads.zero_()
    torch.cuda.synchronize()
    dist.barrier(group)

    t, _, _ = bench(lambda: buffer.lb_reduce_grads(
        perf_redundant_expert_grads, perf_expert_grads, redundancy_mapping, num_sms=num_sms).wait())
    dist_print('Performance:', once_in_node=True)
    dist_print(f' > Redundant experts pulled per rank: {redundant_experts_per_owner} (this rank {redundant_experts_per_owner[rank_idx]})',
               once_in_node=True)
    dist_print(f' > Rank {rank_idx} perf ({num_sms=}): '
               f'{t * 1e6:7.1f} us | '
               f'this rank {num_pulled_bytes / 1e6:.1f} MB @ {num_pulled_bytes / t / 1e9:5.1f} GB/s | '
               f'critical path {num_critical_bytes / 1e6:.1f} MB @ '
               f'{num_critical_bytes / t / 1e9:5.1f} GB/s '
               f'({num_critical_bytes / t / 1e9 / get_nvlink_gbs() * 100:5.1f}% of NVLink)')

    # Destroy. Redundant expert gradient pools are views into the buffer, so release them first.
    del plan, redundant_expert_grad_sets, sweep_redundant_expert_grad_sets, perf_redundant_expert_grads
    buffer.destroy()
    dist.destroy_process_group()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Test expert grad reduction')
    parser.add_argument('--num-processes', type=int, default=8)
    parser.add_argument('--num-experts', type=int, default=64)
    parser.add_argument('--num-redundant-experts', type=int, default=8)
    parser.add_argument('--num-hidden-bytes', type=parse_num_bytes, default=64 << 20,
                        help='Grad bytes per expert in the benchmark, with an optional binary '
                             'suffix: 64M (default), 512K, or plain bytes.')
    parser.add_argument('--num-rounds', type=int, default=2,
                        help='Copy tables per hidden size, on top of the size sweep')
    parser.add_argument('--empty-ratio', type=float, default=0.25)
    parser.add_argument('--pattern', type=str, default='random', choices=('random', 'ring', 'spread'))
    parser.add_argument('--num-sms', type=int, default=0, help='0 uses the bandwidth-based SM estimate')
    parser.add_argument('--seed', type=int, default=0)
    args = parser.parse_args()
    torch.multiprocessing.spawn(test, args=(args.num_processes, args), nprocs=args.num_processes)