diff --git a/python/sglang/srt/debug_utils/comparator/pipeline.py b/python/sglang/srt/debug_utils/comparator/pipeline.py index 20715a09a..6555b83ba 100644 --- a/python/sglang/srt/debug_utils/comparator/pipeline.py +++ b/python/sglang/srt/debug_utils/comparator/pipeline.py @@ -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] diff --git a/python/sglang/srt/debug_utils/comparator/unshard/executor.py b/python/sglang/srt/debug_utils/comparator/unshard/executor.py index 2842ccd26..7ced0660e 100644 --- a/python/sglang/srt/debug_utils/comparator/unshard/executor.py +++ b/python/sglang/srt/debug_utils/comparator/unshard/executor.py @@ -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 diff --git a/python/sglang/srt/debug_utils/comparator/unshard/planner.py b/python/sglang/srt/debug_utils/comparator/unshard/planner.py index 727257fe6..a077f757d 100644 --- a/python/sglang/srt/debug_utils/comparator/unshard/planner.py +++ b/python/sglang/srt/debug_utils/comparator/unshard/planner.py @@ -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: diff --git a/python/sglang/srt/debug_utils/comparator/unshard/types.py b/python/sglang/srt/debug_utils/comparator/unshard/types.py index 114f381f9..ec1b4365b 100644 --- a/python/sglang/srt/debug_utils/comparator/unshard/types.py +++ b/python/sglang/srt/debug_utils/comparator/unshard/types.py @@ -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, diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 8410049c6..768e46e59 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -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, *, diff --git a/test/registered/debug_utils/comparator/unshard/test_execute.py b/test/registered/debug_utils/comparator/unshard/test_execute.py index 736314983..6da2c19d7 100644 --- a/test/registered/debug_utils/comparator/unshard/test_execute.py +++ b/test/registered/debug_utils/comparator/unshard/test_execute.py @@ -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__": diff --git a/test/registered/debug_utils/comparator/unshard/test_plan.py b/test/registered/debug_utils/comparator/unshard/test_plan.py index 011cebd56..60448e792 100644 --- a/test/registered/debug_utils/comparator/unshard/test_plan.py +++ b/test/registered/debug_utils/comparator/unshard/test_plan.py @@ -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)