mukesh
b610cb0d9a
Convert the multidevice allreduce correctness + latency/buffer-kind sweeps to run through the real PyTorch-distributed path (init_process_group(backend="ahbm") -> mp.spawn -> dist.all_reduce) instead of direct ctx.launch, and reorganize the CCL/allreduce tests into a tests/sccl/ package split one test per file. Production change (required for the distributed path on non-square SIP grids): - AhbmCCLBackend now reads explicit system.sips.w/h from the spec, with a square-only sqrt fallback that raises on ambiguity, instead of silently guessing round(sqrt(count)). This fixes the 2x3 / 3x2 torus + mesh cases, which previously resolved to a wrong 2x2 grid. Mirrors the test helper's _sip_topo_dims precedence (explicit w/h > square fallback > raise). Test reorganization (tests/sccl/): - _allreduce_helpers.py: shared plumbing (distributed driver, config writers, direct-launch run_allreduce parity reference, sweep/buffer-kind constants, plot aggregators, topology-diagram + FSIM-comparison emitters). - test_allreduce_ring_torus_mesh.py: correctness across ring/torus/mesh. - test_distributed_default_topology.py: full distributed path on topology.yaml. - test_plot_latency_sweep.py / test_plot_buffer_kind_sweep.py: sweep rows. - test_plot_topology_diagram.py / test_plot_comparison_fsim.py: plot emitters. - test_intercube_root_center.py: moved in (ADR-0032 center-root latency guard). Also: - Move the FSIM comparison plot generator out of scripts/ into the sccl suite. - Delete superseded test files (test_allreduce_multidevice, test_distributed_lrab_hierarchical_allreduce, test_allreduce_buffer_kind_sweep) and repoint conftest aggregators + the ipcq buffer-kind importers. - Regenerate the allreduce_latency_plots derived artifacts from the full sweep. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
"""Full distributed path against topology.yaml as-is (no overrides).
|
|
|
|
The same flow a real DDP training script would use:
|
|
init_process_group(backend="ahbm") → mp.spawn → dist.all_reduce.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from kernbench.runtime_api.context import RuntimeContext
|
|
from kernbench.runtime_api.types import DeviceSelector
|
|
from kernbench.sim_engine.engine import GraphEngine
|
|
from kernbench.topology.builder import resolve_topology
|
|
|
|
from tests.sccl._allreduce_helpers import (
|
|
DEFAULT_N_ELEM,
|
|
TOPOLOGY_PATH,
|
|
_worker,
|
|
_write_ccl_yaml,
|
|
)
|
|
|
|
|
|
def test_distributed_lrab_hierarchical_allreduce(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(_write_ccl_yaml(tmp_path))
|
|
|
|
topo = resolve_topology(str(TOPOLOGY_PATH))
|
|
engine = GraphEngine(topo.topology_obj, enable_data=True)
|
|
spec = topo.topology_obj.spec
|
|
n_sips = int(spec["system"]["sips"]["count"])
|
|
cm = spec["sip"]["cube_mesh"]
|
|
n_cubes = int(cm["w"]) * int(cm["h"])
|
|
|
|
with RuntimeContext(
|
|
engine=engine,
|
|
target_device=DeviceSelector("all"),
|
|
correlation_id="dist_intercube_ar",
|
|
spec=spec,
|
|
) as ctx:
|
|
ctx.distributed.init_process_group(backend="ahbm")
|
|
assert ctx.distributed.get_world_size() == n_sips
|
|
|
|
t_start = engine._env.now
|
|
ctx.multiprocessing.spawn(
|
|
_worker, args=(n_cubes, DEFAULT_N_ELEM, n_sips, ctx),
|
|
nprocs=n_sips,
|
|
)
|
|
t_end = engine._env.now
|
|
print(f"\n[distributed] sim latency = "
|
|
f"{t_end - t_start:.1f} ns ({(t_end - t_start) / 1000:.3f} us)")
|