cfc2d74ec49145c53194dc05b6a4c0a3933298b5
2 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 ywkang
|
cfc2d74ec4 |
Refactor ccl_allreduce bench: rank=SIP only, remove rank=PE legacy path
The unified ccl_allreduce bench previously carried two execution models
in one worker with ``if world_size == n_sips:`` branching:
- TP mode (rank = SIP, ADR-0024/0027): proper ProcessGroup semantics.
- Legacy rank = PE mode: single-driver worker allocating one big tensor
distributed across all PEs via _derive_dp, with kernel-level SPMD via
program_id.
The second model is unnecessary — intra-SIP PE-level collectives are
expressed inside the kernel (tl.send/tl.recv with program_id, IPCQ) and
do not need a host-side ProcessGroup. Removing it lets the bench be a
clean reference implementation of the TP launcher.
benches/ccl_allreduce.py:
- Config resolved once in run() via _resolve_cfg -> _BenchCfg dataclass.
- rank != n_sips now raises RuntimeError explicitly.
- _worker / _allocate_rank_tile / _init_with_rank_value / _report each
have one concern; duplicated init + verification paths collapsed.
- _derive_dp and the second verify+print block deleted.
- 166 lines -> 91 lines.
ccl.yaml:
- mesh_allreduce_4 (world_size: 4) and tree_allreduce_7 (world_size: 7)
algorithm entries removed (rank = PE only).
- Algorithm kernel files (kernbench.ccl.algorithms.mesh_allreduce,
tree_allreduce) kept as-is for direct-dispatch future use.
tests/test_ccl_allreduce_matrix.py:
- Matrix shrinks from 7 cases to 3: ring × {tcm, hbm, sram} at ws =
topology SIP count (= 2). mesh_2x2, tree_binary_7, ring_multi_cube,
and the three ring_*_8 cases removed.
tests/test_ccl_performance.py:
- _run_8rank renamed to _run_ring; world_size: 8 override dropped; now
exercises rank = SIP ring all-reduce.
tests/test_mp_spawn.py, tests/test_ccl_ddp_launcher.py:
- Monkeypatch target updated from bench.worker to bench._worker
(signature now takes BenchCfg instead of (rank, world_size)).
555 passed, 1 intentional skip. Tests that directly call
install_ipcq(world_size_override=N) for kernel-level sanity
(test_ccl_hello_world_guide, test_recv_copy_to_dst, test_tl_recv_async,
test_ccl_deadlock_detection) are unchanged — they never went through
the bench and still exercise the kernel-only path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
||
|
ywkang
|
105f1dc09e |
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> |