mukesh
1d8b9401e5
New intercube allreduce kernel replacing the old flat ring algorithms. Reduces across the 4x4 cube mesh within each SIP (pe0-only, same-lane), then inter-SIP exchange on root cube, then broadcast back. Supports ring_1d, torus_2d, and mesh_2d_no_wrap SIP topologies driven by topology.yaml. Integrated with dist.init_process_group / dist.all_reduce. New files: - src/kernbench/ccl/algorithms/intercube_allreduce.py (kernel) - src/kernbench/ccl/sfr_config.py (configure_sfr_intercube_multisip) - tests/test_allreduce_multidevice.py (config-driven, 3 topologies) - tests/test_distributed_intercube_allreduce.py (full distributed path) - tests/test_intercube_sfr_config.py (SFR wiring verification) Modified: - distributed.py: AhbmCCLBackend uses configure_sfr_intercube_multisip - topologies.py: added torus_2d, mesh_2d_no_wrap - install.py: global_E/W/N/S in _OPPOSITE_DIR - topology.yaml: added system.sips.topology - ccl.yaml: single intercube_allreduce algorithm - benches/ccl_allreduce.py: row_wise cube-mesh tensor layout Removed old flat-ring algorithms and their tests. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""CCL all-reduce bench (ADR-0024 + ADR-0027).
|
|
|
|
Pure TP launcher model: rank = SIP. Each rank owns a ``(N_CUBES, n_elem)``
|
|
tensor sharded row-wise across the cube mesh (pe0 per cube). After
|
|
``dist.all_reduce(op="sum")`` every cube on every rank must hold
|
|
``N_CUBES * sum(1..world_size)``. Rank 0 prints the pass/fail line.
|
|
|
|
Driven by ``ccl.yaml`` (``defaults.algorithm``, ``n_elem``) + ``topology.yaml``
|
|
(SIP count → world_size, cube_mesh → N_CUBES).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
|
|
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
|
|
from kernbench.policy.placement.dp import DPPolicy
|
|
|
|
DEFAULT_N_ELEM = 8
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _BenchCfg:
|
|
algorithm: str
|
|
n_elem: int
|
|
n_cubes: int
|
|
world_size: int
|
|
|
|
|
|
def _resolve_cfg(torch) -> _BenchCfg:
|
|
"""Read ccl.yaml + topology once at host side."""
|
|
merged = resolve_algorithm_config(load_ccl_config())
|
|
ws = torch.distributed.get_world_size()
|
|
spec = torch.spec or {}
|
|
n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1))
|
|
if ws != n_sips:
|
|
raise RuntimeError(
|
|
f"ccl_allreduce bench requires world_size == topology SIP count "
|
|
f"(world_size={ws}, n_sips={n_sips})."
|
|
)
|
|
cm = spec.get("sip", {}).get("cube_mesh", {})
|
|
n_cubes = int(cm.get("w", 4)) * int(cm.get("h", 4))
|
|
return _BenchCfg(
|
|
algorithm=merged["algorithm"],
|
|
n_elem=int(merged.get("n_elem", DEFAULT_N_ELEM)),
|
|
n_cubes=n_cubes,
|
|
world_size=ws,
|
|
)
|
|
|
|
|
|
def _rank_dp(n_cubes: int) -> DPPolicy:
|
|
return DPPolicy(cube="row_wise", pe="replicate", num_cubes=n_cubes, num_pes=1)
|
|
|
|
|
|
def _allocate_rank_tensor(torch, rank: int, cfg: _BenchCfg):
|
|
"""Allocate this rank's ``(n_cubes, n_elem)`` tensor on its SIP."""
|
|
return torch.zeros(
|
|
(cfg.n_cubes, cfg.n_elem), dtype="f16",
|
|
dp=_rank_dp(cfg.n_cubes), name=f"ccl_in_r{rank}",
|
|
)
|
|
|
|
|
|
def _init_with_rank_value(torch, tensor, rank: int, cfg: _BenchCfg) -> None:
|
|
"""Fill all cubes with the scalar ``rank + 1``."""
|
|
arr = np.full((cfg.n_cubes, cfg.n_elem), float(rank + 1), dtype=np.float16)
|
|
tensor.copy_(torch.from_numpy(arr))
|
|
|
|
|
|
def _report(result: np.ndarray, cfg: _BenchCfg) -> None:
|
|
"""Single-line pass/fail printer (rank 0 only)."""
|
|
expected = float(cfg.n_cubes * sum(range(1, cfg.world_size + 1)))
|
|
ok = True
|
|
for cube_id in range(cfg.n_cubes):
|
|
if not np.allclose(result[cube_id], expected, rtol=1e-1, atol=1e-1):
|
|
ok = False
|
|
break
|
|
if ok:
|
|
total = cfg.world_size * cfg.n_cubes
|
|
print(f" {cfg.algorithm} (ws={cfg.world_size}): {total} OK")
|
|
return
|
|
got = float(result.reshape(-1).mean())
|
|
print(
|
|
f" [FAIL] {cfg.algorithm} (ws={cfg.world_size}): "
|
|
f"got mean={got:.3f}, expected={expected:.3f}"
|
|
)
|
|
|
|
|
|
def _worker(rank: int, cfg: _BenchCfg, torch) -> None:
|
|
torch.ahbm.set_device(rank)
|
|
tensor = _allocate_rank_tensor(torch, rank, cfg)
|
|
_init_with_rank_value(torch, tensor, rank, cfg)
|
|
torch.distributed.all_reduce(tensor, op="sum")
|
|
if rank == 0:
|
|
_report(tensor.numpy(), cfg)
|
|
|
|
|
|
def run(torch) -> None:
|
|
torch.distributed.init_process_group(backend="ahbm")
|
|
cfg = _resolve_cfg(torch)
|
|
torch.multiprocessing.spawn(
|
|
_worker, args=(cfg, torch), nprocs=cfg.world_size,
|
|
)
|