ADR-0026: DPPolicy intra-device only + ShardSpec structural coords
DPPolicy no longer carries a cross-SIP axis. SIP-level placement is solely controlled by torch.ahbm.set_device(rank) (ADR-0024); DPPolicy itself describes only the cube × PE layout within one SIP. ShardSpec switches to structural (sip, cube, pe) coordinates; the flat pe_index field/property is fully removed — silent drift between global-flat and SIP-local interpretations was a foot-gun flagged by ADR-0024 D11. Breaking API (explicit TypeError / AttributeError): - DPPolicy(sip=...) / DPPolicy(num_sips=...) -> TypeError - ShardSpec.pe_index -> AttributeError - ShardSpec(pe_index=...) -> TypeError - resolve_dp_policy now takes target_sip= (required), no num_sips. Downstream migration: - PE allocator dict keyed by (sip, cube, pe) tuples, in both _ensure_allocators and _free_tensor. deploy_tensor uses tuple lookup. - _create_tensor passes target_sip=current_sip; post-hoc pe_index shifting removed entirely. - launch._compute_local_shape drops the dp.sip branch. - Internal resolvers (column_wise / row_wise / replicate / tiled_*) return _LocalPeShard (cube-local identifier) instead of ShardSpec — resolve_dp_policy lifts them to full structural coords. Tests: - New tests/test_adr0026_dppolicy_intra_device.py (12 tests) pins the contract end-to-end. - test_sip_parallel.py rewritten: SIP composition now modeled as two resolve_dp_policy(target_sip=...) calls (ADR-0024 launcher style). - Call-site migration: test_tensor, test_va_integration, test_va_offset, test_runtime_api_tensor, test_tl_recv_async, test_ccl_* and benches gemm_single_pe, gpt3_qkv, va_offset_verify, ccl_allreduce (legacy branch) all use intra-device DPPolicy and structural ShardSpec. Result: 523 passed, 1 strict xfail (ring_default_ws — unchanged ADR-0024 Phase B blocker; architectural fix deferred to ADR-0027). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,3 +1,14 @@
|
||||
"""Data-parallel placement policy (ADR-0026: intra-device only).
|
||||
|
||||
``DPPolicy`` describes how a tensor is sharded *within a single SIP* across
|
||||
that SIP's cubes and PEs. Crossing the SIP boundary is not a DPPolicy
|
||||
concern: ADR-0024's ``torch.ahbm.set_device(rank)`` picks the SIP, and
|
||||
Megatron-style TP (ADR-0027) expresses multi-SIP tensors when needed.
|
||||
|
||||
``ShardSpec`` is expressed in structural ``(sip, cube, pe)`` coordinates.
|
||||
The former flat ``pe_index`` field/property is fully removed — callers
|
||||
needing a flat integer key compute it explicitly at the call site.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
@@ -7,25 +18,58 @@ from typing import Literal
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DPPolicy:
|
||||
"""Three-level data-parallel policy: sip-level + cube-level + pe-level.
|
||||
"""Intra-device (cube × PE) data-parallel policy.
|
||||
|
||||
Policies:
|
||||
SIP-level placement is controlled by ``torch.ahbm.set_device(rank)``
|
||||
(ADR-0024). For tensors that must cross SIP boundaries, use
|
||||
Megatron-style parallel layers (ADR-0027). DPPolicy itself never
|
||||
crosses a SIP boundary.
|
||||
|
||||
Policies (per axis):
|
||||
- "replicate": full copy at each unit
|
||||
- "column_wise": split K (column) axis across units
|
||||
- "row_wise": split M (row) axis across units
|
||||
|
||||
Optional overrides (default None = use topology dimensions):
|
||||
- num_pes: override PEs per cube (e.g., 1 for single-PE test)
|
||||
- num_cubes: override cubes per SIP (e.g., 1 for single-cube test)
|
||||
- num_sips: override SIP count
|
||||
Optional overrides (``None`` = use topology dimensions):
|
||||
- num_pes: override PEs per cube
|
||||
- num_cubes: override cubes per SIP
|
||||
"""
|
||||
|
||||
sip: Literal["replicate", "column_wise", "row_wise"] = "replicate"
|
||||
cube: Literal["replicate", "column_wise", "row_wise"] = "replicate"
|
||||
pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
|
||||
num_pes: int | None = None
|
||||
num_cubes: int | None = None
|
||||
num_sips: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShardSpec:
|
||||
"""Structural shard placement — ``(sip, cube, pe)`` coord (ADR-0026).
|
||||
|
||||
Global-flat ``pe_index`` was removed: callers must use structural
|
||||
coords directly. If a flat integer key is needed in a local context
|
||||
(e.g. internal dict lookup), compute it explicitly at the call site
|
||||
and do not expose it in any public API.
|
||||
"""
|
||||
|
||||
sip: int
|
||||
cube: int
|
||||
pe: int
|
||||
offset_bytes: int
|
||||
nbytes: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _LocalPeShard:
|
||||
"""Internal — PE resolver's return type (ADR-0026 D3).
|
||||
|
||||
Holds a cube-local PE identifier (``local_pe``) plus the shard's
|
||||
byte payload. Lifted into ``ShardSpec`` with full ``(sip, cube, pe)``
|
||||
coordinates inside ``resolve_dp_policy``.
|
||||
"""
|
||||
|
||||
local_pe: int
|
||||
offset_bytes: int
|
||||
nbytes: int
|
||||
|
||||
|
||||
def _split_shape(
|
||||
@@ -52,14 +96,13 @@ def resolve_dp_policy(
|
||||
itemsize: int,
|
||||
num_pe: int,
|
||||
num_cubes: int = 1,
|
||||
num_sips: int = 1,
|
||||
target_sip: int,
|
||||
) -> list[ShardSpec]:
|
||||
"""Resolve a DPPolicy into a list[ShardSpec] with three-level resolution.
|
||||
"""Resolve a DPPolicy into a list[ShardSpec] on a single SIP.
|
||||
|
||||
SIP-level → cube-level → pe-level.
|
||||
num_cubes is cubes per SIP (not total).
|
||||
ShardSpec.pe_index uses flat indexing:
|
||||
sip_id * num_cubes * num_pe + cube_id * num_pe + pe_id
|
||||
Two-level resolution (cube × PE) within ``target_sip``. Each returned
|
||||
``ShardSpec`` carries ``sip=target_sip`` and cube/pe local to the SIP.
|
||||
No SIP-level split — DPPolicy is intra-device only (ADR-0026).
|
||||
"""
|
||||
_PE_RESOLVERS = {
|
||||
"replicate": replicate,
|
||||
@@ -70,84 +113,61 @@ def resolve_dp_policy(
|
||||
if resolver is None:
|
||||
raise ValueError(f"Unknown pe-level policy: {policy.pe}")
|
||||
|
||||
cubes_per_sip = num_cubes
|
||||
all_shards: list[ShardSpec] = []
|
||||
|
||||
# Level 1: SIP
|
||||
sip_splits = _split_shape(policy.sip, shape, num_sips, itemsize)
|
||||
# Level 1: cube within SIP
|
||||
cube_splits = _split_shape(policy.cube, shape, num_cubes, itemsize)
|
||||
|
||||
for sip_id, (sip_shape, sip_offset) in enumerate(sip_splits):
|
||||
# Level 2: Cube within SIP
|
||||
cube_splits = _split_shape(policy.cube, sip_shape, cubes_per_sip, itemsize)
|
||||
for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
|
||||
# Level 2: PE within cube — resolver returns _LocalPeShard
|
||||
local_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)
|
||||
|
||||
for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
|
||||
# Level 3: PE within cube
|
||||
pe_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)
|
||||
|
||||
for ps in pe_shards:
|
||||
flat_idx = (
|
||||
sip_id * cubes_per_sip * num_pe
|
||||
+ cube_id * num_pe
|
||||
+ ps.pe_index
|
||||
)
|
||||
all_shards.append(ShardSpec(
|
||||
pe_index=flat_idx,
|
||||
offset_bytes=sip_offset + cube_offset + ps.offset_bytes,
|
||||
nbytes=ps.nbytes,
|
||||
))
|
||||
for ls in local_shards:
|
||||
all_shards.append(ShardSpec(
|
||||
sip=target_sip,
|
||||
cube=cube_id,
|
||||
pe=ls.local_pe,
|
||||
offset_bytes=cube_offset + ls.offset_bytes,
|
||||
nbytes=ls.nbytes,
|
||||
))
|
||||
|
||||
return all_shards
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ShardSpec:
|
||||
pe_index: int
|
||||
offset_bytes: int
|
||||
nbytes: int
|
||||
|
||||
|
||||
def column_wise(
|
||||
*, shape: tuple[int, int], itemsize: int, num_pe: int,
|
||||
) -> list[ShardSpec]:
|
||||
) -> list[_LocalPeShard]:
|
||||
"""Split K axis into num_pe equal parts. Each PE gets (M, K/P)."""
|
||||
M, K = shape
|
||||
chunk_k = K // num_pe
|
||||
chunk_bytes = M * chunk_k * itemsize
|
||||
shards = []
|
||||
for i in range(num_pe):
|
||||
shards.append(ShardSpec(
|
||||
pe_index=i,
|
||||
offset_bytes=i * chunk_bytes,
|
||||
nbytes=chunk_bytes,
|
||||
))
|
||||
return shards
|
||||
return [
|
||||
_LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
|
||||
for i in range(num_pe)
|
||||
]
|
||||
|
||||
|
||||
def row_wise(
|
||||
*, shape: tuple[int, int], itemsize: int, num_pe: int,
|
||||
) -> list[ShardSpec]:
|
||||
) -> list[_LocalPeShard]:
|
||||
"""Split M axis into num_pe equal parts. Each PE gets (M/P, K)."""
|
||||
M, K = shape
|
||||
chunk_m = M // num_pe
|
||||
chunk_bytes = chunk_m * K * itemsize
|
||||
shards = []
|
||||
for i in range(num_pe):
|
||||
shards.append(ShardSpec(
|
||||
pe_index=i,
|
||||
offset_bytes=i * chunk_bytes,
|
||||
nbytes=chunk_bytes,
|
||||
))
|
||||
return shards
|
||||
return [
|
||||
_LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
|
||||
for i in range(num_pe)
|
||||
]
|
||||
|
||||
|
||||
def replicate(
|
||||
*, shape: tuple[int, int], itemsize: int, num_pe: int,
|
||||
) -> list[ShardSpec]:
|
||||
) -> list[_LocalPeShard]:
|
||||
"""Full copy per PE. Each PE gets (M, K)."""
|
||||
M, K = shape
|
||||
full_bytes = M * K * itemsize
|
||||
return [
|
||||
ShardSpec(pe_index=i, offset_bytes=0, nbytes=full_bytes)
|
||||
_LocalPeShard(local_pe=i, offset_bytes=0, nbytes=full_bytes)
|
||||
for i in range(num_pe)
|
||||
]
|
||||
|
||||
@@ -155,20 +175,20 @@ def replicate(
|
||||
def tiled_column_major(
|
||||
*, shape: tuple[int, int], itemsize: int, num_pe: int,
|
||||
tile_m: int, tile_k: int,
|
||||
) -> list[ShardSpec]:
|
||||
) -> list[_LocalPeShard]:
|
||||
"""2D tiling, column-major order (K axis first), round-robin across PEs."""
|
||||
M, K = shape
|
||||
tiles_m = ceil(M / tile_m)
|
||||
tiles_k = ceil(K / tile_k)
|
||||
tile_bytes = tile_m * tile_k * itemsize
|
||||
row_bytes = K * itemsize
|
||||
shards = []
|
||||
shards: list[_LocalPeShard] = []
|
||||
idx = 0
|
||||
for mi in range(tiles_m):
|
||||
for ki in range(tiles_k):
|
||||
offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
|
||||
shards.append(ShardSpec(
|
||||
pe_index=idx % num_pe,
|
||||
shards.append(_LocalPeShard(
|
||||
local_pe=idx % num_pe,
|
||||
offset_bytes=offset,
|
||||
nbytes=tile_bytes,
|
||||
))
|
||||
@@ -179,20 +199,20 @@ def tiled_column_major(
|
||||
def tiled_row_major(
|
||||
*, shape: tuple[int, int], itemsize: int, num_pe: int,
|
||||
tile_m: int, tile_k: int,
|
||||
) -> list[ShardSpec]:
|
||||
) -> list[_LocalPeShard]:
|
||||
"""2D tiling, row-major order (M axis first), round-robin across PEs."""
|
||||
M, K = shape
|
||||
tiles_m = ceil(M / tile_m)
|
||||
tiles_k = ceil(K / tile_k)
|
||||
tile_bytes = tile_m * tile_k * itemsize
|
||||
row_bytes = K * itemsize
|
||||
shards = []
|
||||
shards: list[_LocalPeShard] = []
|
||||
idx = 0
|
||||
for ki in range(tiles_k):
|
||||
for mi in range(tiles_m):
|
||||
offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
|
||||
shards.append(ShardSpec(
|
||||
pe_index=idx % num_pe,
|
||||
shards.append(_LocalPeShard(
|
||||
local_pe=idx % num_pe,
|
||||
offset_bytes=offset,
|
||||
nbytes=tile_bytes,
|
||||
))
|
||||
|
||||
Reference in New Issue
Block a user