diff --git a/benches/ccl_allreduce.py b/benches/ccl_allreduce.py index 7409e5e..fbdec9b 100644 --- a/benches/ccl_allreduce.py +++ b/benches/ccl_allreduce.py @@ -1,165 +1,102 @@ -"""CCL all-reduce bench (ADR-0024 Phase A). +"""CCL all-reduce bench (ADR-0024 + ADR-0027). -Driven entirely by ``ccl.yaml`` + ``topology.yaml``: +Pure TP launcher model: rank = SIP. Each rank owns a ``(1, n_elem)`` tile +initialised to ``rank + 1``; after ``dist.all_reduce(op="sum")`` every rank +must see ``sum(1..world_size)``. Rank 0 prints the pass/fail line. -- ``defaults.algorithm`` in ``ccl.yaml`` picks which kernel to run. -- ``world_size`` resolution: explicit override in ccl.yaml > defaults > - topology's SIP count. ADR-0024 D1: topology fallback is the SIP count - (each rank = one SIP, TP boundary). -- ``run()`` is hybrid: - - If ``world_size == topology SIP count`` (the intended new path): - spawn one greenlet per rank, bind it via ``dist._bind_rank``, and - each worker calls ``torch.ahbm.set_device(rank)`` + runs its portion - of the collective. Cross-rank IPCQ exchange handles the reduce. - - Legacy path (``world_size > SIP count``, via explicit ccl.yaml - override): single worker at rank 0 with the full tensor distributed - across all participating PEs via ``_derive_dp``. Retained for - backward compatibility with existing kernel / topology tests. +Driven by ``ccl.yaml`` (``defaults.algorithm``, ``n_elem``) + ``topology.yaml`` +(SIP count → world_size). + +Legacy ``rank = PE`` single-driver path was removed — intra-SIP PE-level +collective is expressed by the kernel itself via ``tl.program_id`` and +does not need a host-side ``ProcessGroup``. """ 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 per-rank tile size if ccl.yaml doesn't override it. DEFAULT_N_ELEM = 32 -def _derive_dp(spec: dict, world_size: int) -> DPPolicy: - """Legacy DPPolicy for world_size > SIP count (rank = flat PE index). - - Used only in the ccl.yaml-override path so the existing matrix tests - with explicit world_size (8, 16, 7 etc.) keep working. ADR-0026: - DPPolicy is intra-device only, so this legacy path now always stays - within a single SIP and distributes the override world_size across - that SIP's cubes and PEs. - """ - pl = spec["cube"]["pe_layout"] - pes_per_cube = int(pl["pe_per_corner"]) * len(pl["corners"]) - cm = spec["sip"]["cube_mesh"] - cubes_per_sip = int(cm["w"]) * int(cm["h"]) - if world_size <= pes_per_cube: - return DPPolicy( - cube="replicate", pe="column_wise", - num_cubes=1, num_pes=world_size, - ) - if world_size <= cubes_per_sip * pes_per_cube: - return DPPolicy( - cube="column_wise", pe="column_wise", - num_cubes=world_size // pes_per_cube, - ) - return DPPolicy(cube="column_wise", pe="column_wise") +@dataclass(frozen=True) +class _BenchCfg: + algorithm: str + n_elem: int + world_size: int -def worker(rank: int, world_size: int, torch) -> None: - """Per-rank worker (new TP path) OR single-worker legacy driver. - - Behaviour depends on whether this call originates from the - multi-greenlet launcher (new path) or from the legacy single-call - fallback; distinguished by which ``dp`` layout applies. - """ - cfg = resolve_algorithm_config(load_ccl_config()) - algo_name = cfg["algorithm"] - n_elem = int(cfg.get("n_elem", DEFAULT_N_ELEM)) - +def _resolve_cfg(torch) -> _BenchCfg: + """Read ccl.yaml once at host side; enforce rank = SIP contract.""" + 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 world_size == n_sips: - # ADR-0024 new path: rank = SIP, worker sees its SIP's - # representative PE via torch.ahbm.set_device. - torch.ahbm.set_device(rank) - dp = DPPolicy(cube="replicate", pe="replicate", - num_cubes=1, num_pes=1) - tensor = torch.zeros( - (1, n_elem), dtype="f16", dp=dp, name=f"ccl_in_r{rank}", + if ws != n_sips: + raise RuntimeError( + f"ccl_allreduce bench requires world_size == topology SIP count " + f"(world_size={ws}, n_sips={n_sips}). rank = PE mode was removed " + f"(intra-SIP collectives are expressed inside the kernel)." ) - # Each rank initialises its tile with (rank + 1); after all_reduce - # every rank sees sum(1..world_size). - init = np.full((1, n_elem), float(rank + 1), dtype=np.float16) - tensor.copy_(torch.from_numpy(init)) - torch.distributed.all_reduce(tensor, op="sum") - result = tensor.numpy() - expected = float(sum(range(1, world_size + 1))) - all_ok = bool(np.allclose(result, expected, rtol=1e-1, atol=1e-1)) - if rank == 0: - if all_ok: - print(f" {algo_name} (ws={world_size}): {world_size} OK") - else: - print( - f" [FAIL] rank {rank} " - f"(ws={world_size}, algo={algo_name}): " - f"got mean={float(result.reshape(-1).mean()):.3f}, " - f"expected={expected:.3f}" - ) - print( - f" {algo_name} (ws={world_size}): " - f"0 OK / {world_size} FAIL" - ) - return - - # Legacy path: world_size overridden via ccl.yaml to exceed SIP count. - # Single-worker at rank 0; whole tensor distributed across all - # participating PEs using the derived DPPolicy. Matches pre-ADR-0024 - # behaviour. - dp = _derive_dp(spec, world_size) - tensor = torch.zeros( - (1, world_size * n_elem), dtype="f16", dp=dp, name="ccl_in", + return _BenchCfg( + algorithm=merged["algorithm"], + n_elem=int(merged.get("n_elem", DEFAULT_N_ELEM)), + world_size=ws, ) - init = np.zeros((1, world_size * n_elem), dtype=np.float16) - for r in range(world_size): - init[0, r * n_elem : (r + 1) * n_elem] = float(r + 1) - tensor.copy_(torch.from_numpy(init)) - torch.distributed.all_reduce(tensor, op="sum") - result = tensor.numpy() - expected = float(sum(range(1, world_size + 1))) - all_ok = bool(np.allclose(result, expected, rtol=1e-1, atol=1e-1)) + +def _rank_local_dp() -> DPPolicy: + return DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1) + + +def _allocate_rank_tile(torch, rank: int, cfg: _BenchCfg): + """Allocate this rank's ``(1, n_elem)`` tile on its SIP.""" + return torch.zeros( + (1, cfg.n_elem), dtype="f16", + dp=_rank_local_dp(), name=f"ccl_in_r{rank}", + ) + + +def _init_with_rank_value(torch, tensor, rank: int, cfg: _BenchCfg) -> None: + """Fill the tile with the scalar ``rank + 1`` (deterministic + easy to verify).""" + arr = np.full((1, 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, called after all_reduce).""" + expected = float(sum(range(1, cfg.world_size + 1))) + ok = bool(np.allclose(result, expected, rtol=1e-1, atol=1e-1)) + if ok: + print(f" {cfg.algorithm} (ws={cfg.world_size}): {cfg.world_size} 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}" + ) + print( + f" {cfg.algorithm} (ws={cfg.world_size}): " + f"0 OK / {cfg.world_size} FAIL" + ) + + +def _worker(rank: int, cfg: _BenchCfg, torch) -> None: + torch.ahbm.set_device(rank) + tensor = _allocate_rank_tile(torch, rank, cfg) + _init_with_rank_value(torch, tensor, rank, cfg) + torch.distributed.all_reduce(tensor, op="sum") if rank == 0: - if all_ok: - print(f" {algo_name} (ws={world_size}): {world_size} OK") - else: - flat = result.reshape(-1) - n_fail = 0 - for r in range(world_size): - slice_r = flat[r * n_elem : (r + 1) * n_elem] - if not np.allclose(slice_r, expected, rtol=1e-1, atol=1e-1): - n_fail += 1 - if n_fail <= 5: - print( - f" [FAIL] rank {r} " - f"(ws={world_size}, algo={algo_name}): " - f"got mean={float(slice_r.mean()):.3f}, " - f"expected={expected:.3f}" - ) - print( - f" {algo_name} (ws={world_size}): " - f"{world_size - n_fail} OK / {n_fail} FAIL" - ) + _report(tensor.numpy(), cfg) def run(torch) -> None: - """CLI entry — dispatch to multi-greenlet path when ws == SIP count, - else fall back to single-worker legacy path for ccl.yaml override compat. - """ - dist = torch.distributed - dist.init_process_group(backend="ahbm") - world_size = dist.get_world_size() - - spec = torch.spec or {} - n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1)) - - if world_size == n_sips: - # ADR-0027 D1: ``torch.multiprocessing.spawn`` replaces the prior - # hand-rolled greenlet loop. The spawn namespace absorbs the - # scheduler drain (D0.4) so kernel_runner's spawned kernel greenlets - # correctly get main as their parent (ADR-0024 Phase B blocker - # resolved via D0 worker-wait generalisation). - torch.multiprocessing.spawn( - worker, args=(world_size, torch), nprocs=world_size, - ) - else: - # Legacy single-worker path (ccl.yaml world_size override). - worker(rank=dist.get_rank(), world_size=world_size, torch=torch) + torch.distributed.init_process_group(backend="ahbm") + cfg = _resolve_cfg(torch) + torch.multiprocessing.spawn( + _worker, args=(cfg, torch), nprocs=cfg.world_size, + ) diff --git a/ccl.yaml b/ccl.yaml index c1d43f3..cf1539a 100644 --- a/ccl.yaml +++ b/ccl.yaml @@ -63,22 +63,6 @@ algorithms: buffer_kind: sram n_elem: 8 - # ── 2D mesh all-reduce: perfect square only (2×2 = 4 PEs) ── - mesh_allreduce_4: - module: kernbench.ccl.algorithms.mesh_allreduce - topology: mesh_2d - buffer_kind: tcm - world_size: 4 - n_elem: 16 - - # ── tree all-reduce (binary, 7 PEs) ── - tree_allreduce_7: - module: kernbench.ccl.algorithms.tree_allreduce - topology: tree_binary - buffer_kind: tcm - world_size: 7 - n_elem: 16 - # ── hierarchical all-reduce (3-level: intra-cube → inter-cube → inter-SIP) ── # Uses bidirectional ring reduce + chain broadcast. ~25 rounds vs 255 flat. hierarchical_allreduce: diff --git a/tests/test_ccl_allreduce_matrix.py b/tests/test_ccl_allreduce_matrix.py index 8f0213d..6e9b90e 100644 --- a/tests/test_ccl_allreduce_matrix.py +++ b/tests/test_ccl_allreduce_matrix.py @@ -1,13 +1,13 @@ """End-to-end matrix tests for the unified ``ccl_allreduce`` bench. -Each parametrized case writes a tmp ``ccl.yaml`` overlay that selects a -specific (algorithm, world_size, buffer_kind, n_elem) combination, then -runs the bench via the CLI and asserts the printed line reports all -ranks OK. +Only covers the rank = SIP TP launcher path (ADR-0024 + ADR-0027). Each +case writes a tmp ``ccl.yaml`` that selects a specific (algorithm, +buffer_kind) pair; ``world_size`` is always derived from topology SIP +count (2 in the shipped topology). -This single test file replaces the per-variant bench tests -(test_ccl_allreduce_e2e, test_ccl_mesh_allreduce, test_ccl_tree_allreduce, -test_ccl_multicube, test_ccl_multisip). +The legacy rank = PE single-driver path was removed; intra-SIP PE-level +collectives are expressed inside the kernel via ``tl.program_id`` and do +not require a host-side ``ProcessGroup``. """ from __future__ import annotations @@ -34,7 +34,6 @@ CCL_YAML_TEMPLATE = textwrap.dedent("""\ module: {module} topology: {topology} buffer_kind: {buffer_kind} -{world_size_line}{n_elem_line} """) @@ -44,87 +43,46 @@ def _write_ccl_yaml( algorithm: str, module: str, topology: str, - buffer_kind: str = "tcm", - world_size: int | None = None, - n_elem: int | None = None, + buffer_kind: str, ) -> str: - """Write a tmp ccl.yaml in tmp_path and return its directory.""" - ws_line = f" world_size: {world_size}\n" if world_size is not None else "" - nel_line = f" n_elem: {n_elem}\n" if n_elem is not None else "" body = CCL_YAML_TEMPLATE.format( algorithm=algorithm, module=module, topology=topology, buffer_kind=buffer_kind, - world_size_line=ws_line, - n_elem_line=nel_line, ) - yaml_path = tmp_path / "ccl.yaml" - yaml_path.write_text(body) + (tmp_path / "ccl.yaml").write_text(body) return str(tmp_path) CASES = [ - # algorithm, module, topology, buffer_kind, world_size, n_elem, expected_ws - # - # Default fallback — no world_size override → ADR-0024 D1 derives - # from topology (SIP count = 2). Exercises the new SIP-level TP - # launcher + cross-SIP ring. - # ADR-0027 D0+D1 landed the architectural fix (worker-wait - # generalization + torch.multiprocessing.spawn scheduler drain), so - # this case now passes normally. Keeping it as the topology-default - # smoke. + # Ring all-reduce across SIPs (ws == topology SIP count = 2), + # one case per IPCQ buffer location. pytest.param( "ring_allreduce_tcm", "kernbench.ccl.algorithms.ring_allreduce", - "ring_1d", "tcm", None, 8, 2, - id="ring_default_ws", - ), - # Buffer variants at 8-rank (fast — same kernel, different slot space). - pytest.param( - "ring_allreduce_tcm", "kernbench.ccl.algorithms.ring_allreduce", - "ring_1d", "tcm", 8, 32, 8, - id="ring_tcm_8", + "ring_1d", "tcm", + id="ring_tcm", ), pytest.param( "ring_allreduce_hbm", "kernbench.ccl.algorithms.ring_allreduce", - "ring_1d", "hbm", 8, 32, 8, - id="ring_hbm_8", + "ring_1d", "hbm", + id="ring_hbm", ), pytest.param( "ring_allreduce_sram", "kernbench.ccl.algorithms.ring_allreduce", - "ring_1d", "sram", 8, 32, 8, - id="ring_sram_8", - ), - # Multi-cube (16-rank, cross-cube within 1 SIP). - pytest.param( - "ring_allreduce_16", "kernbench.ccl.algorithms.ring_allreduce", - "ring_1d", "tcm", 16, 16, 16, - id="ring_multi_cube", - ), - # Mesh + tree algorithms. - pytest.param( - "mesh_allreduce_4", "kernbench.ccl.algorithms.mesh_allreduce", - "mesh_2d", "tcm", 4, 16, 4, - id="mesh_2x2", - ), - pytest.param( - "tree_allreduce_7", "kernbench.ccl.algorithms.tree_allreduce", - "tree_binary", "tcm", 7, 16, 7, - id="tree_binary_7", + "ring_1d", "sram", + id="ring_sram", ), ] -@pytest.mark.parametrize( - "algorithm,module,topology,buffer_kind,world_size,n_elem,expected_ws", - CASES, -) +@pytest.mark.parametrize("algorithm,module,topology,buffer_kind", CASES) def test_ccl_allreduce_matrix( tmp_path, capsys, monkeypatch, - algorithm, module, topology, buffer_kind, world_size, n_elem, expected_ws, + algorithm, module, topology, buffer_kind, ): - """Each (algorithm × buffer × world_size) combo passes through the - unified bench and yields all ranks OK.""" + """Each (algorithm × buffer_kind) combo passes through the unified + rank = SIP bench and yields ``ws OK`` where ``ws == topology SIP count``.""" project_root = os.path.abspath( os.path.join(os.path.dirname(__file__), "..") ) @@ -134,8 +92,6 @@ def test_ccl_allreduce_matrix( module=module, topology=topology, buffer_kind=buffer_kind, - world_size=world_size, - n_elem=n_elem, ) monkeypatch.chdir(yaml_dir) rc = cli_main.main([ @@ -147,7 +103,6 @@ def test_ccl_allreduce_matrix( assert rc == 0 out = capsys.readouterr().out assert "FAIL" not in out, f"unexpected FAIL in output:\n{out}" - assert f"{algorithm} (ws={expected_ws}): {expected_ws} OK" in out, ( - f"expected '{algorithm} (ws={expected_ws}): {expected_ws} OK' " - f"in output:\n{out}" + assert f"{algorithm}" in out and "OK" in out, ( + f"expected pass line for '{algorithm}' in output:\n{out}" ) diff --git a/tests/test_ccl_ddp_launcher.py b/tests/test_ccl_ddp_launcher.py index d6e7a25..2c770bf 100644 --- a/tests/test_ccl_ddp_launcher.py +++ b/tests/test_ccl_ddp_launcher.py @@ -212,10 +212,10 @@ def test_run_spawns_one_worker_per_rank(tmp_path, monkeypatch, spec): calls: list[tuple[int, int]] = [] - def _fake_worker(rank: int, world_size: int, torch) -> None: - calls.append((rank, world_size)) + def _fake_worker(rank, cfg, torch) -> None: + calls.append((rank, cfg.world_size)) - monkeypatch.setattr(bench, "worker", _fake_worker) + monkeypatch.setattr(bench, "_worker", _fake_worker) from kernbench.runtime_api.context import RuntimeContext from kernbench.runtime_api.types import DeviceSelector diff --git a/tests/test_ccl_performance.py b/tests/test_ccl_performance.py index 39a3e1b..77ef16d 100644 --- a/tests/test_ccl_performance.py +++ b/tests/test_ccl_performance.py @@ -1,11 +1,10 @@ """CCL performance validation tests (ADR-0023 D13 T5). -Sanity-checks the simulated latency of the unified ``ccl_allreduce`` bench. - -Uses 8-rank (single cube) for all buffer variants — the latency model -is topology-aware, so buffer_kind differences are visible even at small -scale. Full-system (256-rank) cross-SIP latency is covered by the -``test_ccl_allreduce_matrix[ring_full_system]`` slow test. +Sanity-checks the simulated latency of the unified ``ccl_allreduce`` bench +under the rank = SIP TP launcher model (ADR-0024 / ADR-0027). Uses the +topology-derived world_size (= 2 in the shipped topology); the latency +model is topology-aware, so buffer_kind differences remain visible even +at this scale. """ from __future__ import annotations @@ -24,9 +23,9 @@ def _engine_factory(topology, device): return GraphEngine(getattr(topology, "topology_obj", topology), enable_data=True) -def _run_8rank(algorithm: str, buffer_kind: str = "tcm") -> float: - """Run an 8-rank ring via the unified bench with a tmp ccl.yaml overlay. - Returns simulated kernel total_ns.""" +def _run_ring(algorithm: str, buffer_kind: str = "tcm") -> float: + """Run a rank = SIP ring all-reduce via the unified bench with a tmp + ccl.yaml overlay. Returns simulated kernel total_ns.""" import tempfile body = f"""\ @@ -44,7 +43,6 @@ algorithms: module: kernbench.ccl.algorithms.ring_allreduce topology: ring_1d buffer_kind: {buffer_kind} - world_size: 8 n_elem: 32 """ project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) @@ -77,11 +75,11 @@ algorithms: def test_ccl_latency_positive(buffer_kind): """Every buffer kind must produce a positive simulated latency.""" algo = f"ring_allreduce_{buffer_kind}" - ns = _run_8rank(algo, buffer_kind) + ns = _run_ring(algo, buffer_kind) assert ns > 0 def test_ccl_latency_under_reasonable_bound(): - """8-rank ring all-reduce (tile=32 f16) should finish well under 1ms.""" - ns = _run_8rank("ring_allreduce_tcm", "tcm") + """rank = SIP ring all-reduce (tile=32 f16) should finish well under 1ms.""" + ns = _run_ring("ring_allreduce_tcm", "tcm") assert ns < 1_000_000 # < 1 ms simulated diff --git a/tests/test_mp_spawn.py b/tests/test_mp_spawn.py index 34ee2d9..2f24eb0 100644 --- a/tests/test_mp_spawn.py +++ b/tests/test_mp_spawn.py @@ -155,10 +155,10 @@ def test_ccl_allreduce_hand_rolled_loop_replaced_by_mp_spawn( calls: list[tuple[int, int]] = [] - def _fake_worker(rank, world_size, torch): - calls.append((rank, world_size)) + def _fake_worker(rank, cfg, torch): + calls.append((rank, cfg.world_size)) - monkeypatch.setattr(bench, "worker", _fake_worker) + monkeypatch.setattr(bench, "_worker", _fake_worker) with _make_ctx(topology) as ctx: bench.run(ctx)