Support multi axis unsharding in dump comparator (#19280)
This commit is contained in:
fzyzcjy
@@ -82,9 +82,8 @@ def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]:
|
||||
|
||||
dim_specs = parse_dims(dims_str)
|
||||
parallel_infos = [normalize_parallel_info(meta) for meta in metas]
|
||||
plan = compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos)
|
||||
|
||||
return [plan] if plan is not None else []
|
||||
return compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos)
|
||||
|
||||
|
||||
def _extract_tensors(
|
||||
@@ -105,15 +104,12 @@ def _execute_plans(
|
||||
return None
|
||||
return tensors[0]
|
||||
|
||||
assert len(plans) <= 1, "multi-plan not supported yet"
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
if isinstance(plan, UnshardPlan):
|
||||
# TODO: incorrect `tensors_by_world_rank` if multi UnshardPlan
|
||||
tensors = execute_unshard_plan(
|
||||
plan, tensors_by_world_rank=dict(enumerate(tensors))
|
||||
)
|
||||
current = execute_unshard_plan(plan, current)
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown {plan=}")
|
||||
|
||||
return tensors
|
||||
assert len(current) == 1
|
||||
return current[0]
|
||||
|
||||
@@ -2,26 +2,25 @@ import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.unshard.types import (
|
||||
ConcatParams,
|
||||
UnshardParams,
|
||||
UnshardPlan,
|
||||
)
|
||||
|
||||
|
||||
def execute_unshard_plan(
|
||||
plan: UnshardPlan,
|
||||
tensors_by_world_rank: dict[int, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
ordered_tensors = [
|
||||
tensors_by_world_rank[world_rank]
|
||||
for world_rank in plan.world_ranks_by_axis_rank
|
||||
]
|
||||
|
||||
return _apply_unshard(plan, ordered_tensors)
|
||||
tensors: list[torch.Tensor],
|
||||
) -> list[torch.Tensor]:
|
||||
result: list[torch.Tensor] = []
|
||||
for group in plan.groups:
|
||||
group_tensors = [tensors[i] for i in group]
|
||||
result.append(_apply_unshard(plan.params, group_tensors))
|
||||
return result
|
||||
|
||||
|
||||
def _apply_unshard(
|
||||
plan: UnshardPlan, ordered_tensors: list[torch.Tensor]
|
||||
params: UnshardParams, ordered_tensors: list[torch.Tensor]
|
||||
) -> torch.Tensor:
|
||||
params = plan.params
|
||||
if isinstance(params, ConcatParams):
|
||||
return _unshard_concat(ordered_tensors, dim=params.dim)
|
||||
# Phase 2: ReduceSumParams, CpZigzagParams
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from typing import Optional
|
||||
from collections import defaultdict
|
||||
from typing import NamedTuple
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.unshard.types import (
|
||||
@@ -8,64 +9,117 @@ from sglang.srt.debug_utils.comparator.unshard.types import (
|
||||
UnshardPlan,
|
||||
)
|
||||
|
||||
# _CoordsList[tensor_index][axis] =
|
||||
# the axis_rank (shard position) of the tensor_index-th tensor along `axis`
|
||||
# (e.g. coords[2] = {TP: 3} means tensor 2 is the 3rd shard in TP axis)
|
||||
_CoordsList = list[dict[ParallelAxis, int]]
|
||||
|
||||
|
||||
class _GroupResult(NamedTuple):
|
||||
groups: list[list[int]]
|
||||
projected_coords: _CoordsList
|
||||
|
||||
|
||||
def compute_unshard_plan(
|
||||
dim_specs: list[DimSpec],
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]],
|
||||
) -> Optional[UnshardPlan]:
|
||||
) -> list[UnshardPlan]:
|
||||
if not parallel_infos:
|
||||
raise ValueError("parallel_infos must not be empty")
|
||||
|
||||
sharded_axes: dict[ParallelAxis, tuple[int, DimSpec]] = {}
|
||||
for dim_idx, spec in enumerate(dim_specs):
|
||||
if spec.parallel is not None:
|
||||
sharded_axes[spec.parallel] = (dim_idx, spec)
|
||||
sharded_axis_infos: dict[ParallelAxis, tuple[int, DimSpec]] = {
|
||||
spec.parallel: (dim_idx, spec)
|
||||
for dim_idx, spec in enumerate(dim_specs)
|
||||
if spec.parallel is not None
|
||||
}
|
||||
if not sharded_axis_infos:
|
||||
return []
|
||||
|
||||
if len(sharded_axes) > 1:
|
||||
raise NotImplementedError(
|
||||
f"Multi-axis unshard is not supported. "
|
||||
f"Got {len(sharded_axes)} sharded axes: {sorted(a.value for a in sharded_axes)}"
|
||||
_validate(sharded_axes=set(sharded_axis_infos), parallel_infos=parallel_infos)
|
||||
|
||||
current_coords: _CoordsList = [
|
||||
{axis: info[axis].axis_rank for axis in sharded_axis_infos}
|
||||
for info in parallel_infos
|
||||
]
|
||||
|
||||
plans: list[UnshardPlan] = []
|
||||
for axis, (dim_index, spec) in sharded_axis_infos.items():
|
||||
result = _group_and_project(
|
||||
current_coords=current_coords,
|
||||
target_axis=axis,
|
||||
)
|
||||
|
||||
if not sharded_axes:
|
||||
return None
|
||||
plans.append(
|
||||
UnshardPlan(
|
||||
axis=axis,
|
||||
params=_resolve_unshard_params(spec=spec, dim_index=dim_index),
|
||||
groups=result.groups,
|
||||
)
|
||||
)
|
||||
|
||||
axis_name, (dim_idx, spec) = next(iter(sharded_axes.items()))
|
||||
current_coords = result.projected_coords
|
||||
|
||||
expected_size: Optional[int] = None
|
||||
rank_to_world: dict[int, int] = {}
|
||||
return plans
|
||||
|
||||
for world_rank, pinfo in enumerate(parallel_infos):
|
||||
if axis_name not in pinfo:
|
||||
continue
|
||||
|
||||
ainfo = pinfo[axis_name]
|
||||
def _validate(
|
||||
*,
|
||||
sharded_axes: set[ParallelAxis],
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]],
|
||||
) -> None:
|
||||
"""Check that every rank has all sharded axes, sizes are consistent, and ranks are complete."""
|
||||
axis_sizes: dict[ParallelAxis, int] = {}
|
||||
|
||||
if expected_size is None:
|
||||
expected_size = ainfo.axis_size
|
||||
elif ainfo.axis_size != expected_size:
|
||||
for world_rank, parallel_info in enumerate(parallel_infos):
|
||||
for axis in sharded_axes:
|
||||
if axis not in parallel_info:
|
||||
raise ValueError(
|
||||
f"world_rank={world_rank} missing parallel_info for "
|
||||
f"sharded axis {axis.value!r}"
|
||||
)
|
||||
|
||||
axis_info = parallel_info[axis]
|
||||
if axis not in axis_sizes:
|
||||
axis_sizes[axis] = axis_info.axis_size
|
||||
elif axis_info.axis_size != axis_sizes[axis]:
|
||||
raise ValueError(
|
||||
f"Inconsistent axis_size for {axis.value}: "
|
||||
f"expected {axis_sizes[axis]}, got {axis_info.axis_size} "
|
||||
f"at world_rank={world_rank}"
|
||||
)
|
||||
|
||||
for axis, expected_size in axis_sizes.items():
|
||||
seen_ranks = {info[axis].axis_rank for info in parallel_infos}
|
||||
if seen_ranks != set(range(expected_size)):
|
||||
raise ValueError(
|
||||
f"Inconsistent axis_size for {axis_name.value}: "
|
||||
f"expected {expected_size}, got {ainfo.axis_size} "
|
||||
f"at world_rank={world_rank}"
|
||||
f"axis_rank coverage for {axis.value} is incomplete: "
|
||||
f"got {sorted(seen_ranks)}, expected 0..{expected_size - 1}"
|
||||
)
|
||||
|
||||
rank_to_world.setdefault(ainfo.axis_rank, world_rank)
|
||||
|
||||
if expected_size is None:
|
||||
raise ValueError(f"No parallel_info found for sharded axis {axis_name.value!r}")
|
||||
def _group_and_project(
|
||||
*,
|
||||
current_coords: _CoordsList,
|
||||
target_axis: ParallelAxis,
|
||||
) -> _GroupResult:
|
||||
"""Group tensors by other-axes coords, sort within group by target_axis rank."""
|
||||
# buckets[coords_excluding_target] = [(axis_rank, tensor_index), ...]
|
||||
# e.g. when target_axis=CP: buckets[{(TP,0)}] = [(0, 1), (1, 3)]
|
||||
# means tensor 1 (CP rank 0) and tensor 3 (CP rank 1) share TP rank 0
|
||||
buckets: dict[frozenset, list[tuple[int, int]]] = defaultdict(list)
|
||||
|
||||
if set(rank_to_world.keys()) != set(range(expected_size)):
|
||||
raise ValueError(
|
||||
f"axis_rank coverage for {axis_name.value} is incomplete: "
|
||||
f"got {sorted(rank_to_world.keys())}, expected 0..{expected_size - 1}"
|
||||
)
|
||||
for idx, coords in enumerate(current_coords):
|
||||
key = frozenset((k, v) for k, v in coords.items() if k != target_axis)
|
||||
buckets[key].append((coords[target_axis], idx))
|
||||
|
||||
return UnshardPlan(
|
||||
axis=spec.parallel,
|
||||
params=_resolve_unshard_params(spec=spec, dim_index=dim_idx),
|
||||
world_ranks_by_axis_rank=[rank_to_world[i] for i in range(expected_size)],
|
||||
)
|
||||
groups: list[list[int]] = []
|
||||
projected: _CoordsList = []
|
||||
for key in sorted(buckets, key=lambda k: sorted((a.value, v) for a, v in k)):
|
||||
entries = sorted(buckets[key])
|
||||
groups.append([idx for _, idx in entries])
|
||||
projected.append(dict(key))
|
||||
|
||||
return _GroupResult(groups=groups, projected_coords=projected)
|
||||
|
||||
|
||||
def _resolve_unshard_params(*, spec: DimSpec, dim_index: int) -> UnshardParams:
|
||||
|
||||
@@ -27,7 +27,12 @@ UnshardParams = ConcatParams
|
||||
class UnshardPlan(_FrozenBase):
|
||||
axis: ParallelAxis
|
||||
params: UnshardParams
|
||||
world_ranks_by_axis_rank: list[int]
|
||||
# groups[i] = indices in the input tensor list, which will be operated (e.g. concat) into i-th output tensor.
|
||||
#
|
||||
# Multistep example (CP=2, TP=2, 4 input tensors):
|
||||
# plan[0] (CP): groups=[[0,2],[1,3]] — 4 tensors → 2 tensors
|
||||
# plan[1] (TP): groups=[[0,1]] — 2 tensors → 1 tensor
|
||||
groups: list[list[int]]
|
||||
|
||||
|
||||
# Union of all plan types. Future pipeline components (e.g. reduction,
|
||||
|
||||
@@ -701,6 +701,114 @@ class TestEntrypointGroupingLogical:
|
||||
assert summary.total == 2
|
||||
assert summary.passed == 2
|
||||
|
||||
def test_cp_tp_unshard(self, tmp_path, capsys):
|
||||
"""CP=2 + TP=2: multi-axis shards are unsharded before comparison."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 16)
|
||||
full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
for side_dir, full_tensor in [
|
||||
(baseline_dir, full_baseline),
|
||||
(target_dir, full_target),
|
||||
]:
|
||||
_create_cp_tp_sharded_dumps(
|
||||
side_dir,
|
||||
full_tensor=full_tensor,
|
||||
name="hidden",
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
head_dim=2,
|
||||
dims_str="b s(cp) h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=0.01,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.name == "hidden"
|
||||
|
||||
def test_cp_tp_different_sizes(self, tmp_path, capsys):
|
||||
"""Baseline CP=2+TP=2 vs target CP=1+TP=4: both sides independently unshard."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 16)
|
||||
full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
_create_cp_tp_sharded_dumps(
|
||||
baseline_dir,
|
||||
full_tensor=full_baseline,
|
||||
name="hidden",
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
seq_dim=1,
|
||||
head_dim=2,
|
||||
dims_str="b s(cp) h(tp)",
|
||||
)
|
||||
|
||||
_create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=full_target,
|
||||
name="hidden",
|
||||
tp_size=4,
|
||||
shard_dim=2,
|
||||
dims_str="b s h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=0.01,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
_assert_single_comparison_passed(records)
|
||||
|
||||
def test_ep_cp_tp_three_axis_unshard(self, tmp_path, capsys):
|
||||
"""EP=2 + CP=2 + TP=2: three-axis shards are unsharded before comparison."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 16, 32)
|
||||
full_target = full_baseline + torch.randn(4, 8, 16, 32) * 0.001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
for side_dir, full_tensor in [
|
||||
(baseline_dir, full_baseline),
|
||||
(target_dir, full_target),
|
||||
]:
|
||||
_create_ep_cp_tp_sharded_dumps(
|
||||
side_dir,
|
||||
full_tensor=full_tensor,
|
||||
name="hidden",
|
||||
ep_size=2,
|
||||
cp_size=2,
|
||||
tp_size=2,
|
||||
expert_dim=1,
|
||||
seq_dim=2,
|
||||
head_dim=3,
|
||||
dims_str="b e(ep) s(cp) h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(
|
||||
baseline_dir / _FIXED_EXP_NAME,
|
||||
target_dir / _FIXED_EXP_NAME,
|
||||
diff_threshold=0.01,
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.name == "hidden"
|
||||
|
||||
|
||||
# --------------------------- Assertion helpers -------------------
|
||||
|
||||
@@ -825,6 +933,84 @@ def _create_rank_dump(
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
def _create_cp_tp_sharded_dumps(
|
||||
directory: Path,
|
||||
*,
|
||||
full_tensor: torch.Tensor,
|
||||
name: str,
|
||||
cp_size: int,
|
||||
tp_size: int,
|
||||
seq_dim: int,
|
||||
head_dim: int,
|
||||
dims_str: str,
|
||||
num_steps: int = 1,
|
||||
) -> Path:
|
||||
"""Create CP+TP multi-axis sharded dump files from a full tensor."""
|
||||
cp_chunks = list(full_tensor.chunk(cp_size, dim=seq_dim))
|
||||
rank = 0
|
||||
for cp_rank in range(cp_size):
|
||||
tp_chunks = list(cp_chunks[cp_rank].chunk(tp_size, dim=head_dim))
|
||||
for tp_rank in range(tp_size):
|
||||
_create_rank_dump(
|
||||
directory,
|
||||
rank=rank,
|
||||
name=name,
|
||||
tensor=tp_chunks[tp_rank],
|
||||
dims=dims_str,
|
||||
parallel_info={
|
||||
"cp_rank": cp_rank,
|
||||
"cp_size": cp_size,
|
||||
"tp_rank": tp_rank,
|
||||
"tp_size": tp_size,
|
||||
},
|
||||
num_steps=num_steps,
|
||||
)
|
||||
rank += 1
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
def _create_ep_cp_tp_sharded_dumps(
|
||||
directory: Path,
|
||||
*,
|
||||
full_tensor: torch.Tensor,
|
||||
name: str,
|
||||
ep_size: int,
|
||||
cp_size: int,
|
||||
tp_size: int,
|
||||
expert_dim: int,
|
||||
seq_dim: int,
|
||||
head_dim: int,
|
||||
dims_str: str,
|
||||
num_steps: int = 1,
|
||||
) -> Path:
|
||||
"""Create EP+CP+TP three-axis sharded dump files from a full tensor."""
|
||||
ep_chunks = list(full_tensor.chunk(ep_size, dim=expert_dim))
|
||||
rank = 0
|
||||
for ep_rank in range(ep_size):
|
||||
cp_chunks = list(ep_chunks[ep_rank].chunk(cp_size, dim=seq_dim))
|
||||
for cp_rank in range(cp_size):
|
||||
tp_chunks = list(cp_chunks[cp_rank].chunk(tp_size, dim=head_dim))
|
||||
for tp_rank in range(tp_size):
|
||||
_create_rank_dump(
|
||||
directory,
|
||||
rank=rank,
|
||||
name=name,
|
||||
tensor=tp_chunks[tp_rank],
|
||||
dims=dims_str,
|
||||
parallel_info={
|
||||
"ep_rank": ep_rank,
|
||||
"ep_size": ep_size,
|
||||
"cp_rank": cp_rank,
|
||||
"cp_size": cp_size,
|
||||
"tp_rank": tp_rank,
|
||||
"tp_size": tp_size,
|
||||
},
|
||||
num_steps=num_steps,
|
||||
)
|
||||
rank += 1
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
def _create_tp_sharded_dumps(
|
||||
directory: Path,
|
||||
*,
|
||||
|
||||
@@ -3,8 +3,11 @@ import sys
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims import parse_dims
|
||||
from sglang.srt.debug_utils.comparator.unshard.executor import execute_unshard_plan
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
|
||||
from sglang.srt.debug_utils.comparator.unshard.executor import (
|
||||
_apply_unshard,
|
||||
execute_unshard_plan,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
|
||||
from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -18,37 +21,240 @@ class TestExecuteUnshardPlan:
|
||||
shards = list(full_tensor.chunk(4, dim=1))
|
||||
|
||||
dim_specs = parse_dims("b h(tp) d")
|
||||
parallel_infos = [{"tp": AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)]
|
||||
plan = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert plan is not None
|
||||
tensors_by_rank = {i: shards[i] for i in range(4)}
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
|
||||
]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
|
||||
result = execute_unshard_plan(plan, tensors_by_rank)
|
||||
assert torch.allclose(result, full_tensor)
|
||||
result = execute_unshard_plan(plans[0], shards)
|
||||
assert len(result) == 1
|
||||
assert torch.allclose(result[0], full_tensor)
|
||||
|
||||
def test_scrambled_world_ranks_correct_result(self) -> None:
|
||||
full_tensor = torch.randn(4, 8)
|
||||
shards = list(full_tensor.chunk(4, dim=0))
|
||||
|
||||
parallel_infos = [
|
||||
{"tp": AxisInfo(axis_rank=2, axis_size=4)},
|
||||
{"tp": AxisInfo(axis_rank=0, axis_size=4)},
|
||||
{"tp": AxisInfo(axis_rank=3, axis_size=4)},
|
||||
{"tp": AxisInfo(axis_rank=1, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
|
||||
]
|
||||
dim_specs = parse_dims("h(tp) d")
|
||||
plan = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert plan is not None
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
|
||||
tensors_by_rank = {
|
||||
0: shards[2],
|
||||
1: shards[0],
|
||||
2: shards[3],
|
||||
3: shards[1],
|
||||
}
|
||||
tensors_ordered_by_world_rank = [
|
||||
shards[2], # world_rank=0, axis_rank=2
|
||||
shards[0], # world_rank=1, axis_rank=0
|
||||
shards[3], # world_rank=2, axis_rank=3
|
||||
shards[1], # world_rank=3, axis_rank=1
|
||||
]
|
||||
|
||||
result = execute_unshard_plan(plan, tensors_by_rank)
|
||||
assert torch.allclose(result, full_tensor)
|
||||
result = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
|
||||
assert len(result) == 1
|
||||
assert torch.allclose(result[0], full_tensor)
|
||||
|
||||
def test_single_step_reduces_tensor_count(self) -> None:
|
||||
"""8 tensors with 2 groups of 4 produce 2 output tensors."""
|
||||
full_a = torch.randn(4, 8)
|
||||
full_b = torch.randn(4, 8)
|
||||
shards_a = list(full_a.chunk(4, dim=0))
|
||||
shards_b = list(full_b.chunk(4, dim=0))
|
||||
|
||||
dim_specs = parse_dims("s(cp) h(tp)")
|
||||
parallel_infos = []
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(4):
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=4),
|
||||
}
|
||||
)
|
||||
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 2
|
||||
|
||||
tensors: list[torch.Tensor] = []
|
||||
for cp_rank in range(2):
|
||||
source = shards_a if cp_rank == 0 else shards_b
|
||||
for tp_rank in range(4):
|
||||
tensors.append(source[tp_rank])
|
||||
|
||||
intermediate = execute_unshard_plan(plans[0], tensors)
|
||||
assert len(intermediate) == 4
|
||||
|
||||
final = execute_unshard_plan(plans[1], intermediate)
|
||||
assert len(final) == 1
|
||||
|
||||
def test_cp_tp_concat(self) -> None:
|
||||
"""CP=2 + TP=2: multi-step unshard reconstructs original tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
cp_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos = []
|
||||
for cp_rank in range(2):
|
||||
tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=2))
|
||||
for tp_rank in range(2):
|
||||
tensors.append(tp_chunks[tp_rank])
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b s(cp) h(tp)")
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 2
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
def test_cp_tp_scrambled(self) -> None:
|
||||
"""Scrambled world_ranks for CP=2 + TP=2 still reconstruct correctly."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
cp_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
shard_map: dict[tuple[int, int], torch.Tensor] = {}
|
||||
for cp_rank in range(2):
|
||||
tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=2))
|
||||
for tp_rank in range(2):
|
||||
shard_map[(cp_rank, tp_rank)] = tp_chunks[tp_rank]
|
||||
|
||||
scrambled_assignment = [
|
||||
(1, 1), # world_rank=0
|
||||
(0, 0), # world_rank=1
|
||||
(1, 0), # world_rank=2
|
||||
(0, 1), # world_rank=3
|
||||
]
|
||||
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos = []
|
||||
for cp_rank, tp_rank in scrambled_assignment:
|
||||
tensors.append(shard_map[(cp_rank, tp_rank)])
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b s(cp) h(tp)")
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 2
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
def test_unsupported_params_type_raises(self) -> None:
|
||||
"""_apply_unshard raises ValueError for unknown params type."""
|
||||
|
||||
class _FakeParams:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported unshard"):
|
||||
_apply_unshard(_FakeParams(), [torch.randn(2, 2)])
|
||||
|
||||
def test_cp_tp_ep_three_axis_concat(self) -> None:
|
||||
"""CP=2 + TP=2 + EP=2: three-step unshard reconstructs original tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16, 32)
|
||||
|
||||
ep_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
shard_map: dict[tuple[int, int, int], torch.Tensor] = {}
|
||||
for ep_rank in range(2):
|
||||
cp_chunks = list(ep_chunks[ep_rank].chunk(2, dim=2))
|
||||
for cp_rank in range(2):
|
||||
tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=3))
|
||||
for tp_rank in range(2):
|
||||
shard_map[(ep_rank, cp_rank, tp_rank)] = tp_chunks[tp_rank]
|
||||
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for ep_rank in range(2):
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
tensors.append(shard_map[(ep_rank, cp_rank, tp_rank)])
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 3
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
def test_cp_tp_ep_scrambled_three_axis(self) -> None:
|
||||
"""Scrambled ranks for CP=2 + TP=2 + EP=2 still reconstruct correctly."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16, 32)
|
||||
|
||||
ep_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
shard_map: dict[tuple[int, int, int], torch.Tensor] = {}
|
||||
for ep_rank in range(2):
|
||||
cp_chunks = list(ep_chunks[ep_rank].chunk(2, dim=2))
|
||||
for cp_rank in range(2):
|
||||
tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=3))
|
||||
for tp_rank in range(2):
|
||||
shard_map[(ep_rank, cp_rank, tp_rank)] = tp_chunks[tp_rank]
|
||||
|
||||
scrambled_assignment = [
|
||||
(1, 0, 1), # world_rank=0
|
||||
(0, 1, 0), # world_rank=1
|
||||
(1, 1, 0), # world_rank=2
|
||||
(0, 0, 0), # world_rank=3
|
||||
(0, 1, 1), # world_rank=4
|
||||
(1, 0, 0), # world_rank=5
|
||||
(0, 0, 1), # world_rank=6
|
||||
(1, 1, 1), # world_rank=7
|
||||
]
|
||||
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for ep_rank, cp_rank, tp_rank in scrambled_assignment:
|
||||
tensors.append(shard_map[(ep_rank, cp_rank, tp_rank)])
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 3
|
||||
|
||||
current = tensors
|
||||
for plan in plans:
|
||||
current = execute_unshard_plan(plan, current)
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(current[0], full_tensor)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -16,12 +16,12 @@ class TestComputeUnshardPlan:
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
|
||||
]
|
||||
plan = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert plan is not None
|
||||
assert plan.axis == ParallelAxis.TP
|
||||
assert plan.params.dim == 2
|
||||
assert plan.world_ranks_by_axis_rank == [0, 1, 2, 3]
|
||||
assert len(plans) == 1
|
||||
assert plans[0].axis == ParallelAxis.TP
|
||||
assert plans[0].params.dim == 2
|
||||
assert plans[0].groups == [[0, 1, 2, 3]]
|
||||
|
||||
def test_inconsistent_axis_size_raises(self) -> None:
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
@@ -35,7 +35,7 @@ class TestComputeUnshardPlan:
|
||||
def test_missing_axis_in_parallel_info_raises(self) -> None:
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}]
|
||||
with pytest.raises(ValueError, match="No parallel_info found"):
|
||||
with pytest.raises(ValueError, match="missing parallel_info"):
|
||||
compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
def test_empty_parallel_infos_raises(self) -> None:
|
||||
@@ -52,29 +52,193 @@ class TestComputeUnshardPlan:
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
|
||||
]
|
||||
plan = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert plan is not None
|
||||
assert plan.world_ranks_by_axis_rank == [1, 3, 0, 2]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
assert plans[0].groups == [[1, 3, 0, 2]]
|
||||
|
||||
def test_no_sharded_axes_returns_none(self) -> None:
|
||||
def test_no_sharded_axes_returns_empty(self) -> None:
|
||||
dim_specs = parse_dims("b s d")
|
||||
parallel_infos = [{}]
|
||||
plan = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert plan is None
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert plans == []
|
||||
|
||||
def test_multi_axis_raises(self) -> None:
|
||||
dim_specs = parse_dims("h(tp) s(cp)")
|
||||
def test_multi_axis_plan(self) -> None:
|
||||
"""Multi-axis (TP + CP) produces a 2-step plan."""
|
||||
dim_specs = parse_dims("s(cp) h(tp)")
|
||||
parallel_infos = [
|
||||
{
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="Multi-axis unshard"):
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
assert plans[0].axis == ParallelAxis.CP
|
||||
assert plans[1].axis == ParallelAxis.TP
|
||||
|
||||
def test_cp_tp_plan(self) -> None:
|
||||
"""CP=2 + TP=4 produces correct 2-step plan with correct groups."""
|
||||
dim_specs = parse_dims("s(cp) h(tp)")
|
||||
parallel_infos = []
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(4):
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=4),
|
||||
}
|
||||
)
|
||||
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
|
||||
cp_plan = plans[0]
|
||||
assert cp_plan.axis == ParallelAxis.CP
|
||||
assert len(cp_plan.groups) == 4
|
||||
for group in cp_plan.groups:
|
||||
assert len(group) == 2
|
||||
|
||||
tp_plan = plans[1]
|
||||
assert tp_plan.axis == ParallelAxis.TP
|
||||
assert len(tp_plan.groups) == 1
|
||||
assert len(tp_plan.groups[0]) == 4
|
||||
|
||||
def test_cp_tp_scrambled_ranks(self) -> None:
|
||||
"""Scrambled rank assignment still produces correct plan."""
|
||||
dim_specs = parse_dims("s(cp) h(tp)")
|
||||
parallel_infos = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 2
|
||||
|
||||
cp_plan = plans[0]
|
||||
assert cp_plan.axis == ParallelAxis.CP
|
||||
assert len(cp_plan.groups) == 2
|
||||
for group in cp_plan.groups:
|
||||
assert len(group) == 2
|
||||
|
||||
tp_plan = plans[1]
|
||||
assert tp_plan.axis == ParallelAxis.TP
|
||||
assert len(tp_plan.groups) == 1
|
||||
assert len(tp_plan.groups[0]) == 2
|
||||
|
||||
def test_axis_rank_coverage_incomplete_raises(self) -> None:
|
||||
"""TP size=4 but only ranks 0,1,3 provided (missing rank 2)."""
|
||||
dim_specs = parse_dims("h(tp)")
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
|
||||
]
|
||||
with pytest.raises(ValueError, match="axis_rank coverage.*incomplete"):
|
||||
compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
def test_reduction_not_implemented_raises(self) -> None:
|
||||
dim_specs = parse_dims("h(tp,partial)")
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="reduction"):
|
||||
compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
def test_ordering_not_natural_raises(self) -> None:
|
||||
dim_specs = parse_dims("s(cp,zigzag)")
|
||||
parallel_infos = [
|
||||
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="ordering"):
|
||||
compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
def test_ordering_natural_accepted(self) -> None:
|
||||
dim_specs = parse_dims("s(cp,natural)")
|
||||
parallel_infos = [
|
||||
{ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
assert plans[0].axis == ParallelAxis.CP
|
||||
|
||||
def test_three_axis_plan(self) -> None:
|
||||
"""EP=2 + CP=2 + TP=2 produces a 3-step plan."""
|
||||
dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for ep_rank in range(2):
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
plans = compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
assert len(plans) == 3
|
||||
assert plans[0].axis == ParallelAxis.EP
|
||||
assert plans[1].axis == ParallelAxis.CP
|
||||
assert plans[2].axis == ParallelAxis.TP
|
||||
|
||||
# Step 0 (EP): 8 tensors → 4 (groups of 2)
|
||||
assert len(plans[0].groups) == 4
|
||||
for group in plans[0].groups:
|
||||
assert len(group) == 2
|
||||
|
||||
# Step 1 (CP): 4 tensors → 2 (groups of 2)
|
||||
assert len(plans[1].groups) == 2
|
||||
for group in plans[1].groups:
|
||||
assert len(group) == 2
|
||||
|
||||
# Step 2 (TP): 2 tensors → 1 (single group of 2)
|
||||
assert len(plans[2].groups) == 1
|
||||
assert len(plans[2].groups[0]) == 2
|
||||
|
||||
def test_replicated_axis_raises(self) -> None:
|
||||
"""A world_rank missing a sharded axis raises ValueError."""
|
||||
dim_specs = parse_dims("s(cp) h(tp)")
|
||||
parallel_infos = [
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
|
||||
},
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
|
||||
# missing TP — replicated
|
||||
},
|
||||
]
|
||||
with pytest.raises(ValueError, match="missing parallel_info"):
|
||||
compute_unshard_plan(dim_specs, parallel_infos)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user