ywkang
105f1dc09e
Implements ADR-0027 Phase 2 end-to-end. All 559 tests pass (was 523 + 1 xfail; ring_default_ws strict-xfail is now resolved). D0 — Worker-wait generalization (context.py): - _pending_worker_waits queue on RuntimeContext. - ctx.wait(h) in worker context defers to main via g.parent.switch(). Fast-path for already-completed handles. - Worker API is unchanged: tensor deploy, launch, etc. still look synchronous; they're transparently cooperatively scheduled. - Solves ADR-0024 Phase B kernel-greenlet orphan bug (env.run now only ever drives from main; kernel _parent is always main). D0.5 — Host-read barrier (tensor.py): - Explicit _HOST_READ_BARRIERS registry (T5.g closed-set via code review, not reflection-magic). - numpy/data/__getitem__/__repr__ drain pending worker-waits before host-observable read. - copy_: source-side barrier via source.numpy(). Target-side write barrier is intentionally NOT applied — global pending target barrier prematurely drains cross-rank collectives → deadlock. - Collective pending is excluded from barrier drain condition (collective is cross-rank; its own yield in all_reduce covers the invariant naturally). D1 — torch.multiprocessing.spawn (runtime_api/multiprocessing.py): - API signature parity with real PyTorch spawn; execution is cooperative greenlet scheduler (process isolation etc. are explicit non-goals per D1.0). - _drain_pending drains worker-waits then collectives in one barrier, loop-until-empty. - Round-based exception handling with SystemExit sibling abort + SpawnException(errors) wrapping root-cause ranks. - RuntimeContext attaches ctx.multiprocessing in __post_init__. - benches/ccl_allreduce.py hand-rolled loop collapses to one torch.multiprocessing.spawn call. D2–D6 — kernbench.tp package: - parallel_state: initialize_model_parallel, get_*_rank, get_*_world_size, with weak active-ctx registry in context.py. - layers: ColumnParallelLinear, RowParallelLinear (shape-only primitives — fp16 gemm via tl.load + tl.dot + tl.store). - kernels: _gemm_kernel used by TP layers (self-contained; no bench dependency). - primitives / mappings stubs per D6/D8. Data-path fixes (surfaced by TP gemm + all_reduce sequence): - sim_engine/op_log.py: dma_write snapshot is skipped for TCM sources (PE scratch is repopulated by Phase 2 math/gemm replay — capturing Phase-1-time snapshot picked up STALE data from prior kernel's output aliased at the same scratch addr, causing the later kernel's dma_write to overwrite Phase 2 result with stale value). - sim_engine/op_log.py + sim_engine/data_executor.py: per-operand space recorded on GemmCmd and composite gemm records so HBM-resident operands (tl.load output) don't default to TCM during replay. - runtime_api/context.py: ctx.zeros writes zero-init to MemoryStore at VA keys so kernels reading via VA see deterministic init even without explicit copy_(). Tests (Phase 1 + Phase 2): - test_worker_wait_drain (T3): orphan invariant + resume + multi-rank drain + idempotency + exception propagation. - test_mp_spawn (T4): spawn shape + bind + SpawnException scope. - test_host_read_barrier (T5): barrier contract per entry-point + closed-set registry check. - test_tp_parallel_state (T1): initialize + rank lookup. - test_tp_layers (T2): shape + deterministic numerical correctness (concat-matmul equality for RowParallel, not mean-only). - test_tp_mlp (T6): full 2-layer MLP with deterministic weight numerical match + rank-consistency post all-reduce. - test_ccl_allreduce_matrix: ring_default_ws xfail removed (T7). Regression: 523 pre + 35 new + 1 ex-xfail = 559 passed, 1 intentional skip (T3.e historical failure documentation). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
235 lines
8.4 KiB
Python
235 lines
8.4 KiB
Python
"""ADR-0027 T2: TP layer shape + numerical correctness (D4/D5).
|
|
|
|
Phase 1: ``kernbench.tp.layers`` doesn't exist → import failure. Phase 2
|
|
lands D4/D5 and T2 passes with deterministic non-zero weight patterns.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
|
|
def _make_ctx(topology):
|
|
from kernbench.runtime_api.context import RuntimeContext
|
|
from kernbench.runtime_api.types import DeviceSelector
|
|
from kernbench.sim_engine.engine import GraphEngine
|
|
|
|
engine = GraphEngine(topology.topology_obj, enable_data=True)
|
|
return RuntimeContext(
|
|
engine=engine,
|
|
target_device=DeviceSelector("all"),
|
|
correlation_id="test_t2",
|
|
spec=topology.topology_obj.spec,
|
|
)
|
|
|
|
|
|
# ── Shape / structural ───────────────────────────────────────────────
|
|
|
|
|
|
def test_column_parallel_weight_shape_per_rank(topology):
|
|
"""ColumnParallelLinear weight per rank is (in_features, out // ws)."""
|
|
import kernbench.tp as tp
|
|
from kernbench.runtime_api.tensor import Tensor
|
|
|
|
with _make_ctx(topology) as ctx:
|
|
ctx.distributed.init_process_group(backend="ahbm")
|
|
ws = ctx.distributed.get_world_size()
|
|
tp.initialize_model_parallel(ws)
|
|
|
|
def _worker(rank: int):
|
|
ctx.ahbm.set_device(rank)
|
|
fc = tp.ColumnParallelLinear(
|
|
in_features=256, out_features=512, torch=ctx,
|
|
)
|
|
assert fc.weight.shape == (256, 512 // ws)
|
|
|
|
ctx.multiprocessing.spawn(_worker, args=(), nprocs=ws)
|
|
|
|
|
|
def test_row_parallel_weight_shape_per_rank(topology):
|
|
"""RowParallelLinear weight per rank is (in_features // ws, out_features)."""
|
|
import kernbench.tp as tp
|
|
|
|
with _make_ctx(topology) as ctx:
|
|
ctx.distributed.init_process_group(backend="ahbm")
|
|
ws = ctx.distributed.get_world_size()
|
|
tp.initialize_model_parallel(ws)
|
|
|
|
def _worker(rank: int):
|
|
ctx.ahbm.set_device(rank)
|
|
fc = tp.RowParallelLinear(
|
|
in_features=512, out_features=256, torch=ctx,
|
|
)
|
|
assert fc.weight.shape == (512 // ws, 256)
|
|
|
|
ctx.multiprocessing.spawn(_worker, args=(), nprocs=ws)
|
|
|
|
|
|
# ── T2.a: ColumnParallel deterministic numerical ─────────────────────
|
|
|
|
|
|
def test_column_parallel_forward_matches_matmul(topology):
|
|
"""T2.a: ColumnParallelLinear.forward output == x @ W_rank (rtol 1e-2)."""
|
|
import kernbench.tp as tp
|
|
from kernbench.runtime_api.tensor import Tensor
|
|
|
|
with _make_ctx(topology) as ctx:
|
|
ctx.distributed.init_process_group(backend="ahbm")
|
|
ws = ctx.distributed.get_world_size()
|
|
tp.initialize_model_parallel(ws)
|
|
|
|
M = 4
|
|
D_in, D_out = 32, 32 * ws
|
|
|
|
def _worker(rank: int):
|
|
ctx.ahbm.set_device(rank)
|
|
fc = tp.ColumnParallelLinear(
|
|
in_features=D_in, out_features=D_out, torch=ctx,
|
|
)
|
|
# Deterministic non-zero weight: rank-scaled constant.
|
|
k_local = D_out // ws
|
|
weight_np = np.full(
|
|
(D_in, k_local), 0.01 * (rank + 1), dtype=np.float16,
|
|
)
|
|
src = Tensor(shape=(D_in, k_local), dtype="f16", name="host_w")
|
|
src._host_buffer = weight_np
|
|
fc.weight.copy_(src)
|
|
|
|
# Input: full-replicated constant.
|
|
x_np = np.full((M, D_in), 0.5, dtype=np.float16)
|
|
x = ctx.zeros(
|
|
(M, D_in), dtype="f16",
|
|
dp=_replicate_dp(), name=f"t2a_x_r{rank}",
|
|
)
|
|
hx = Tensor(shape=x_np.shape, dtype="f16", name="host_x")
|
|
hx._host_buffer = x_np
|
|
x.copy_(hx)
|
|
|
|
y = fc.forward(x)
|
|
out = y.numpy()
|
|
|
|
expected = x_np.astype(np.float32) @ weight_np.astype(np.float32)
|
|
assert out.shape == (M, k_local)
|
|
assert np.allclose(out.astype(np.float32), expected,
|
|
rtol=1e-2, atol=1e-2), (
|
|
f"rank {rank}: output does not match x @ W_local"
|
|
)
|
|
|
|
ctx.multiprocessing.spawn(_worker, args=(), nprocs=ws)
|
|
|
|
|
|
# ── T2.b: RowParallel observable equality ────────────────────────────
|
|
|
|
|
|
def test_row_parallel_forward_concat_matmul_equality(topology):
|
|
"""T2.b (primary): RowParallel output == concat(x) @ concat(W) (all-reduced)."""
|
|
import kernbench.tp as tp
|
|
from kernbench.runtime_api.tensor import Tensor
|
|
|
|
with _make_ctx(topology) as ctx:
|
|
ctx.distributed.init_process_group(backend="ahbm")
|
|
ws = ctx.distributed.get_world_size()
|
|
tp.initialize_model_parallel(ws)
|
|
|
|
M = 4
|
|
D_in, D_out = 32 * ws, 32 # must divide ws evenly
|
|
results: dict[int, np.ndarray] = {}
|
|
|
|
def _worker(rank: int):
|
|
ctx.ahbm.set_device(rank)
|
|
fc = tp.RowParallelLinear(
|
|
in_features=D_in, out_features=D_out, torch=ctx,
|
|
)
|
|
# Per-rank W_k = constant 0.01 * (rank + 1)
|
|
n_local = D_in // ws
|
|
weight_np = np.full(
|
|
(n_local, D_out), 0.01 * (rank + 1), dtype=np.float16,
|
|
)
|
|
src = Tensor(shape=weight_np.shape, dtype="f16", name="host_w")
|
|
src._host_buffer = weight_np
|
|
fc.weight.copy_(src)
|
|
|
|
# Input x_k = constant 0.1 * (rank + 1) (pretending it was
|
|
# column-sharded from upstream).
|
|
x_np = np.full((M, n_local), 0.1 * (rank + 1), dtype=np.float16)
|
|
x = ctx.zeros(
|
|
(M, n_local), dtype="f16",
|
|
dp=_replicate_dp(), name=f"t2b_x_r{rank}",
|
|
)
|
|
hx = Tensor(shape=x_np.shape, dtype="f16", name="host_x")
|
|
hx._host_buffer = x_np
|
|
x.copy_(hx)
|
|
|
|
y = fc.forward(x)
|
|
results[rank] = y.numpy().astype(np.float32)
|
|
|
|
ctx.multiprocessing.spawn(_worker, args=(), nprocs=ws)
|
|
|
|
# Host-side reference: compute sum_r (x_r @ W_r) = y (same on all ranks).
|
|
expected = np.zeros((M, D_out), dtype=np.float32)
|
|
n_local = D_in // ws
|
|
for r in range(ws):
|
|
x_r = np.full((M, n_local), 0.1 * (r + 1), dtype=np.float32)
|
|
w_r = np.full((n_local, D_out), 0.01 * (r + 1), dtype=np.float32)
|
|
expected += x_r @ w_r
|
|
|
|
for r, out in results.items():
|
|
assert np.allclose(out, expected, rtol=1e-2, atol=1e-2), (
|
|
f"rank {r}: all-reduced output != expected partial sum"
|
|
)
|
|
|
|
|
|
# ── T2.c: rank-consistency post all-reduce ───────────────────────────
|
|
|
|
|
|
def test_row_parallel_rank_identity_post_all_reduce(topology):
|
|
"""T2.c: after all_reduce, all ranks see elementwise-identical output."""
|
|
import kernbench.tp as tp
|
|
from kernbench.runtime_api.tensor import Tensor
|
|
|
|
with _make_ctx(topology) as ctx:
|
|
ctx.distributed.init_process_group(backend="ahbm")
|
|
ws = ctx.distributed.get_world_size()
|
|
tp.initialize_model_parallel(ws)
|
|
|
|
M = 2
|
|
D_in, D_out = 16 * ws, 16
|
|
results: dict[int, np.ndarray] = {}
|
|
|
|
def _worker(rank: int):
|
|
ctx.ahbm.set_device(rank)
|
|
fc = tp.RowParallelLinear(
|
|
in_features=D_in, out_features=D_out, torch=ctx,
|
|
)
|
|
n_local = D_in // ws
|
|
weight_np = np.full((n_local, D_out), 0.01, dtype=np.float16)
|
|
src = Tensor(shape=weight_np.shape, dtype="f16", name="host_w")
|
|
src._host_buffer = weight_np
|
|
fc.weight.copy_(src)
|
|
|
|
x_np = np.full((M, n_local), 0.1, dtype=np.float16)
|
|
x = ctx.zeros(
|
|
(M, n_local), dtype="f16",
|
|
dp=_replicate_dp(), name=f"t2c_x_r{rank}",
|
|
)
|
|
hx = Tensor(shape=x_np.shape, dtype="f16", name="host_x")
|
|
hx._host_buffer = x_np
|
|
x.copy_(hx)
|
|
|
|
y = fc.forward(x)
|
|
results[rank] = y.numpy()
|
|
|
|
ctx.multiprocessing.spawn(_worker, args=(), nprocs=ws)
|
|
|
|
ref = results[0]
|
|
for r, out in results.items():
|
|
assert np.allclose(out, ref, rtol=1e-2, atol=1e-2), (
|
|
f"rank {r} output differs from rank 0 — all_reduce failed to make "
|
|
f"outputs elementwise identical"
|
|
)
|
|
|
|
|
|
def _replicate_dp():
|
|
from kernbench.policy.placement.dp import DPPolicy
|
|
return DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
|