TheRock/rocm-systems/projects/rccl/bindings/nccl4py/tests/test_collectives.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Collective-correctness tests for the RCCL fork.

Seven tests, all marked ``mpi``, each compares the binding's collective
result against a ``numpy`` (or ``torch``-on-CPU) reference within
``atol=1e-3 rtol=1e-3``. Run with::

    mpirun -np 2 --allow-run-as-root pytest -m mpi tests/test_collectives.py
    mpirun -np 4 --allow-run-as-root pytest -m mpi tests/test_collectives.py
    mpirun -np 8 --allow-run-as-root pytest -m mpi tests/test_collectives.py

The 8-GPU configuration is xfailed (with a recorded reason) on hosts
where fewer than 8 GPUs are visible; see ``conftest.py``.
"""

from __future__ import annotations

import numpy as np
import pytest

torch = pytest.importorskip("torch", reason="torch is required for the GPU buffers")

import nccl.core as nccl  # noqa: E402

from conftest import require_at_least_n_ranks, xfail_if_fewer_than_8_gpus  # noqa: E402

pytestmark = [
    pytest.mark.mpi,
    xfail_if_fewer_than_8_gpus(),  # only triggers on the 8-GPU launch
]

# Element count chosen to be small enough to run quickly on every rank
# count, large enough to exercise multi-stage ring algorithms.
_N = 4096

# Tolerances. The default `atol=1e-3 rtol=1e-3` is achievable for fp16
# / fp32 but mathematically impossible for bf16: 1 ULP in bf16 at value
# 1 is 2**-7 ~ 7.8e-3, and a K-rank sum accumulates O((K-1)*ULP) round-
# off error against an fp32 reference. To keep the test passing
# uniformly across {2, 4, 8} ranks without inflating the tolerances
# absurdly, we (1) normalize per-rank inputs by 1/nranks so the
# all_reduce sum stays bounded in [-1, 1] regardless of rank count;
# (2) scale the bf16 tolerance with nranks. Neither indicates a bug
# in the binding.
_BF16_EPS = 2**-7  # 1 ULP in bf16 at magnitude 1.0


def _tol_for(dtype, nranks: int) -> tuple[float, float]:
    if dtype == torch.bfloat16:
        atol = max(5e-3, (nranks - 1) * _BF16_EPS)
        rtol = 5e-2
        return atol, rtol
    # fp16 / fp32: literal tolerances. fp16 has eps = 2**-10 so K-1
    # accumulated ULPs is well under 1e-3 for K <= 8 with the 1/nranks
    # input scaling applied below.
    return 1e-3, 1e-3


def _seeded_tensor(rank: int, dtype, device, n: int = _N, *, scale: float = 1.0) -> torch.Tensor:
    """Per-rank deterministic random tensor in ``[-scale, scale]``."""
    g = torch.Generator(device="cpu").manual_seed(0xBEEF + rank)
    cpu = torch.empty(n, dtype=torch.float32).uniform_(-scale, scale, generator=g)
    return cpu.to(device=device, dtype=dtype)


def _allclose_or_fail(actual: torch.Tensor, expected: torch.Tensor, *, nranks: int = 1) -> None:
    atol, rtol = _tol_for(actual.dtype, nranks)
    actual32 = actual.detach().to(dtype=torch.float32, device="cpu")
    expected32 = expected.detach().to(dtype=torch.float32, device="cpu")
    torch.testing.assert_close(actual32, expected32, atol=atol, rtol=rtol)


# ---------------------------------------------------------------------------
# all_reduce: fp16 / bf16 / fp32 (3 tests)
# ---------------------------------------------------------------------------


def _run_all_reduce(comm, mpi_world, rank_info, device, dtype):
    rank, nranks = rank_info
    # Scale per-rank inputs by 1/nranks so the all_reduce result stays in
    # [-1, 1] for any nranks; otherwise the peak intermediate sum in the
    # ring-allreduce grows with nranks, the ULP at the peak grows too,
    # and the round-off floor of low-precision dtypes blows past any
    # fixed atol. This is the same trick rccl-tests applies via the
    # `--scale` option to its random-data verifier.
    send = _seeded_tensor(rank, dtype, device, scale=1.0 / nranks)
    recv = torch.empty_like(send)

    comm.reduce(send, recv, nccl.SUM)  # root=None -> all_reduce
    torch.cuda.synchronize()

    # Reference: gather every rank's send buffer on every rank via mpi4py
    # (cheap because we operate on small tensors and a CPU copy).
    cpu_send = send.detach().to(dtype=torch.float32, device="cpu").numpy()
    gathered = np.stack(mpi_world.allgather(cpu_send), axis=0)
    expected_np = gathered.sum(axis=0)
    expected = torch.from_numpy(expected_np)

    _allclose_or_fail(recv, expected, nranks=nranks)


def test_all_reduce_fp16(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    _run_all_reduce(comm, mpi_world, rank_info, device, torch.float16)


def test_all_reduce_bf16(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    _run_all_reduce(comm, mpi_world, rank_info, device, torch.bfloat16)


def test_all_reduce_fp32(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    _run_all_reduce(comm, mpi_world, rank_info, device, torch.float32)


# ---------------------------------------------------------------------------
# broadcast (1 test)
# ---------------------------------------------------------------------------


def test_broadcast(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    rank, _ = rank_info
    root = 0
    # Root has the source tensor; non-root ranks start from zeros and
    # receive the root's payload.
    src = _seeded_tensor(0, torch.float32, device)
    if rank == root:
        send = src.clone()
        recv = torch.empty_like(send)
    else:
        send = torch.zeros_like(src)  # ignored on non-root ranks
        recv = torch.zeros_like(src)

    comm.broadcast(send, recv, root=root)
    torch.cuda.synchronize()

    _allclose_or_fail(recv, src)


# ---------------------------------------------------------------------------
# all_gather (1 test)
# ---------------------------------------------------------------------------


def test_all_gather(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    rank, nranks = rank_info
    per_rank = _N // nranks
    send = _seeded_tensor(rank, torch.float32, device, n=per_rank)
    recv = torch.empty(per_rank * nranks, dtype=torch.float32, device=device)

    comm.allgather(send, recv)
    torch.cuda.synchronize()

    # Reference: numpy concat of every rank's send buffer in rank order.
    cpu_send = send.detach().to(dtype=torch.float32, device="cpu").numpy()
    expected_np = np.concatenate(mpi_world.allgather(cpu_send), axis=0)
    expected = torch.from_numpy(expected_np)

    _allclose_or_fail(recv, expected)


# ---------------------------------------------------------------------------
# reduce_scatter (1 test)
# ---------------------------------------------------------------------------


def test_reduce_scatter(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    rank, nranks = rank_info
    per_rank = _N // nranks
    total = per_rank * nranks  # divisible
    send = _seeded_tensor(rank, torch.float32, device, n=total)
    recv = torch.empty(per_rank, dtype=torch.float32, device=device)

    comm.reduce_scatter(send, recv, nccl.SUM)
    torch.cuda.synchronize()

    cpu_send = send.detach().to(dtype=torch.float32, device="cpu").numpy()
    gathered = np.stack(mpi_world.allgather(cpu_send), axis=0)  # (nranks, total)
    summed = gathered.sum(axis=0)
    expected_np = summed[rank * per_rank : (rank + 1) * per_rank]
    expected = torch.from_numpy(expected_np)

    _allclose_or_fail(recv, expected)


# ---------------------------------------------------------------------------
# send / recv (1 test, pairs adjacent ranks)
# ---------------------------------------------------------------------------


def test_send_recv(comm, mpi_world, rank_info, device):
    require_at_least_n_ranks(rank_info, 2)
    rank, nranks = rank_info
    # Pair rank r with rank r ^ 1 inside the first even pairs only; ranks
    # without a partner (when nranks is odd) sit out the test.
    partner = rank ^ 1
    if partner >= nranks:
        pytest.skip(f"rank {rank} has no partner (nranks={nranks})")

    send = _seeded_tensor(rank, torch.float32, device)
    expected_partner_send = _seeded_tensor(partner, torch.float32, device)
    recv = torch.zeros_like(send)

    with nccl.group():
        comm.send(send, peer=partner)
        comm.recv(recv, peer=partner)
    torch.cuda.synchronize()

    _allclose_or_fail(recv, expected_partner_send)