265 lines
8.1 KiB
Python
265 lines
8.1 KiB
Python
"""
|
|
Correctness test for the JIT custom all-reduce (v2) kernel.
|
|
|
|
The test compares the JIT custom all-reduce output against NCCL all-reduce
|
|
for various tensor sizes and dtypes, in both eager and CUDA-graph modes.
|
|
|
|
Usage:
|
|
python -m pytest test_jit_custom_all_reduce.py -v
|
|
|
|
This file doubles as the torchrun worker script. The test class launches
|
|
torchrun --nproc_per_node=N <this_file>
|
|
and asserts that all worker processes exit successfully.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import itertools
|
|
import logging
|
|
import multiprocessing as mp
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
import sglang.srt.distributed.parallel_state as ps
|
|
from sglang.jit_kernel.all_reduce import (
|
|
AllReduceAlgo,
|
|
_jit_custom_all_reduce_pull_module,
|
|
_jit_custom_all_reduce_push_module,
|
|
)
|
|
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
|
CustomAllReduceV2,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
|
|
register_cuda_ci(
|
|
est_time=300,
|
|
suite="stage-b-kernel-unit-8-gpu-h200",
|
|
)
|
|
register_cuda_ci(
|
|
est_time=300,
|
|
suite="nightly-kernel-8-gpu-h200",
|
|
nightly=True,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test parameters (shared between test class and worker)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
TEST_SIZES = [
|
|
16,
|
|
32,
|
|
512,
|
|
1024,
|
|
1024 + 16, # weird case
|
|
4 * 1024,
|
|
32 * 1024,
|
|
256 * 1024,
|
|
2 * 1024 * 1024, # 2M elements
|
|
4 * 1024 * 1024, # 4M elements
|
|
]
|
|
TEST_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
|
SHOTS = [
|
|
AllReduceAlgo.ONE_SHOT_PULL,
|
|
AllReduceAlgo.ONE_SHOT_PUSH,
|
|
AllReduceAlgo.TWO_SHOT_PULL,
|
|
]
|
|
USE_GRAPH_OPTIONS = [True, False]
|
|
TEST_CONFIG = itertools.product(TEST_SIZES, TEST_DTYPES, SHOTS, USE_GRAPH_OPTIONS)
|
|
TEST_LAYERS = 4
|
|
TEST_LOOP = 16
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test class (runs via pytest, launches torchrun subprocesses)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _run_torchrun(nproc: int, timeout: int = 300) -> None:
|
|
"""Launch this script as a torchrun worker and assert success."""
|
|
cmd = [
|
|
"torchrun",
|
|
f"--nproc_per_node={nproc}",
|
|
__file__,
|
|
]
|
|
result = subprocess.run(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
timeout=timeout,
|
|
)
|
|
assert result.returncode == 0, (
|
|
f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
|
|
f"{result.stdout}"
|
|
)
|
|
|
|
|
|
def _compile_one(dtype: torch.dtype, world_size: int):
|
|
_jit_custom_all_reduce_push_module(dtype, world_size)
|
|
_jit_custom_all_reduce_pull_module(dtype, world_size)
|
|
|
|
|
|
def _precompile_kernels() -> None:
|
|
# NOTE: even when device count < 8, we should be able to compile all
|
|
process_map: Dict[Tuple[torch.dtype, int], mp.Process] = {}
|
|
COMPILE_SPACE = itertools.product(TEST_DTYPES, [2, 3, 4, 5, 6, 7, 8])
|
|
mp.set_start_method("spawn")
|
|
for config in COMPILE_SPACE:
|
|
process_map[config] = mp.Process(target=_compile_one, args=config)
|
|
for process in process_map.values():
|
|
process.start()
|
|
for (dtype, world_size), process in process_map.items():
|
|
process.join()
|
|
if process.exitcode != 0:
|
|
raise RuntimeError(f"Custom All Reduce {world_size=} {dtype=} failed")
|
|
|
|
|
|
@pytest.mark.parametrize("nproc", [1, 2, 3, 4, 5, 6, 7, 8])
|
|
def test_custom_allreduce(nproc: int) -> None:
|
|
if nproc == 1: # NOTE: special case to speed up tests
|
|
return _precompile_kernels()
|
|
|
|
device_count = torch.cuda.device_count()
|
|
if device_count < nproc:
|
|
pytest.skip(
|
|
f"Requires at least {nproc} GPUs, but only {device_count} available"
|
|
)
|
|
_run_torchrun(nproc)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Worker logic (executed by each torchrun process)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def init_distributed():
|
|
"""Initialize distributed groups via torchrun env vars.
|
|
|
|
Returns (rank, device, cpu_group, nccl_group, comm).
|
|
"""
|
|
local_rank = int(os.environ["LOCAL_RANK"])
|
|
world_size = int(os.environ["WORLD_SIZE"])
|
|
rank = local_rank
|
|
device = torch.device(f"cuda:{rank}")
|
|
torch.cuda.set_device(device)
|
|
|
|
dist.init_process_group(backend="gloo")
|
|
ps._WORLD = coord = ps.init_world_group(
|
|
ranks=list(range(world_size)),
|
|
local_rank=local_rank,
|
|
backend="nccl",
|
|
)
|
|
|
|
cpu_group = coord.cpu_group
|
|
nccl_group = coord.device_group
|
|
assert nccl_group is not None
|
|
|
|
max_size = max(TEST_SIZES) * 4
|
|
comm = CustomAllReduceV2(cpu_group, device, max_size, max_size)
|
|
if comm.disabled:
|
|
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
|
|
|
return rank, device, cpu_group, nccl_group, comm
|
|
|
|
|
|
@torch.inference_mode()
|
|
def worker_test(
|
|
device: torch.device,
|
|
nccl_group: dist.ProcessGroup,
|
|
comm: CustomAllReduceV2,
|
|
size: int,
|
|
dtype: torch.dtype,
|
|
use_graph: bool,
|
|
algo: AllReduceAlgo,
|
|
) -> Optional[RuntimeError]:
|
|
comm.override_algo = algo
|
|
|
|
def get_run_graph_fn():
|
|
graph = torch.cuda.CUDAGraph()
|
|
graph_inp = torch.zeros((TEST_LAYERS, size), dtype=dtype, device=device)
|
|
out_jits = []
|
|
with comm.capture():
|
|
with torch.cuda.graph(graph):
|
|
for i in range(TEST_LAYERS):
|
|
out_jits.append(comm.custom_all_reduce(graph_inp[i]))
|
|
out_jit = torch.stack(out_jits)
|
|
torch.cuda.synchronize()
|
|
|
|
def run_graph(x: torch.Tensor) -> torch.Tensor:
|
|
graph_inp.copy_(x)
|
|
graph.replay()
|
|
return out_jit.clone()
|
|
|
|
return run_graph
|
|
|
|
def get_run_eager_fn():
|
|
def run_eager(x: torch.Tensor) -> torch.Tensor:
|
|
eager_inp = x.clone()
|
|
out_eagers = []
|
|
for i in range(TEST_LAYERS):
|
|
out_eagers.append(comm.custom_all_reduce(eager_inp[i]))
|
|
torch.cuda.synchronize()
|
|
return torch.stack(out_eagers)
|
|
|
|
return run_eager
|
|
|
|
run_fn = get_run_graph_fn() if use_graph else get_run_eager_fn()
|
|
num_errors = 0
|
|
for _ in range(TEST_LOOP):
|
|
# NOTE: 15 * 8 < 128, which is the precision limit for bf16
|
|
inp = torch.randint(0, 16, (TEST_LAYERS, size), dtype=dtype, device=device)
|
|
assert comm.should_custom_ar(inp[0])
|
|
out_ref = inp.clone()
|
|
dist.all_reduce(out_ref, group=nccl_group)
|
|
out_jit = run_fn(inp)
|
|
num_errors += not torch.all(out_jit == out_ref)
|
|
if num_errors > 0:
|
|
return RuntimeError(
|
|
f"Test failed for {size=}, {dtype=}, {algo=}, "
|
|
f"{use_graph=} with {num_errors} errors. "
|
|
)
|
|
return None
|
|
|
|
|
|
def worker_main() -> None:
|
|
"""Entry point for each torchrun worker process."""
|
|
rank, device, cpu_group, nccl_group, comm = init_distributed()
|
|
world_size = dist.get_world_size()
|
|
|
|
torch.cuda.set_stream(torch.cuda.Stream())
|
|
|
|
logging.disable(logging.INFO) # Suppress internal logging for cleaner test output
|
|
items = list(enumerate(TEST_CONFIG))
|
|
for i, (size, dtype, algo, use_graph) in items:
|
|
error = worker_test(device, nccl_group, comm, size, dtype, use_graph, algo)
|
|
if error is not None:
|
|
print(
|
|
f"Worker {rank} failed for {size=}, {dtype=}, "
|
|
f"{algo=}, {use_graph=}, iteration={i}\n"
|
|
f"Error: {error}"
|
|
)
|
|
# communicate the result to rank 0 for logging
|
|
result = torch.tensor([int(error is not None)])
|
|
dist.all_reduce(result, group=cpu_group)
|
|
failed = bool(result.item())
|
|
if failed:
|
|
raise RuntimeError(
|
|
f"Test failed on rank {rank} for config: "
|
|
f"{size=}, {dtype=}, {algo=}, {use_graph=}"
|
|
)
|
|
|
|
comm.close()
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if "LOCAL_RANK" in os.environ:
|
|
worker_main()
|
|
else:
|
|
sys.exit(pytest.main([__file__, "-x", "-vv", "-s"]))
|