ADR-0027: Megatron TP API + worker-wait generalization + mp.spawn
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>
This commit is contained in:
@@ -101,12 +101,19 @@ class DataExecutor:
|
||||
p = op.params
|
||||
if "src_a_addr" not in p:
|
||||
return # composite record without full params
|
||||
space = p.get("addr_space", "tcm")
|
||||
default_space = p.get("addr_space", "tcm")
|
||||
# ADR-0027: per-operand + output spaces (fall back to single space
|
||||
# for legacy records without explicit space keys).
|
||||
src_a_space = p.get("src_a_space", default_space)
|
||||
src_b_space = p.get("src_b_space", default_space)
|
||||
dst_space = p.get("dst_space", default_space)
|
||||
dtype_in = p.get("dtype_in", "f16")
|
||||
dtype_out = p.get("dtype_out", dtype_in)
|
||||
|
||||
a = self.store.read(space, p["src_a_addr"], shape=p.get("shape_a"), dtype=dtype_in)
|
||||
b = self.store.read(space, p["src_b_addr"], shape=p.get("shape_b"), dtype=dtype_in)
|
||||
a = self.store.read(src_a_space, p["src_a_addr"],
|
||||
shape=p.get("shape_a"), dtype=dtype_in)
|
||||
b = self.store.read(src_b_space, p["src_b_addr"],
|
||||
shape=p.get("shape_b"), dtype=dtype_in)
|
||||
|
||||
# Compute in higher precision if specified
|
||||
dtype_acc = p.get("dtype_acc", "f32")
|
||||
@@ -114,7 +121,7 @@ class DataExecutor:
|
||||
b_f = b.astype(_resolve_dtype(dtype_acc))
|
||||
result = np.matmul(a_f, b_f).astype(_resolve_dtype(dtype_out))
|
||||
|
||||
self.store.write(space, p["dst_addr"], result)
|
||||
self.store.write(dst_space, p["dst_addr"], result)
|
||||
|
||||
def _execute_math(self, op: OpRecord) -> None:
|
||||
"""Execute math op: unary, binary, or reduction."""
|
||||
|
||||
@@ -79,16 +79,24 @@ class OpLogger:
|
||||
snaps.append(None)
|
||||
params["input_snapshots"] = snaps
|
||||
elif op_name == "dma_write":
|
||||
try:
|
||||
arr = self._memory_store.read(
|
||||
params["src_space"], params["src_addr"],
|
||||
shape=params.get("shape"), dtype=params.get("dtype"),
|
||||
)
|
||||
params["snapshot"] = (
|
||||
arr.copy() if hasattr(arr, "copy") else arr
|
||||
)
|
||||
except Exception:
|
||||
params["snapshot"] = None
|
||||
# ADR-0027 fix: only snapshot HBM sources. TCM (PE scratch)
|
||||
# sources are repopulated by Phase 2 math/gemm replay —
|
||||
# capturing a Phase-1-time snapshot here would pick up stale
|
||||
# data from a PRIOR kernel's Phase 2 output that aliased the
|
||||
# same scratch address, causing the later kernel's replay
|
||||
# to write that stale value instead of the fresh math
|
||||
# result. See ADR-0027 postmortem (TP gemm → all_reduce).
|
||||
if params.get("src_space") == "hbm":
|
||||
try:
|
||||
arr = self._memory_store.read(
|
||||
params["src_space"], params["src_addr"],
|
||||
shape=params.get("shape"), dtype=params.get("dtype"),
|
||||
)
|
||||
params["snapshot"] = (
|
||||
arr.copy() if hasattr(arr, "copy") else arr
|
||||
)
|
||||
except Exception:
|
||||
params["snapshot"] = None
|
||||
self._records.append(OpRecord(
|
||||
t_start=pending["t_start"],
|
||||
t_end=t,
|
||||
@@ -167,6 +175,13 @@ def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]:
|
||||
"dtype_in": msg.a.dtype,
|
||||
"dtype_out": msg.out.dtype,
|
||||
"m": msg.m, "k": msg.k, "n": msg.n,
|
||||
# ADR-0027: preserve per-operand + output MemoryStore spaces so
|
||||
# Phase 2 replay can resolve HBM-resident operands (e.g. tl.load
|
||||
# results keep space="hbm"). Absent → DataExecutor falls back
|
||||
# to the legacy single-space mode via ``addr_space``.
|
||||
"src_a_space": getattr(msg.a, "space", "tcm"),
|
||||
"src_b_space": getattr(msg.b, "space", "tcm"),
|
||||
"dst_space": getattr(msg.out, "space", "tcm"),
|
||||
}
|
||||
if isinstance(msg, MathCmd):
|
||||
return "math", msg.op, {
|
||||
@@ -181,10 +196,27 @@ def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]:
|
||||
"axis": msg.axis,
|
||||
}
|
||||
if isinstance(msg, CompositeCmd):
|
||||
return "gemm" if msg.op == "gemm" else "math", f"composite_{msg.op}", {
|
||||
params: dict[str, Any] = {
|
||||
"op": msg.op,
|
||||
"out_addr": msg.out_addr,
|
||||
"out_nbytes": msg.out_nbytes,
|
||||
}
|
||||
# ADR-0027: preserve operand info so Phase 2 DataExecutor can replay
|
||||
# the composite's numerical effect (treat it like a GemmCmd).
|
||||
if msg.op == "gemm" and msg.a is not None and msg.b is not None:
|
||||
params.update({
|
||||
"src_a_addr": msg.a.addr,
|
||||
"src_b_addr": msg.b.addr,
|
||||
"shape_a": msg.a.shape,
|
||||
"shape_b": msg.b.shape,
|
||||
"dtype_in": msg.a.dtype,
|
||||
"dtype_out": msg.a.dtype,
|
||||
"src_a_space": getattr(msg.a, "space", "hbm"),
|
||||
"src_b_space": getattr(msg.b, "space", "hbm"),
|
||||
"dst_space": "hbm",
|
||||
# dst_addr alias so DataExecutor._execute_gemm picks it up.
|
||||
"dst_addr": msg.out_addr,
|
||||
})
|
||||
return "gemm" if msg.op == "gemm" else "math", f"composite_{msg.op}", params
|
||||
# Fallback for unknown data_op messages
|
||||
return "unknown", type(msg).__name__, {}
|
||||
|
||||
Reference in New Issue
Block a user