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:
@@ -48,8 +48,8 @@ def test_from_numpy_creates_host_tensor():
|
||||
assert h._handle is None
|
||||
# Submit a no-op so run_bench has at least one handle.
|
||||
torch.zeros((1, 8), dtype="f16",
|
||||
dp=DPPolicy(sip="replicate", cube="replicate", pe="replicate",
|
||||
num_sips=1, num_cubes=1, num_pes=1),
|
||||
dp=DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=1, num_pes=1),
|
||||
name="dummy")
|
||||
|
||||
_run_with(body)
|
||||
@@ -63,8 +63,8 @@ def test_copy_and_numpy_single_pe():
|
||||
a single-PE (no real sharding) tensor."""
|
||||
|
||||
def body(torch):
|
||||
dp = DPPolicy(sip="replicate", cube="replicate", pe="replicate",
|
||||
num_sips=1, num_cubes=1, num_pes=1)
|
||||
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=1, num_pes=1)
|
||||
t = torch.zeros((1, 16), dtype="f16", dp=dp, name="t")
|
||||
src = np.arange(16, dtype=np.float16).reshape(1, 16)
|
||||
t.copy_(torch.from_numpy(src))
|
||||
@@ -83,8 +83,8 @@ def test_copy_and_numpy_multi_pe_column_wise():
|
||||
|
||||
def body(torch):
|
||||
n_pe = 8
|
||||
dp = DPPolicy(sip="replicate", cube="replicate", pe="column_wise",
|
||||
num_sips=1, num_cubes=1, num_pes=n_pe)
|
||||
dp = DPPolicy(cube="replicate", pe="column_wise",
|
||||
num_cubes=1, num_pes=n_pe)
|
||||
t = torch.zeros((1, n_pe * 4), dtype="f16", dp=dp, name="t")
|
||||
src = np.arange(n_pe * 4, dtype=np.float16).reshape(1, n_pe * 4)
|
||||
t.copy_(torch.from_numpy(src))
|
||||
@@ -107,8 +107,8 @@ def test_copy_and_numpy_multi_cube():
|
||||
n_pe_per_cube = 8
|
||||
n_cubes = 2
|
||||
total = n_cubes * n_pe_per_cube # 16
|
||||
dp = DPPolicy(sip="replicate", cube="column_wise", pe="column_wise",
|
||||
num_sips=1, num_cubes=n_cubes)
|
||||
dp = DPPolicy(cube="column_wise", pe="column_wise",
|
||||
num_cubes=n_cubes)
|
||||
t = torch.zeros((1, total * 4), dtype="f16", dp=dp, name="t")
|
||||
src = np.arange(total * 4, dtype=np.float16).reshape(1, total * 4)
|
||||
t.copy_(torch.from_numpy(src))
|
||||
@@ -126,8 +126,8 @@ def test_copy_shape_mismatch_raises():
|
||||
"""copy_ with mismatched shapes raises ValueError."""
|
||||
|
||||
def body(torch):
|
||||
dp = DPPolicy(sip="replicate", cube="replicate", pe="replicate",
|
||||
num_sips=1, num_cubes=1, num_pes=1)
|
||||
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=1, num_pes=1)
|
||||
t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t")
|
||||
src = np.zeros((1, 16), dtype=np.float16)
|
||||
with pytest.raises(ValueError, match="copy_ shape mismatch"):
|
||||
@@ -143,8 +143,8 @@ def test_setitem_getitem_single_pe():
|
||||
"""Scalar and slice assignment on a single-PE tensor round-trips."""
|
||||
|
||||
def body(torch):
|
||||
dp = DPPolicy(sip="replicate", cube="replicate", pe="replicate",
|
||||
num_sips=1, num_cubes=1, num_pes=1)
|
||||
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=1, num_pes=1)
|
||||
t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t")
|
||||
|
||||
# Scalar broadcast
|
||||
@@ -169,8 +169,8 @@ def test_setitem_getitem_multi_pe_shard_aligned():
|
||||
def body(torch):
|
||||
n_pe = 8
|
||||
n_elem = 4 # per shard
|
||||
dp = DPPolicy(sip="replicate", cube="replicate", pe="column_wise",
|
||||
num_sips=1, num_cubes=1, num_pes=n_pe)
|
||||
dp = DPPolicy(cube="replicate", pe="column_wise",
|
||||
num_cubes=1, num_pes=n_pe)
|
||||
t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t")
|
||||
|
||||
# Write each shard with its rank value
|
||||
@@ -197,8 +197,8 @@ def test_setitem_cross_shard_raises():
|
||||
def body(torch):
|
||||
n_pe = 4
|
||||
n_elem = 4
|
||||
dp = DPPolicy(sip="replicate", cube="replicate", pe="column_wise",
|
||||
num_sips=1, num_cubes=1, num_pes=n_pe)
|
||||
dp = DPPolicy(cube="replicate", pe="column_wise",
|
||||
num_cubes=1, num_pes=n_pe)
|
||||
t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t")
|
||||
with pytest.raises(NotImplementedError, match="spans multiple shards"):
|
||||
t[0, 2:6] = 1.0 # crosses shard 0 (0:4) and shard 1 (4:8)
|
||||
|
||||
Reference in New Issue
Block a user