From b7af58b9afab832a9979a789f7c5085817d6c4be Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 25 Feb 2026 09:48:43 +0800 Subject: [PATCH] Support replication axis in dump comparator (#19282) --- .../comparator/aligner/unshard/executor.py | 68 +++++- .../comparator/aligner/unshard/planner.py | 42 ++-- .../comparator/aligner/unshard/types.py | 13 +- .../debug_utils/comparator/output_types.py | 49 +++- .../srt/debug_utils/comparator/pipeline.py | 31 ++- .../comparator/aligner/test_reorder.py | 2 +- .../aligner/unshard/test_execute.py | 231 +++++++++++++++++- .../comparator/aligner/unshard/test_plan.py | 162 +++++++++++- .../tensor_comparison/test_types.py | 67 +++++ .../debug_utils/comparator/test_entrypoint.py | 175 +++++++++++++ 10 files changed, 778 insertions(+), 62 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py index f48c770c9..c89f73d50 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py @@ -2,30 +2,80 @@ import torch from sglang.srt.debug_utils.comparator.aligner.unshard.types import ( ConcatParams, + PickParams, UnshardParams, UnshardPlan, ) +from sglang.srt.debug_utils.comparator.dims import ParallelAxis +from sglang.srt.debug_utils.comparator.output_types import ( + AlignWarning, + ReplicatedMismatchWarning, +) def execute_unshard_plan( plan: UnshardPlan, tensors: list[torch.Tensor], -) -> list[torch.Tensor]: +) -> tuple[list[torch.Tensor], list[AlignWarning]]: + all_warnings: list[AlignWarning] = [] result: list[torch.Tensor] = [] - for group in plan.groups: + + for group_idx, group in enumerate(plan.groups): group_tensors = [tensors[i] for i in group] - result.append(_apply_unshard(plan.params, group_tensors)) - return result + tensor, warnings = _apply_unshard( + plan.params, + group_tensors, + axis=plan.axis, + group_index=group_idx, + ) + result.append(tensor) + all_warnings.extend(warnings) + + return result, all_warnings def _apply_unshard( - params: UnshardParams, ordered_tensors: list[torch.Tensor] -) -> torch.Tensor: + params: UnshardParams, + ordered_tensors: list[torch.Tensor], + *, + axis: ParallelAxis, + group_index: int, +) -> tuple[torch.Tensor, list[AlignWarning]]: + if isinstance(params, PickParams): + warnings = _verify_replicated_group( + ordered_tensors, + axis=axis, + group_index=group_index, + ) + return ordered_tensors[0], warnings + if isinstance(params, ConcatParams): - return _unshard_concat(ordered_tensors, dim=params.dim) + return torch.cat(ordered_tensors, dim=params.dim), [] + # Phase 2: ReduceSumParams, CpZigzagParams raise ValueError(f"Unsupported unshard operation: {type(params).__name__}") -def _unshard_concat(tensors: list[torch.Tensor], dim: int) -> torch.Tensor: - return torch.cat(tensors, dim=dim) +def _verify_replicated_group( + ordered_tensors: list[torch.Tensor], + *, + axis: ParallelAxis, + group_index: int, +) -> list[ReplicatedMismatchWarning]: + warnings: list[ReplicatedMismatchWarning] = [] + baseline = ordered_tensors[0] + + for i in range(1, len(ordered_tensors)): + other = ordered_tensors[i] + if not torch.allclose(baseline, other, atol=1e-6): + warnings.append( + ReplicatedMismatchWarning( + axis=axis.value, + group_index=group_index, + differing_index=i, + baseline_index=0, + max_abs_diff=(baseline - other).abs().max().item(), + ) + ) + + return warnings diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py index b5404744e..f8badb588 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py @@ -4,6 +4,7 @@ from typing import NamedTuple from sglang.srt.debug_utils.comparator.aligner.unshard.types import ( AxisInfo, ConcatParams, + PickParams, UnshardParams, UnshardPlan, ) @@ -32,31 +33,38 @@ def compute_unshard_plan( for dim_idx, spec in enumerate(dim_specs) if spec.parallel is not None } - if not sharded_axis_infos: + sharded_axes: set[ParallelAxis] = set(sharded_axis_infos) + + all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info} + replicated_axes: set[ParallelAxis] = all_axes - sharded_axes + + if not sharded_axes and not replicated_axes: return [] - _validate(sharded_axes=set(sharded_axis_infos), parallel_infos=parallel_infos) + _validate( + axes_to_validate=sharded_axes | replicated_axes, + parallel_infos=parallel_infos, + ) current_coords: _CoordsList = [ - {axis: info[axis].axis_rank for axis in sharded_axis_infos} + {axis: info[axis].axis_rank for axis in sharded_axes | replicated_axes} for info in parallel_infos ] + axis_and_params: list[tuple[ParallelAxis, UnshardParams]] = [ + (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value) + ] + [ + (axis, _resolve_unshard_params(spec=spec, dim_index=dim_index)) + for axis, (dim_index, spec) in sharded_axis_infos.items() + ] + plans: list[UnshardPlan] = [] - for axis, (dim_index, spec) in sharded_axis_infos.items(): + for axis, params in axis_and_params: result = _group_and_project( current_coords=current_coords, target_axis=axis, ) - - plans.append( - UnshardPlan( - axis=axis, - params=_resolve_unshard_params(spec=spec, dim_index=dim_index), - groups=result.groups, - ) - ) - + plans.append(UnshardPlan(axis=axis, params=params, groups=result.groups)) current_coords = result.projected_coords return plans @@ -64,18 +72,18 @@ def compute_unshard_plan( def _validate( *, - sharded_axes: set[ParallelAxis], + axes_to_validate: set[ParallelAxis], parallel_infos: list[dict[ParallelAxis, AxisInfo]], ) -> None: - """Check that every rank has all sharded axes, sizes are consistent, and ranks are complete.""" + """Check that every rank has all axes, sizes are consistent, and ranks are complete.""" axis_sizes: dict[ParallelAxis, int] = {} for world_rank, parallel_info in enumerate(parallel_infos): - for axis in sharded_axes: + for axis in axes_to_validate: if axis not in parallel_info: raise ValueError( f"world_rank={world_rank} missing parallel_info for " - f"sharded axis {axis.value!r}" + f"axis {axis.value!r}" ) axis_info = parallel_info[axis] diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py b/python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py index e0d837056..a8b860f26 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py @@ -1,6 +1,8 @@ from __future__ import annotations -from typing import Literal +from typing import Annotated, Literal, Union + +from pydantic import Field from sglang.srt.debug_utils.comparator.dims import ParallelAxis from sglang.srt.debug_utils.comparator.utils import _FrozenBase @@ -16,7 +18,14 @@ class ConcatParams(_FrozenBase): dim: int -UnshardParams = ConcatParams +class PickParams(_FrozenBase): + op: Literal["pick"] = "pick" + + +UnshardParams = Annotated[ + Union[ConcatParams, PickParams], + Field(discriminator="op"), +] class UnshardPlan(_FrozenBase): diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index 28f325ab2..300954561 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -1,7 +1,7 @@ from abc import abstractmethod from typing import Annotated, Literal, Union -from pydantic import Discriminator, TypeAdapter +from pydantic import Discriminator, Field, TypeAdapter from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import ( format_comparison, @@ -12,9 +12,38 @@ from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( from sglang.srt.debug_utils.comparator.utils import _StrictBase +class ReplicatedMismatchWarning(_StrictBase): + kind: Literal["replicated_mismatch"] = "replicated_mismatch" + axis: str + group_index: int + differing_index: int + baseline_index: int + max_abs_diff: float + + def to_text(self) -> str: + return ( + f"Replicated along {self.axis}: group {self.group_index}, " + f"index {self.differing_index} differs from {self.baseline_index} " + f"(max_abs_diff={self.max_abs_diff:.6e})" + ) + + +AlignWarning = ( + ReplicatedMismatchWarning # future: Annotated[Union[...], Discriminator("kind")] +) + + class _OutputRecord(_StrictBase): + align_warnings: list[AlignWarning] = Field(default_factory=list) + @abstractmethod - def to_text(self) -> str: ... + def _format_body(self) -> str: ... + + def to_text(self) -> str: + body = self._format_body() + if self.align_warnings: + body += "\n" + "\n".join(f" ⚠ {w.to_text()}" for w in self.align_warnings) + return body class ConfigRecord(_OutputRecord): @@ -25,7 +54,7 @@ class ConfigRecord(_OutputRecord): start_step: int end_step: int - def to_text(self) -> str: + def _format_body(self) -> str: return ( f"Config: baseline={self.baseline_path} target={self.target_path}\n" f"diff_threshold={self.diff_threshold} " @@ -39,10 +68,12 @@ class SkipRecord(_OutputRecord): reason: str @property - def category(self): + def category(self) -> str: + if self.align_warnings: + return "failed" return "skipped" - def to_text(self) -> str: + def _format_body(self) -> str: return f"Skip: {self.name} ({self.reason})" @@ -50,10 +81,12 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord): type: Literal["comparison"] = "comparison" @property - def category(self): + def category(self) -> str: + if self.align_warnings: + return "failed" return "passed" if self.diff is not None and self.diff.passed else "failed" - def to_text(self) -> str: + def _format_body(self) -> str: return format_comparison(self) @@ -64,7 +97,7 @@ class SummaryRecord(_OutputRecord): failed: int skipped: int - def to_text(self) -> str: + def _format_body(self) -> str: return ( f"Summary: {self.passed} passed, {self.failed} failed, " f"{self.skipped} skipped (total {self.total})" diff --git a/python/sglang/srt/debug_utils/comparator/pipeline.py b/python/sglang/srt/debug_utils/comparator/pipeline.py index 5bb3c1736..becf7f595 100644 --- a/python/sglang/srt/debug_utils/comparator/pipeline.py +++ b/python/sglang/srt/debug_utils/comparator/pipeline.py @@ -20,6 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unshard.planner import ( from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan from sglang.srt.debug_utils.comparator.dims import parse_dims from sglang.srt.debug_utils.comparator.output_types import ( + AlignWarning, ComparisonRecord, SkipRecord, ) @@ -50,12 +51,13 @@ def process_tensor_group( t_extracted = _extract_tensors(t_tensors) del b_tensors, t_tensors - b_tensor = _execute_plans(b_extracted, b_plans) - t_tensor = _execute_plans(t_extracted, t_plans) + b_tensor, b_warns = _execute_plans(b_extracted, b_plans) + t_tensor, t_warns = _execute_plans(t_extracted, t_plans) + all_warnings: list[AlignWarning] = b_warns + t_warns if b_tensor is None or t_tensor is None: reason = "baseline_load_failed" if b_tensor is None else "target_load_failed" - return SkipRecord(name=name, reason=reason) + return SkipRecord(name=name, reason=reason, align_warnings=all_warnings) info = compare_tensors( x_baseline=b_tensor, @@ -64,7 +66,7 @@ def process_tensor_group( diff_threshold=diff_threshold, ) - return ComparisonRecord(**info.model_dump()) + return ComparisonRecord(**info.model_dump(), align_warnings=all_warnings) def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]: @@ -112,27 +114,32 @@ def _extract_tensors( def _execute_plans( tensors: list[torch.Tensor], plans: list[Plan], -) -> Optional[torch.Tensor]: +) -> tuple[Optional[torch.Tensor], list[AlignWarning]]: if not tensors: - return None + return None, [] if not plans: if len(tensors) != 1: - return None - return tensors[0] + return None, [] + return tensors[0], [] + warnings: list[AlignWarning] = [] current = tensors for plan in plans: - current = _execute_plan(current, plan) + current, new_warnings = _execute_plan(current, plan) + warnings.extend(new_warnings) assert len(current) == 1 - return current[0] + return current[0], warnings -def _execute_plan(tensors, plan): +def _execute_plan( + tensors: list[torch.Tensor], + plan: Plan, +) -> tuple[list[torch.Tensor], list[AlignWarning]]: if isinstance(plan, UnshardPlan): return execute_unshard_plan(plan, tensors) elif isinstance(plan, ReorderPlan): - return execute_reorder_plan(plan, tensors) + return execute_reorder_plan(plan, tensors), [] else: raise NotImplementedError(f"Unknown {plan=}") diff --git a/test/registered/debug_utils/comparator/aligner/test_reorder.py b/test/registered/debug_utils/comparator/aligner/test_reorder.py index 1b4b2f02c..a78c5e3d2 100644 --- a/test/registered/debug_utils/comparator/aligner/test_reorder.py +++ b/test/registered/debug_utils/comparator/aligner/test_reorder.py @@ -148,7 +148,7 @@ class TestCpZigzagTpE2E: if isinstance(plan, ReorderPlan): current = execute_reorder_plan(plan, current) else: - current = execute_unshard_plan(plan, current) + current, _ = execute_unshard_plan(plan, current) assert len(current) == 1 assert torch.allclose(current[0], full_tensor) diff --git a/test/registered/debug_utils/comparator/aligner/unshard/test_execute.py b/test/registered/debug_utils/comparator/aligner/unshard/test_execute.py index 083c7b4e3..e862ddbad 100644 --- a/test/registered/debug_utils/comparator/aligner/unshard/test_execute.py +++ b/test/registered/debug_utils/comparator/aligner/unshard/test_execute.py @@ -5,12 +5,16 @@ import torch from sglang.srt.debug_utils.comparator.aligner.unshard.executor import ( _apply_unshard, + _verify_replicated_group, execute_unshard_plan, ) from sglang.srt.debug_utils.comparator.aligner.unshard.planner import ( compute_unshard_plan, ) -from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo +from sglang.srt.debug_utils.comparator.aligner.unshard.types import ( + AxisInfo, + PickParams, +) from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims from sglang.test.ci.ci_register import register_cpu_ci @@ -29,9 +33,10 @@ class TestExecuteUnshardPlan: plans = compute_unshard_plan(dim_specs, parallel_infos) assert len(plans) == 1 - result = execute_unshard_plan(plans[0], shards) + result, warnings = execute_unshard_plan(plans[0], shards) assert len(result) == 1 assert torch.allclose(result[0], full_tensor) + assert warnings == [] def test_scrambled_world_ranks_correct_result(self) -> None: full_tensor = torch.randn(4, 8) @@ -54,9 +59,10 @@ class TestExecuteUnshardPlan: shards[1], # world_rank=3, axis_rank=1 ] - result = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank) + result, warnings = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank) assert len(result) == 1 assert torch.allclose(result[0], full_tensor) + assert warnings == [] def test_single_step_reduces_tensor_count(self) -> None: """8 tensors with 2 groups of 4 produce 2 output tensors.""" @@ -85,10 +91,10 @@ class TestExecuteUnshardPlan: for tp_rank in range(4): tensors.append(source[tp_rank]) - intermediate = execute_unshard_plan(plans[0], tensors) + intermediate, _ = execute_unshard_plan(plans[0], tensors) assert len(intermediate) == 4 - final = execute_unshard_plan(plans[1], intermediate) + final, _ = execute_unshard_plan(plans[1], intermediate) assert len(final) == 1 def test_cp_tp_concat(self) -> None: @@ -116,7 +122,7 @@ class TestExecuteUnshardPlan: current = tensors for plan in plans: - current = execute_unshard_plan(plan, current) + current, _ = execute_unshard_plan(plan, current) assert len(current) == 1 assert torch.allclose(current[0], full_tensor) @@ -157,7 +163,7 @@ class TestExecuteUnshardPlan: current = tensors for plan in plans: - current = execute_unshard_plan(plan, current) + current, _ = execute_unshard_plan(plan, current) assert len(current) == 1 assert torch.allclose(current[0], full_tensor) @@ -169,7 +175,12 @@ class TestExecuteUnshardPlan: pass with pytest.raises(ValueError, match="Unsupported unshard"): - _apply_unshard(_FakeParams(), [torch.randn(2, 2)]) + _apply_unshard( + _FakeParams(), + [torch.randn(2, 2)], + axis=ParallelAxis.TP, + group_index=0, + ) def test_cp_tp_ep_three_axis_concat(self) -> None: """CP=2 + TP=2 + EP=2: three-step unshard reconstructs original tensor.""" @@ -205,7 +216,7 @@ class TestExecuteUnshardPlan: current = tensors for plan in plans: - current = execute_unshard_plan(plan, current) + current, _ = execute_unshard_plan(plan, current) assert len(current) == 1 assert torch.allclose(current[0], full_tensor) @@ -253,11 +264,211 @@ class TestExecuteUnshardPlan: current = tensors for plan in plans: - current = execute_unshard_plan(plan, current) + current, _ = execute_unshard_plan(plan, current) assert len(current) == 1 assert torch.allclose(current[0], full_tensor) +class TestPickOperation: + def test_pick_single_group(self) -> None: + """PickParams picks the first tensor from a single group.""" + tensor = torch.randn(4, 8) + dim_specs = parse_dims("h d") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)}, + {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)}, + ] + + plans = compute_unshard_plan(dim_specs, parallel_infos) + assert len(plans) == 1 + assert isinstance(plans[0].params, PickParams) + + result, warnings = execute_unshard_plan(plans[0], [tensor, tensor.clone()]) + assert len(result) == 1 + assert torch.allclose(result[0], tensor) + assert warnings == [] + + def test_pick_multiple_groups(self) -> None: + """PickParams with multiple groups picks one from each.""" + dim_specs = parse_dims("h(tp)") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ + { + 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), + 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=1, axis_size=2), + }, + ] + + plans = compute_unshard_plan(dim_specs, parallel_infos) + pick_plans = [p for p in plans if isinstance(p.params, PickParams)] + assert len(pick_plans) == 1 + assert pick_plans[0].axis == ParallelAxis.CP + + tensor = torch.randn(4) + tensors = [tensor.clone() for _ in range(4)] + + result, warnings = execute_unshard_plan(pick_plans[0], tensors) + assert len(result) == 2 + assert warnings == [] + + def test_replicated_tp_sharded_cp_e2e(self) -> None: + """CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip.""" + 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: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(2): + for tp_rank in range(2): + tensors.append(cp_chunks[cp_rank].clone()) + 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) d") + 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_fully_replicated_e2e(self) -> None: + """CP2 TP2, dims='b h d': fully replicated → 2 pick steps → 1 tensor.""" + torch.manual_seed(42) + full_tensor = torch.randn(4, 8, 16) + + tensors: list[torch.Tensor] = [] + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(2): + for tp_rank in range(2): + tensors.append(full_tensor.clone()) + 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 h d") + plans = compute_unshard_plan(dim_specs, parallel_infos) + assert len(plans) == 2 + assert all(isinstance(p.params, PickParams) for p in plans) + + current = tensors + for plan in plans: + current, _ = execute_unshard_plan(plan, current) + + assert len(current) == 1 + assert torch.allclose(current[0], full_tensor) + + +class TestVerifyReplicatedGroup: + def test_warns_on_mismatch(self) -> None: + """_verify_replicated_group produces warning when replicas differ.""" + tensor_a = torch.ones(4) + tensor_b = torch.ones(4) + 0.1 + + warnings = _verify_replicated_group( + [tensor_a, tensor_b], + axis=ParallelAxis.TP, + group_index=0, + ) + assert len(warnings) == 1 + assert warnings[0].axis == "tp" + assert warnings[0].group_index == 0 + assert warnings[0].differing_index == 1 + assert warnings[0].baseline_index == 0 + assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5) + + def test_no_warn_when_identical(self) -> None: + """_verify_replicated_group produces no warning for identical replicas.""" + tensor = torch.randn(4, 8) + + warnings = _verify_replicated_group( + [tensor, tensor.clone()], + axis=ParallelAxis.TP, + group_index=0, + ) + assert warnings == [] + + def test_multiple_mismatches(self) -> None: + """_verify_replicated_group reports each differing replica.""" + baseline = torch.zeros(4) + other_a = torch.ones(4) + other_b = torch.ones(4) * 2 + + warnings = _verify_replicated_group( + [baseline, other_a, other_b], + axis=ParallelAxis.CP, + group_index=1, + ) + assert len(warnings) == 2 + assert warnings[0].differing_index == 1 + assert warnings[1].differing_index == 2 + assert warnings[1].max_abs_diff == pytest.approx(2.0, abs=1e-5) + + def test_execute_returns_warnings(self) -> None: + """execute_unshard_plan returns warnings for replicated mismatch.""" + dim_specs = parse_dims("h d") + parallel_infos = [ + {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)}, + {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)}, + ] + plans = compute_unshard_plan(dim_specs, parallel_infos) + + tensor_a = torch.zeros(4) + tensor_b = torch.ones(4) + + result, warnings = execute_unshard_plan(plans[0], [tensor_a, tensor_b]) + assert len(result) == 1 + assert len(warnings) == 1 + assert torch.allclose(result[0], tensor_a) + + def test_atol_boundary_within(self) -> None: + """Difference exactly at atol (1e-6) → torch.allclose passes → no warning.""" + baseline = torch.zeros(4) + other = torch.full((4,), 1e-6) + + warnings = _verify_replicated_group( + [baseline, other], + axis=ParallelAxis.TP, + group_index=0, + ) + assert warnings == [] + + def test_atol_boundary_exceeded(self) -> None: + """Difference just above atol (1e-6 + 1e-9) → torch.allclose fails → warning.""" + baseline = torch.zeros(4) + other = torch.full((4,), 1e-6 + 1e-9) + + warnings = _verify_replicated_group( + [baseline, other], + axis=ParallelAxis.TP, + group_index=0, + ) + assert len(warnings) == 1 + assert warnings[0].differing_index == 1 + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/aligner/unshard/test_plan.py b/test/registered/debug_utils/comparator/aligner/unshard/test_plan.py index d6dfaf727..0f5a6577f 100644 --- a/test/registered/debug_utils/comparator/aligner/unshard/test_plan.py +++ b/test/registered/debug_utils/comparator/aligner/unshard/test_plan.py @@ -5,7 +5,11 @@ import pytest from sglang.srt.debug_utils.comparator.aligner.unshard.planner import ( compute_unshard_plan, ) -from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo +from sglang.srt.debug_utils.comparator.aligner.unshard.types import ( + AxisInfo, + ConcatParams, + PickParams, +) from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims from sglang.test.ci.ci_register import register_cpu_ci @@ -228,7 +232,7 @@ class TestComputeUnshardPlan: assert len(plans[2].groups) == 1 assert len(plans[2].groups[0]) == 2 - def test_replicated_axis_raises(self) -> None: + def test_sharded_axis_missing_from_rank_raises(self) -> None: """A world_rank missing a sharded axis raises ValueError.""" dim_specs = parse_dims("s(cp) h(tp)") parallel_infos = [ @@ -238,7 +242,159 @@ class TestComputeUnshardPlan: }, { ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2), - # missing TP — replicated + # missing TP — sharded axis absent from rank + }, + ] + with pytest.raises(ValueError, match="missing parallel_info"): + compute_unshard_plan(dim_specs, parallel_infos) + + +class TestReplicatedAxes: + def test_replicated_tp_with_sharded_cp(self) -> None: + """CP2 TP2, dims='b s(cp) d' → PickPlan(TP) + ConcatPlan(CP).""" + dim_specs = parse_dims("b s(cp) d") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ + { + 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), + }, + ] + plans = compute_unshard_plan(dim_specs, parallel_infos) + + assert len(plans) == 2 + assert plans[0].axis == ParallelAxis.TP + assert isinstance(plans[0].params, PickParams) + assert len(plans[0].groups) == 2 + for group in plans[0].groups: + assert len(group) == 2 + + assert plans[1].axis == ParallelAxis.CP + assert isinstance(plans[1].params, ConcatParams) + assert plans[1].params.dim == 1 + + def test_fully_replicated(self) -> None: + """CP2 TP2, dims='b h d' → PickPlan(CP) + PickPlan(TP).""" + dim_specs = parse_dims("b h d") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ + { + 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), + }, + ] + plans = compute_unshard_plan(dim_specs, parallel_infos) + + assert len(plans) == 2 + assert all(isinstance(p.params, PickParams) for p in plans) + axes = {p.axis for p in plans} + assert axes == {ParallelAxis.CP, ParallelAxis.TP} + + def test_multiple_replicated_one_sharded(self) -> None: + """CP2 TP2 EP2, dims='h(tp)' → PickPlan(CP) + PickPlan(EP) + ConcatPlan(TP).""" + dim_specs = parse_dims("h(tp)") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [] + for cp_rank in range(2): + for ep_rank in range(2): + for tp_rank in range(2): + parallel_infos.append( + { + ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2), + ParallelAxis.EP: AxisInfo(axis_rank=ep_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 + pick_plans = [p for p in plans if isinstance(p.params, PickParams)] + concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)] + assert len(pick_plans) == 2 + assert len(concat_plans) == 1 + assert concat_plans[0].axis == ParallelAxis.TP + + replicated_axes = {p.axis for p in pick_plans} + assert replicated_axes == {ParallelAxis.CP, ParallelAxis.EP} + + def test_replicated_scrambled_ranks(self) -> None: + """Scrambled world_rank order with replicated axis.""" + dim_specs = parse_dims("h(tp)") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ + { + ParallelAxis.CP: AxisInfo(axis_rank=1, 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=0, 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), + }, + ] + plans = compute_unshard_plan(dim_specs, parallel_infos) + + assert len(plans) == 2 + assert plans[0].axis == ParallelAxis.CP + assert isinstance(plans[0].params, PickParams) + assert plans[1].axis == ParallelAxis.TP + assert isinstance(plans[1].params, ConcatParams) + + def test_replicated_axis_inconsistent_size_raises(self) -> None: + """Replicated axis with inconsistent sizes raises ValueError.""" + dim_specs = parse_dims("h(tp)") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ + { + 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=4), + ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2), + }, + ] + with pytest.raises(ValueError, match="Inconsistent axis_size"): + compute_unshard_plan(dim_specs, parallel_infos) + + def test_replicated_axis_missing_from_rank_raises(self) -> None: + """A rank missing a replicated axis that other ranks have raises ValueError.""" + dim_specs = parse_dims("h(tp)") + parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [ + { + ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2), + ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2), + }, + { + # missing CP — replicated axis absent from this rank + ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2), }, ] with pytest.raises(ValueError, match="missing parallel_info"): diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_types.py b/test/registered/debug_utils/comparator/tensor_comparison/test_types.py index cffd08d28..7fb446af7 100644 --- a/test/registered/debug_utils/comparator/tensor_comparison/test_types.py +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_types.py @@ -6,6 +6,7 @@ import pytest from sglang.srt.debug_utils.comparator.output_types import ( ComparisonRecord, ConfigRecord, + ReplicatedMismatchWarning, SkipRecord, SummaryRecord, parse_record_json, @@ -118,5 +119,71 @@ class TestRecordTypes: assert restored == record +def _make_warning(**overrides) -> ReplicatedMismatchWarning: + defaults: dict = dict( + axis="tp", + group_index=0, + differing_index=1, + baseline_index=0, + max_abs_diff=0.1, + ) + defaults.update(overrides) + return ReplicatedMismatchWarning(**defaults) + + +class TestAlignWarnings: + def test_comparison_record_failed_when_diff_passed_but_warnings(self): + """ComparisonRecord with diff.passed=True but align_warnings → category=='failed'.""" + record = ComparisonRecord( + name="hidden", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(passed=True), + align_warnings=[_make_warning()], + ) + assert record.category == "failed" + + def test_skip_record_failed_when_warnings(self): + """SkipRecord with align_warnings → category=='failed' instead of 'skipped'.""" + record = SkipRecord( + name="x", + reason="no_baseline", + align_warnings=[_make_warning()], + ) + assert record.category == "failed" + + def test_align_warnings_json_round_trip(self): + """align_warnings survive model_dump_json → parse_record_json round-trip.""" + warning = _make_warning( + axis="cp", + group_index=2, + differing_index=3, + baseline_index=0, + max_abs_diff=0.42, + ) + record = ComparisonRecord( + name="mlp", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + align_warnings=[warning], + ) + + restored = parse_record_json(record.model_dump_json()) + assert isinstance(restored, ComparisonRecord) + assert len(restored.align_warnings) == 1 + + restored_warning = restored.align_warnings[0] + assert restored_warning.axis == "cp" + assert restored_warning.group_index == 2 + assert restored_warning.differing_index == 3 + assert restored_warning.baseline_index == 0 + assert restored_warning.max_abs_diff == pytest.approx(0.42) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 03c78a4b5..eb50eaa4d 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -878,6 +878,138 @@ class TestEntrypointGroupingLogical: assert comp.name == "hidden" +class TestEntrypointReplicatedAxis: + """Test replicated-axis scenarios through the full entrypoint pipeline.""" + + def test_replicated_axis_identical_replicas_passed(self, tmp_path, capsys): + """CP2 TP2, TP replicated and identical → passed, no align_warnings.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8, 6) + full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001 + + 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_replicated_tp_sharded_cp_dumps( + side_dir, + full_tensor=full_tensor, + name="attn_out", + cp_size=2, + tp_size=2, + seq_dim=1, + dims_str="b s(cp) d", + ) + + 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.align_warnings == [] + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.passed == 1 + + def test_replicated_mismatch_fails(self, tmp_path, capsys): + """CP2 TP2, TP replicas differ (> atol) → failed with align_warnings.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8, 6) + full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001 + + 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_replicated_tp_sharded_cp_dumps( + side_dir, + full_tensor=full_tensor, + name="attn_out", + cp_size=2, + tp_size=2, + seq_dim=1, + dims_str="b s(cp) d", + tp_noise=0.5, + ) + + args = _make_args( + baseline_dir / _FIXED_EXP_NAME, + target_dir / _FIXED_EXP_NAME, + diff_threshold=0.01, + ) + + records = _run_and_parse(args, capsys) + comparisons = _get_comparisons(records) + assert len(comparisons) == 1 + assert comparisons[0].category == "failed" + assert len(comparisons[0].align_warnings) > 0 + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.failed == 1 + + def test_summary_counts_failed_from_align_warnings_only(self, tmp_path, capsys): + """Diff itself passes but TP replicas differ → summary.failed=1 from align_warnings.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8, 6) + full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001 + + baseline_dir = tmp_path / "baseline" + target_dir = tmp_path / "target" + + _create_replicated_tp_sharded_cp_dumps( + baseline_dir, + full_tensor=full_baseline, + name="attn_out", + cp_size=2, + tp_size=2, + seq_dim=1, + dims_str="b s(cp) d", + tp_noise=0.5, + ) + _create_replicated_tp_sharded_cp_dumps( + target_dir, + full_tensor=full_target, + name="attn_out", + cp_size=2, + tp_size=2, + seq_dim=1, + dims_str="b s(cp) d", + tp_noise=0.5, + ) + + args = _make_args( + baseline_dir / _FIXED_EXP_NAME, + target_dir / _FIXED_EXP_NAME, + diff_threshold=0.5, + ) + + records = _run_and_parse(args, capsys) + comparisons = _get_comparisons(records) + assert len(comparisons) == 1 + + comp = comparisons[0] + assert comp.diff is not None + assert comp.diff.passed + assert len(comp.align_warnings) > 0 + assert comp.category == "failed" + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.failed == 1 + assert summary.passed == 0 + + # --------------------------- Assertion helpers ------------------- @@ -1138,6 +1270,49 @@ def _create_cp_zigzag_tp_sharded_dumps( return directory / _FIXED_EXP_NAME +def _create_replicated_tp_sharded_cp_dumps( + directory: Path, + *, + full_tensor: torch.Tensor, + name: str, + cp_size: int, + tp_size: int, + seq_dim: int, + dims_str: str, + tp_noise: float = 0.0, +) -> Path: + """Create CP-sharded + TP-replicated dump files from a full tensor. + + CP direction: chunks along seq_dim (sharded). + TP direction: clones (replicated), with optional noise to simulate mismatch. + """ + cp_chunks: list[torch.Tensor] = list(full_tensor.chunk(cp_size, dim=seq_dim)) + + rank: int = 0 + for cp_rank in range(cp_size): + for tp_rank in range(tp_size): + shard = cp_chunks[cp_rank].clone() + if tp_noise > 0 and tp_rank > 0: + shard = shard + torch.randn_like(shard) * tp_noise + + _create_rank_dump( + directory, + rank=rank, + name=name, + tensor=shard, + dims=dims_str, + parallel_info={ + "cp_rank": cp_rank, + "cp_size": cp_size, + "tp_rank": tp_rank, + "tp_size": tp_size, + }, + ) + rank += 1 + + return directory / _FIXED_EXP_NAME + + def _create_tp_sharded_dumps( directory: Path, *,