recipes/docker-images/glm53-kindling/experimental/arx/arx.py

"""All-reduce over RoCE for GB10 tensor-parallel groups, without NCCL.

NCCL costs about 80-90 us per decode-sized all-reduce on these boxes; this
costs 13 us at 8 KB and 27 us at 64 KB. Each rank writes its partial to its
peers with RDMA over both ConnectX roots (see arx_vllm.cu). Larger all-reduces,
which are prefill-sized, stay on NCCL, whose bandwidth wins there.

Enabled with VLLM_ARX_ALLREDUCE=1. The RDMA devices and GID index come from
NCCL_IB_HCA (exactly two, one per root, in the same subnet order on every
rank) and NCCL_IB_GID_INDEX.

Ring mode (VLLM_ARX_RING=1) is for groups of 2, 3 or 4 cabled in a ring with
no switch, with rank r's "next" port cabled to rank r+1. arx and arxbig open
QPs only to r-1 and r+1, and at 4 the rank in between relays raw data for the
rank beyond, so results match mesh mode bit for bit (see the .cu files). At 3
both other ranks are neighbours and nothing is relayed. arxbig's ring
reduce-scatter needs 2 or 4 ranks, so a ring of 3 runs arx alone
(VLLM_ARXBIG=0). ARX_RING_PREV_HCAS and ARX_RING_NEXT_HCAS name the two RDMA
devices of the port facing each neighbour, in the order whose root r shares a
subnet with root r at that neighbour: root 0 then root 1 when every box is
cabled the same way, the other order on a box whose cabling is mirrored. Each
port has its own subnet, so each device uses the RoCE v2 GID of the IPv4
address on its own netdev.
"""
import os

import torch
import torch.distributed as dist
from torch.distributed import ProcessGroup

from vllm.logger import init_logger

logger = init_logger(__name__)

_ext = None
_taken = False  # the extension holds one connection per process


def _hcas(var: str) -> list[str]:
    return [h.strip("=^").split(":")[0] for h in os.environ.get(var, "").split(",") if h]


def roce_v2_gid(dev: str) -> int:
    """GID index of dev's RoCE v2 GID for an IPv4 address on dev's own netdev.

    The kernel fills the GID table from the netdev's addresses, and ndevs
    names the netdev an entry came from, so a VLAN or other upper device's
    entries are skipped.
    """
    base = f"/sys/class/infiniband/{dev}"
    netdevs = set(os.listdir(f"{base}/device/net"))
    port = f"{base}/ports/1"
    found = []
    for name in sorted(os.listdir(f"{port}/gids"), key=int):
        try:
            with open(f"{port}/gids/{name}") as f:
                gid = f.read().strip()
            with open(f"{port}/gid_attrs/types/{name}") as f:
                kind = f.read().strip()
            with open(f"{port}/gid_attrs/ndevs/{name}") as f:
                ndev = f.read().strip()
        except OSError:  # empty entries refuse to report their attributes
            continue
        if kind == "RoCE v2" and ndev in netdevs and gid.startswith("0000:0000:0000:0000:0000:ffff:"):
            found.append(int(name))
    if not found:
        raise RuntimeError(f"arx: {dev} has no RoCE v2 GID for an IPv4 address on {sorted(netdevs)}")
    if len(found) > 1:
        logger.warning("arx: %s has IPv4 GIDs at %s; using %d", dev, found, found[0])
    return found[0]


def fabric(world: int, who: str):
    """(devices, GID index per device, ring) for this rank, or None to use NCCL."""
    if os.environ.get("VLLM_ARX_RING") == "1":
        prev, nxt = _hcas("ARX_RING_PREV_HCAS"), _hcas("ARX_RING_NEXT_HCAS")
        if len(prev) != 2 or len(nxt) != 2:
            logger.warning("%s ring mode needs two devices each in ARX_RING_PREV_HCAS and ARX_RING_NEXT_HCAS, "
                           "got %r and %r; using NCCL", who, prev, nxt)
            return None
        if world not in (2, 3, 4):
            raise ValueError(f"{who} ring mode supports groups of 2 to 4 ranks, not {world}")
        devs = prev + nxt
        return devs, [roce_v2_gid(d) for d in devs], True
    hcas = _hcas("NCCL_IB_HCA")
    if len(hcas) != 2:
        logger.warning("%s needs two RDMA devices in NCCL_IB_HCA, got %r; using NCCL", who, hcas)
        return None
    gid = int(os.environ.get("NCCL_IB_GID_INDEX", "5"))
    return hcas, [gid, gid], False


def _jit_load(name, *args, **kwargs):
    """torch's cpp_extension.load(), safe after a build was killed midway (#53).

    torch marks a build with a `lock` file and waits forever on one it finds, so a
    build killed midway hangs every later load with nothing in the log. Builds here
    also hold an flock, which the kernel drops when its holder dies: whoever holds
    it is the only builder, and any `lock` file it finds is stale."""
    import fcntl

    from torch.utils.cpp_extension import _get_build_directory, load

    build = _get_build_directory(name, False)
    with open(os.path.join(build, "lock.flock"), "w") as f:
        fcntl.flock(f, fcntl.LOCK_EX)
        try:
            os.remove(os.path.join(build, "lock"))
        except FileNotFoundError:
            pass
        return load(name, *args, **kwargs)


def _load():
    global _ext
    if _ext is None:
        _ext = _jit_load(
            "arx_vllm",
            [os.path.join(os.path.dirname(__file__), "arx_vllm.cu")],
            extra_cuda_cflags=["-O3", "-std=c++17", "-gencode=arch=compute_121a,code=sm_121a"],
            extra_ldflags=["-libverbs"],
        )
    return _ext


# (address, bytes) ranges for the next all-reduce to prefetch into L2 while it waits.
_prefetch: tuple = ()


def set_prefetch(tensors) -> None:
    """The next arx all-reduce asks L2 for these tensors' bytes while it waits for peers."""
    global _prefetch
    _prefetch = tuple((t.data_ptr(), t.numel() * t.element_size()) for t in tensors if t is not None)


class ArxCommunicator:
    def __init__(self, group: ProcessGroup, device: torch.device):
        global _taken
        self.disabled = True
        if _taken:
            logger.warning("arx already serves another group in this process; using NCCL here")
            return
        rank, world = dist.get_rank(group), dist.get_world_size(group)
        fab = fabric(world, "arx")
        if fab is None:
            return
        devs, gids, ring = fab
        with torch.cuda.device(device):
            ext = _load()
            info = ext.prepare(rank, world, devs, gids, ring)
            infos: list[bytes] = [b""] * world
            dist.all_gather_object(infos, info, group=group)
            ext.connect(infos)
        dist.barrier(group=group)
        _taken = True
        self.ext = ext
        self.max_bytes = min(int(os.environ.get("VLLM_ARX_MAX_BYTES", 256 << 10)), ext.max_bytes())
        # Two-shot (reduce-scatter + all-gather) for all-reduces above max_bytes,
        # up to this size: 2-2.7x faster than NCCL from 512 KB. Mesh only; 0 is off.
        self.two_shot_max = min(int(os.environ.get("VLLM_ARX_TWO_SHOT_MAX_KB", "0")) << 10, ext.max_bytes2())
        self.world = world
        self.disabled = False
        logger.info("arx all-reduce: rank %d/%d%s on %s, gids %s, up to %d bytes%s", rank, world,
                    " ring (prev, next)" if ring else "", devs, gids, self.max_bytes,
                    f", two-shot up to {self.two_shot_max} bytes" if self.two_shot_max else "")

    def should_use(self, t: torch.Tensor, two_shot: bool = False) -> bool:
        """Whether arx takes t: up to max_bytes, or up to two_shot_max when the caller allows it."""
        if not (t.dtype == torch.bfloat16 and t.is_cuda and t.is_contiguous() and t.numel() % 8 == 0):
            return False
        nbytes = t.numel() * 2
        return nbytes <= self.max_bytes or (
            two_shot and nbytes <= self.two_shot_max and t.numel() % (8 * self.world) == 0)

    def all_reduce(self, t: torch.Tensor) -> torch.Tensor:
        global _prefetch
        out = torch.empty_like(t)
        if t.numel() * 2 > self.max_bytes:
            self.ext.allreduce2(t, out)
            return out
        pf, _prefetch = _prefetch, ()
        self.ext.allreduce(t, out, [p for p, _ in pf], [b for _, b in pf])
        return out


_big_ext = None


class ArxBig:
    """Prefill-sized all-gathers over RoCE (arxbig.cu): ~10% faster than NCCL's.

    Enabled with VLLM_ARXBIG=1 next to arx. all_gather returns a view of a
    pinned buffer that the fourth all-gather after it overwrites, so a caller
    that keeps the result longer must copy it.
    """

    def __init__(self, group: ProcessGroup, device: torch.device):
        global _big_ext
        self.disabled = True
        self.rs = False
        self.gather = False
        rank, world = dist.get_rank(group), dist.get_world_size(group)
        fab = fabric(world, "arxbig")
        if fab is None:
            return
        devs, gids, ring = fab
        slot = int(os.environ.get("VLLM_ARXBIG_SLOT_MB", "132")) << 20
        slot -= slot % (world * 16)
        with torch.cuda.device(device):
            if _big_ext is None:
                _big_ext = _jit_load("arxbig", [os.path.join(os.path.dirname(__file__), "arxbig.cu")],
                                     extra_cuda_cflags=["-O3", "-std=c++17", "-gencode=arch=compute_121a,code=sm_121a"],
                                     extra_ldflags=["-libverbs"])
            self.rs = os.environ.get("VLLM_ARXBIG_RS") == "1"
            info = _big_ext.prepare(rank, world, devs, gids, ring, slot, slot if self.rs else 0)
            infos: list[bytes] = [b""] * world
            dist.all_gather_object(infos, info, group=group)
            _big_ext.connect(infos)
            dist.barrier(group=group)
        self.ext, self.world, self.slot = _big_ext, world, slot
        self.min_bytes = int(os.environ.get("VLLM_ARXBIG_MIN_KB", "1024")) << 10
        self.gather = os.environ.get("VLLM_ARXBIG_AG") == "1"
        self.disabled = False
        logger.info("arxbig all-gather: rank %d/%d%s, %d MiB slots", rank, world, " ring" if ring else "", slot >> 20)

    def should_gather(self, t: torch.Tensor) -> bool:
        total = t.numel() * t.element_size() * self.world
        return (t.is_cuda and t.dim() >= 1 and (t.numel() * t.element_size()) % 16 == 0
                and self.min_bytes <= total <= self.slot)

    def all_gather(self, t: torch.Tensor) -> torch.Tensor:
        return self.ext.all_gather(t.contiguous())