diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py index e9781711e..c4a39b814 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional import torch @@ -21,9 +21,11 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import ( execute_token_aligner, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import ( + UnsharderResult, execute_unsharder_plan, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan +from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult from sglang.srt.debug_utils.comparator.utils import Pair @@ -31,6 +33,7 @@ from sglang.srt.debug_utils.comparator.utils import Pair class AlignerResult: tensors: Optional[Pair[torch.Tensor]] failed_side_xy: Optional[str] # "x" or "y"; None if success + replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list) def execute_aligner_plan( @@ -39,18 +42,26 @@ def execute_aligner_plan( plan: AlignerPlan, ) -> AlignerResult: """Execute unified unshard/reorder + token-align.""" + all_checks: list[ReplicatedCheckResult] = [] # Per-side: unshard + reorder -> dict[step, tensor] - step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans( + step_tensors_x, checks_x = _execute_step_plans( tensors=tensors_pair.x, step_plans=plan.per_step_plans.x ) - step_tensors_y: dict[int, torch.Tensor] = _execute_step_plans( + all_checks.extend(checks_x) + + step_tensors_y, checks_y = _execute_step_plans( tensors=tensors_pair.y, step_plans=plan.per_step_plans.y ) + all_checks.extend(checks_y) if not step_tensors_x or not step_tensors_y: failed_side_xy: str = "x" if not step_tensors_x else "y" - return AlignerResult(tensors=None, failed_side_xy=failed_side_xy) + return AlignerResult( + tensors=None, + failed_side_xy=failed_side_xy, + replicated_checks=all_checks, + ) # Cross-side: token alignment (or direct extraction for single-step) if plan.token_aligner_plan is not None: @@ -72,55 +83,64 @@ def execute_aligner_plan( y=execute_axis_aligner_plan(tensor=combined.y, plan=aligner_plan, side="y"), ) - return AlignerResult(tensors=combined, failed_side_xy=None) + return AlignerResult( + tensors=combined, + failed_side_xy=None, + replicated_checks=all_checks, + ) def _execute_step_plans( tensors: list[torch.Tensor], step_plans: list[AlignerPerStepPlan], -) -> dict[int, torch.Tensor]: +) -> tuple[dict[int, torch.Tensor], list[ReplicatedCheckResult]]: result: dict[int, torch.Tensor] = {} + all_checks: list[ReplicatedCheckResult] = [] for step_plan in step_plans: step_tensors: list[torch.Tensor] = [ tensors[i] for i in step_plan.input_object_indices ] - tensor: Optional[torch.Tensor] = execute_sub_plans( + tensor, checks = execute_sub_plans( tensors=step_tensors, plans=step_plan.sub_plans ) + all_checks.extend(checks) if tensor is not None: result[step_plan.step] = tensor - return result + return result, all_checks def execute_sub_plans( tensors: list[torch.Tensor], plans: list[AlignerPerStepSubPlan], -) -> Optional[torch.Tensor]: +) -> tuple[Optional[torch.Tensor], list[ReplicatedCheckResult]]: if not tensors: - return None + return None, [] if not plans: if len(tensors) != 1: - return None - return tensors[0] + return None, [] + return tensors[0], [] - current = tensors + current: list[torch.Tensor] = tensors + all_checks: list[ReplicatedCheckResult] = [] for plan in plans: - current = execute_sub_plan(tensors=current, plan=plan) + current, checks = execute_sub_plan(tensors=current, plan=plan) + all_checks.extend(checks) assert len(current) == 1 - return current[0] + return current[0], all_checks def execute_sub_plan( tensors: list[torch.Tensor], plan: AlignerPerStepSubPlan, -) -> list[torch.Tensor]: +) -> tuple[list[torch.Tensor], list[ReplicatedCheckResult]]: if isinstance(plan, UnsharderPlan): - return execute_unsharder_plan(plan, tensors) + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, tensors) + return unsharder_result.tensors, unsharder_result.replicated_checks elif isinstance(plan, ReordererPlan): - return execute_reorderer_plan(plan, tensors) + return execute_reorderer_plan(plan, tensors), [] else: raise NotImplementedError(f"Unknown {plan=}") diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py index a6f94f919..51267e6a6 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py @@ -236,7 +236,7 @@ def _load_and_align_aux_tensor( dim_names: list[str] = resolve_dim_names(dims_str) tensors = [apply_dim_names(t, dim_names) for t in tensors] - result = execute_sub_plans(tensors=tensors, plans=sub_plans) + result, _replicated_checks = execute_sub_plans(tensors=tensors, plans=sub_plans) assert result is not None return result.rename(None) # strip named dims before returning to plugin diff --git a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py index c52a99153..a87cfd700 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass, field from typing import Optional import torch @@ -14,27 +15,37 @@ from sglang.srt.debug_utils.comparator.dims import ( ParallelAxis, resolve_dim_by_name, ) -from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning -from sglang.srt.debug_utils.comparator.warning_sink import warning_sink +from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult +from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import compute_diff + +_REPLICATED_ATOL: float = 1e-6 + + +@dataclass(frozen=True) +class UnsharderResult: + tensors: list[torch.Tensor] + replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list) def execute_unsharder_plan( plan: UnsharderPlan, tensors: list[torch.Tensor], -) -> list[torch.Tensor]: - result: list[torch.Tensor] = [] +) -> UnsharderResult: + result_tensors: list[torch.Tensor] = [] + all_checks: list[ReplicatedCheckResult] = [] for group_idx, group in enumerate(plan.groups): group_tensors = [tensors[i] for i in group] - tensor = _apply_unshard( + tensor, checks = _apply_unshard( plan.params, group_tensors, axis=plan.axis, group_index=group_idx, ) - result.append(tensor) + result_tensors.append(tensor) + all_checks.extend(checks) - return result + return UnsharderResult(tensors=result_tensors, replicated_checks=all_checks) def _apply_unshard( @@ -43,25 +54,28 @@ def _apply_unshard( *, axis: ParallelAxis, group_index: int, -) -> torch.Tensor: +) -> tuple[torch.Tensor, list[ReplicatedCheckResult]]: if isinstance(params, PickParams): - _verify_replicated_group( + checks: list[ReplicatedCheckResult] = _verify_replicated_group( ordered_tensors, axis=axis, group_index=group_index, ) - return ordered_tensors[0] + return ordered_tensors[0], checks if isinstance(params, ConcatParams): dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name) - return torch.cat(ordered_tensors, dim=dim) + return torch.cat(ordered_tensors, dim=dim), [] if isinstance(params, CpThdConcatParams): thd_dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name) - return _thd_concat( - ordered_tensors, - dim=thd_dim, - seq_lens_per_rank=params.seq_lens_per_rank, + return ( + _thd_concat( + ordered_tensors, + dim=thd_dim, + seq_lens_per_rank=params.seq_lens_per_rank, + ), + [], ) if isinstance(params, ReduceSumParams): @@ -70,7 +84,7 @@ def _apply_unshard( names: tuple[Optional[str], ...] = ordered_tensors[0].names if names[0] is not None: result = result.refine_names(*names) - return result + return result, [] raise ValueError(f"Unsupported unshard operation: {type(params).__name__}") @@ -80,21 +94,31 @@ def _verify_replicated_group( *, axis: ParallelAxis, group_index: int, -) -> None: - baseline = ordered_tensors[0].rename(None) +) -> list[ReplicatedCheckResult]: + baseline: torch.Tensor = ordered_tensors[0].rename(None).float() + checks: list[ReplicatedCheckResult] = [] for i in range(1, len(ordered_tensors)): - other = ordered_tensors[i].rename(None) - if not torch.allclose(baseline, other, atol=1e-6): - warning_sink.add( - ReplicatedMismatchWarning( - axis=axis.value, - group_index=group_index, - differing_index=i, - baseline_index=0, - max_abs_diff=(baseline - other).abs().max().item(), - ) + other: torch.Tensor = ordered_tensors[i].rename(None).float() + diff_info = compute_diff( + x_baseline=baseline, + x_target=other, + diff_threshold=_REPLICATED_ATOL, + ) + passed: bool = diff_info.max_abs_diff <= _REPLICATED_ATOL + checks.append( + ReplicatedCheckResult( + axis=axis.value, + group_index=group_index, + compared_index=i, + baseline_index=0, + passed=passed, + atol=_REPLICATED_ATOL, + diff=diff_info, ) + ) + + return checks def _thd_concat( diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 54daf41ea..eca4dc8df 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -160,6 +160,7 @@ def _compare_bundle_pair_tensor_type( aligner_result: AlignerResult = execute_aligner_plan( tensors_pair=tensors_pair, plan=plan ) + replicated_checks = aligner_result.replicated_checks if aligner_result.tensors is None: assert aligner_result.failed_side_xy is not None @@ -183,7 +184,11 @@ def _compare_bundle_pair_tensor_type( diff_threshold=diff_threshold, seq_dim=seq_dim, ) - record = ComparisonRecord(**info.model_dump(), aligner_plan=plan) + record = ComparisonRecord( + **info.model_dump(), + aligner_plan=plan, + replicated_checks=replicated_checks, + ) if viz_output_dir is not None: _try_generate_viz( diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index 1aa784d91..c83db102e 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -8,8 +8,10 @@ from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_valida from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( format_comparison, + format_replicated_checks, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( + DiffInfo, TensorComparisonInfo, ) from sglang.srt.debug_utils.comparator.utils import _StrictBase @@ -20,22 +22,6 @@ if TYPE_CHECKING: ) -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})" - ) - - class GeneralWarning(_StrictBase): kind: Literal["general"] = "general" category: str @@ -45,10 +31,19 @@ class GeneralWarning(_StrictBase): return self.message -AnyWarning = Annotated[ - Union[ReplicatedMismatchWarning, GeneralWarning], - Discriminator("kind"), -] +# Type alias — currently only GeneralWarning exists. +# When adding new warning types, convert back to Union + Discriminator("kind"). +AnyWarning = GeneralWarning + + +class ReplicatedCheckResult(_StrictBase): + axis: str + group_index: int + compared_index: int + baseline_index: int + passed: bool + atol: float + diff: DiffInfo class _OutputRecord(_StrictBase): @@ -126,15 +121,20 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord): type: Literal["comparison"] = "comparison" aligner_plan: Optional[AlignerPlan] = None + replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list) @property def category(self) -> str: if self.warnings: return "failed" + if any(not check.passed for check in self.replicated_checks): + return "failed" return "passed" if self.diff is not None and self.diff.passed else "failed" def _format_body(self) -> str: body: str = format_comparison(self) + if self.replicated_checks: + body += "\n" + format_replicated_checks(self.replicated_checks) if self.aligner_plan is not None: body += "\n" + _format_aligner_plan(self.aligner_plan) return body diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py index 46016bfbd..296429f92 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py @@ -57,7 +57,7 @@ def compare_tensor_pair( downcast_dtype: Optional[torch.dtype] = None if not shape_mismatch: - diff = _compute_diff( + diff = compute_diff( x_baseline=x_baseline_f, x_target=x_target_f, diff_threshold=diff_threshold, @@ -74,7 +74,7 @@ def compare_tensor_pair( Pair(x=baseline_original_dtype, y=target_original_dtype) ) if downcast_dtype is not None: - diff_downcast = _compute_diff( + diff_downcast = compute_diff( x_baseline=x_baseline_f.to(downcast_dtype), x_target=x_target_f.to(downcast_dtype), diff_threshold=diff_threshold, @@ -121,7 +121,7 @@ def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]: return {p: torch.quantile(x_float, p / 100.0).item() for p in DEFAULT_PERCENTILES} -def _compute_diff( +def compute_diff( x_baseline: torch.Tensor, x_target: torch.Tensor, diff_threshold: float = 1e-3, diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py index fdc14fc57..419c962d1 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py @@ -1,9 +1,16 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( DiffInfo, TensorComparisonInfo, TensorStats, ) +if TYPE_CHECKING: + from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult + def format_comparison(info: TensorComparisonInfo) -> str: lines: list[str] = [] @@ -54,6 +61,22 @@ def format_comparison(info: TensorComparisonInfo) -> str: return "\n".join(lines) +def format_replicated_checks(checks: list[ReplicatedCheckResult]) -> str: + lines: list[str] = ["Replicated checks:"] + + for check in checks: + marker: str = "✅" if check.passed else "❌" + lines.append( + f" {marker} axis={check.axis} group={check.group_index} " + f"idx={check.compared_index} vs {check.baseline_index}: " + f"rel_diff={check.diff.rel_diff:.6e} " + f"max_abs_diff={check.diff.max_abs_diff:.6e} " + f"mean_abs_diff={check.diff.mean_abs_diff:.6e}" + ) + + return "\n".join(lines) + + def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]: lines: list[str] = [] diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py index 6c4638273..151d52f22 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py @@ -1,5 +1,4 @@ import sys -from typing import Optional import pytest import torch @@ -32,22 +31,25 @@ register_cpu_ci(est_time=15, suite="default", nightly=True) class TestExecuteSubPlans: def test_empty_tensors_returns_none(self) -> None: - result: Optional[torch.Tensor] = execute_sub_plans(tensors=[], plans=[]) + result, checks = execute_sub_plans(tensors=[], plans=[]) assert result is None + assert checks == [] def test_no_plans_single_tensor_passthrough(self) -> None: tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0]) - result: Optional[torch.Tensor] = execute_sub_plans(tensors=[tensor], plans=[]) + result, checks = execute_sub_plans(tensors=[tensor], plans=[]) assert result is not None assert torch.equal(result, tensor) + assert checks == [] def test_no_plans_multiple_tensors_returns_none(self) -> None: tensors: list[torch.Tensor] = [ torch.tensor([1.0]), torch.tensor([2.0]), ] - result: Optional[torch.Tensor] = execute_sub_plans(tensors=tensors, plans=[]) + result, checks = execute_sub_plans(tensors=tensors, plans=[]) assert result is None + assert checks == [] def test_with_unsharder_plan(self) -> None: t0: torch.Tensor = torch.tensor([[1.0, 2.0]]).refine_names("b", "h") @@ -59,13 +61,12 @@ class TestExecuteSubPlans: groups=[[0, 1]], ) - result: Optional[torch.Tensor] = execute_sub_plans( - tensors=[t0, t1], plans=[plan] - ) + result, checks = execute_sub_plans(tensors=[t0, t1], plans=[plan]) assert result is not None expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) assert torch.equal(result.rename(None), expected) + assert checks == [] class TestExecuteSubPlan: @@ -74,7 +75,7 @@ class TestExecuteSubPlan: pass with pytest.raises(NotImplementedError, match="Unknown"): - execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan()) # type: ignore[arg-type] + execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan()) class TestExecuteStepPlans: @@ -90,11 +91,10 @@ class TestExecuteStepPlans: sub_plans=[], ) - result: dict[int, torch.Tensor] = _execute_step_plans( - tensors=tensors, step_plans=[step_plan] - ) + result, checks = _execute_step_plans(tensors=tensors, step_plans=[step_plan]) assert result == {} + assert checks == [] def test_single_step_passthrough(self) -> None: tensor: torch.Tensor = torch.tensor([1.0, 2.0]) @@ -105,12 +105,11 @@ class TestExecuteStepPlans: sub_plans=[], ) - result: dict[int, torch.Tensor] = _execute_step_plans( - tensors=[tensor], step_plans=[step_plan] - ) + result, checks = _execute_step_plans(tensors=[tensor], step_plans=[step_plan]) assert 5 in result assert torch.equal(result[5], tensor) + assert checks == [] class TestExecuteAlignerPlan: diff --git a/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py b/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py index da0764e9b..bdfb4a7e4 100644 --- a/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py @@ -20,7 +20,6 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( UnsharderPlan, ) from sglang.srt.debug_utils.comparator.dims import ParallelAxis -from sglang.srt.debug_utils.comparator.warning_sink import warning_sink from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) @@ -227,10 +226,8 @@ class TestThdCpZigzagE2E: params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=seq_lens_per_rank), groups=[[0, 1]], ) - with warning_sink.context(): - unsharded: list[torch.Tensor] = execute_unsharder_plan( - unshard_plan, rank_tensors - ) + unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors) + unsharded: list[torch.Tensor] = unsharder_result.tensors assert len(unsharded) == 1 # Step 2: THD reorder @@ -266,10 +263,8 @@ class TestThdCpZigzagE2E: ), groups=[list(range(cp_size))], ) - with warning_sink.context(): - unsharded: list[torch.Tensor] = execute_unsharder_plan( - unshard_plan, rank_tensors - ) + unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors) + unsharded: list[torch.Tensor] = unsharder_result.tensors assert len(unsharded) == 1 # Step 2: THD reorder diff --git a/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py b/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py index e8aed9ad5..1bc96ac97 100644 --- a/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py @@ -18,7 +18,6 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import ( ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis, parse_dims -from sglang.srt.debug_utils.comparator.warning_sink import warning_sink from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) @@ -157,12 +156,11 @@ class TestCpZigzagTpE2E: assert len(reorderer_plans) == 1 current: list[torch.Tensor] = [t.refine_names(*dim_names) for t in tensors] - with warning_sink.context(): - for plan in all_plans: - if isinstance(plan, ReordererPlan): - current = execute_reorderer_plan(plan, current) - else: - current = execute_unsharder_plan(plan, current) + for plan in all_plans: + if isinstance(plan, ReordererPlan): + current = execute_reorderer_plan(plan, current) + else: + current = execute_unsharder_plan(plan, current).tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py index e72a25964..57283e1ce 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py @@ -4,6 +4,7 @@ import pytest import torch from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import ( + UnsharderResult, _apply_unshard, _verify_replicated_group, execute_unsharder_plan, @@ -23,7 +24,7 @@ from sglang.srt.debug_utils.comparator.dims import ( ParallelAxis, parse_dims, ) -from sglang.srt.debug_utils.comparator.warning_sink import warning_sink +from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) @@ -49,11 +50,12 @@ class TestExecuteUnsharderPlan: assert len(plans) == 1 named_shards: list[torch.Tensor] = _name_tensors(shards, dim_specs) - with warning_sink.context() as warnings: - result = execute_unsharder_plan(plans[0], named_shards) - assert len(result) == 1 - assert torch.allclose(result[0].rename(None), full_tensor) - assert warnings == [] + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], named_shards + ) + assert len(unsharder_result.tensors) == 1 + assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor) + assert unsharder_result.replicated_checks == [] def test_scrambled_world_ranks_correct_result(self) -> None: full_tensor = torch.randn(4, 8) @@ -79,11 +81,12 @@ class TestExecuteUnsharderPlan: dim_specs, ) - with warning_sink.context() as warnings: - result = execute_unsharder_plan(plans[0], tensors_ordered_by_world_rank) - assert len(result) == 1 - assert torch.allclose(result[0].rename(None), full_tensor) - assert warnings == [] + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], tensors_ordered_by_world_rank + ) + assert len(unsharder_result.tensors) == 1 + assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor) + assert unsharder_result.replicated_checks == [] def test_single_step_reduces_tensor_count(self) -> None: """8 tensors with 2 groups of 4 produce 2 output tensors.""" @@ -113,13 +116,15 @@ class TestExecuteUnsharderPlan: tensors.append(source[tp_rank]) named_tensors: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - intermediate = execute_unsharder_plan(plans[0], named_tensors) - assert len(intermediate) == 4 + intermediate_result: UnsharderResult = execute_unsharder_plan( + plans[0], named_tensors + ) + assert len(intermediate_result.tensors) == 4 - with warning_sink.context(): - final = execute_unsharder_plan(plans[1], intermediate) - assert len(final) == 1 + final_result: UnsharderResult = execute_unsharder_plan( + plans[1], intermediate_result.tensors + ) + assert len(final_result.tensors) == 1 def test_cp_tp_concat(self) -> None: """CP=2 + TP=2: multi-step unshard reconstructs original tensor.""" @@ -145,9 +150,9 @@ class TestExecuteUnsharderPlan: assert len(plans) == 2 current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) @@ -187,9 +192,9 @@ class TestExecuteUnsharderPlan: assert len(plans) == 2 current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) @@ -241,9 +246,9 @@ class TestExecuteUnsharderPlan: assert len(plans) == 3 current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) @@ -290,9 +295,9 @@ class TestExecuteUnsharderPlan: assert len(plans) == 3 current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) @@ -312,11 +317,12 @@ class TestPickOperation: assert len(plans) == 1 assert isinstance(plans[0].params, PickParams) - with warning_sink.context() as warnings: - result = execute_unsharder_plan(plans[0], [tensor, tensor.clone()]) - assert len(result) == 1 - assert torch.allclose(result[0].rename(None), tensor) - assert warnings == [] + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], [tensor, tensor.clone()] + ) + assert len(unsharder_result.tensors) == 1 + assert torch.allclose(unsharder_result.tensors[0].rename(None), tensor) + assert all(c.passed for c in unsharder_result.replicated_checks) def test_pick_multiple_groups(self) -> None: """PickParams with multiple groups picks one from each.""" @@ -348,10 +354,11 @@ class TestPickOperation: tensor = torch.randn(4) tensors = [tensor.clone() for _ in range(4)] - with warning_sink.context() as warnings: - result = execute_unsharder_plan(pick_plans[0], tensors) - assert len(result) == 2 - assert warnings == [] + unsharder_result: UnsharderResult = execute_unsharder_plan( + pick_plans[0], tensors + ) + assert len(unsharder_result.tensors) == 2 + assert all(c.passed for c in unsharder_result.replicated_checks) def test_replicated_tp_sharded_cp_e2e(self) -> None: """CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip.""" @@ -376,9 +383,9 @@ class TestPickOperation: assert len(plans) == 2 current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) @@ -406,44 +413,44 @@ class TestPickOperation: assert all(isinstance(p.params, PickParams) for p in plans) current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) class TestVerifyReplicatedGroup: - def test_warns_on_mismatch(self) -> None: - """_verify_replicated_group produces warning when replicas differ.""" + def test_fails_on_mismatch(self) -> None: + """_verify_replicated_group returns failed check when replicas differ.""" tensor_a = torch.ones(4) tensor_b = torch.ones(4) + 0.1 - with warning_sink.context() as 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) + checks: list[ReplicatedCheckResult] = _verify_replicated_group( + [tensor_a, tensor_b], + axis=ParallelAxis.TP, + group_index=0, + ) + assert len(checks) == 1 + assert checks[0].axis == "tp" + assert checks[0].group_index == 0 + assert checks[0].compared_index == 1 + assert checks[0].baseline_index == 0 + assert not checks[0].passed + assert checks[0].diff.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.""" + def test_passes_when_identical(self) -> None: + """_verify_replicated_group returns passed check for identical replicas.""" tensor = torch.randn(4, 8) - with warning_sink.context() as warnings: - _verify_replicated_group( - [tensor, tensor.clone()], - axis=ParallelAxis.TP, - group_index=0, - ) - assert warnings == [] + checks: list[ReplicatedCheckResult] = _verify_replicated_group( + [tensor, tensor.clone()], + axis=ParallelAxis.TP, + group_index=0, + ) + assert len(checks) == 1 + assert checks[0].passed def test_multiple_mismatches(self) -> None: """_verify_replicated_group reports each differing replica.""" @@ -451,19 +458,20 @@ class TestVerifyReplicatedGroup: other_a = torch.ones(4) other_b = torch.ones(4) * 2 - with warning_sink.context() as 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) + checks: list[ReplicatedCheckResult] = _verify_replicated_group( + [baseline, other_a, other_b], + axis=ParallelAxis.CP, + group_index=1, + ) + assert len(checks) == 2 + assert checks[0].compared_index == 1 + assert not checks[0].passed + assert checks[1].compared_index == 2 + assert not checks[1].passed + assert checks[1].diff.max_abs_diff == pytest.approx(2.0, abs=1e-5) - def test_execute_returns_warnings(self) -> None: - """execute_unsharder_plan emits warnings for replicated mismatch.""" + def test_execute_returns_replicated_checks(self) -> None: + """execute_unsharder_plan returns replicated checks for mismatch.""" dim_specs = parse_dims("h d") parallel_infos = [ {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)}, @@ -474,56 +482,58 @@ class TestVerifyReplicatedGroup: tensor_a = torch.zeros(4) tensor_b = torch.ones(4) - with warning_sink.context() as warnings: - result = execute_unsharder_plan(plans[0], [tensor_a, tensor_b]) - assert len(result) == 1 - assert len(warnings) == 1 - assert torch.allclose(result[0].rename(None), tensor_a) + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], [tensor_a, tensor_b] + ) + assert len(unsharder_result.tensors) == 1 + assert len(unsharder_result.replicated_checks) == 1 + assert not unsharder_result.replicated_checks[0].passed + assert torch.allclose(unsharder_result.tensors[0].rename(None), tensor_a) def test_atol_boundary_within(self) -> None: - """Difference exactly at atol (1e-6) -> torch.allclose passes -> no warning.""" + """Difference exactly at atol (1e-6) -> passed.""" baseline = torch.zeros(4) other = torch.full((4,), 1e-6) - with warning_sink.context() as warnings: - _verify_replicated_group( - [baseline, other], - axis=ParallelAxis.TP, - group_index=0, - ) - assert warnings == [] + checks: list[ReplicatedCheckResult] = _verify_replicated_group( + [baseline, other], + axis=ParallelAxis.TP, + group_index=0, + ) + assert len(checks) == 1 + assert checks[0].passed def test_atol_boundary_exceeded(self) -> None: - """Difference just above atol (1e-6 + 1e-9) -> torch.allclose fails -> warning.""" + """Difference just above atol (1e-6 + 1e-9) -> failed.""" baseline = torch.zeros(4) other = torch.full((4,), 1e-6 + 1e-9) - with warning_sink.context() as warnings: - _verify_replicated_group( - [baseline, other], - axis=ParallelAxis.TP, - group_index=0, - ) - assert len(warnings) == 1 - assert warnings[0].differing_index == 1 + checks: list[ReplicatedCheckResult] = _verify_replicated_group( + [baseline, other], + axis=ParallelAxis.TP, + group_index=0, + ) + assert len(checks) == 1 + assert not checks[0].passed + assert checks[0].compared_index == 1 - def test_recompute_pseudo_mismatch_warns(self) -> None: - """_verify_replicated_group produces warning for RECOMPUTE_PSEUDO axis mismatch.""" + def test_recompute_pseudo_mismatch(self) -> None: + """_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch.""" tensor_a = torch.ones(4) tensor_b = torch.ones(4) + 0.1 - with warning_sink.context() as warnings: - _verify_replicated_group( - [tensor_a, tensor_b], - axis=ParallelAxis.RECOMPUTE_PSEUDO, - group_index=0, - ) - assert len(warnings) == 1 - assert warnings[0].axis == "recompute_pseudo" - 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) + checks: list[ReplicatedCheckResult] = _verify_replicated_group( + [tensor_a, tensor_b], + axis=ParallelAxis.RECOMPUTE_PSEUDO, + group_index=0, + ) + assert len(checks) == 1 + assert checks[0].axis == "recompute_pseudo" + assert checks[0].group_index == 0 + assert checks[0].compared_index == 1 + assert checks[0].baseline_index == 0 + assert not checks[0].passed + assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5) class TestThdCpConcat: @@ -537,12 +547,11 @@ class TestThdCpConcat: params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]), groups=[[0, 1]], ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1]) - assert len(result) == 1 + assert len(unsharder_result.tensors) == 1 expected = torch.tensor([1, 2, 3, 4, 5, 6]) - assert torch.equal(result[0].rename(None), expected) + assert torch.equal(unsharder_result.tensors[0].rename(None), expected) def test_multi_seq(self) -> None: """Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46].""" @@ -563,11 +572,10 @@ class TestThdCpConcat: params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]), groups=[[0, 1]], ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1]) - assert len(result) == 1 - unsharded: torch.Tensor = result[0].rename(None) + assert len(unsharder_result.tensors) == 1 + unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None) # seqA: r0(50) + r1(50) = 100 tokens, values 0..99 assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1])) @@ -595,11 +603,10 @@ class TestThdCpConcat: params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]), groups=[[0, 1]], ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1]) - assert len(result) == 1 - unsharded: torch.Tensor = result[0].rename(None) + assert len(unsharder_result.tensors) == 1 + unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None) assert unsharded.shape == (10, hidden) assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1])) @@ -625,11 +632,10 @@ class TestThdCpConcat: params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]), groups=[[0, 1]], ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [rank0, rank1]) + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1]) - assert len(result) == 1 - unsharded: torch.Tensor = result[0].rename(None) + assert len(unsharder_result.tensors) == 1 + unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None) assert unsharded.shape == (batch, 10, hidden) # seqA: r0(3) + r1(3) = 6 tokens per batch @@ -657,11 +663,12 @@ class TestReduceSum: assert isinstance(plans[0].params, ReduceSumParams) named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs) - with warning_sink.context(): - result = execute_unsharder_plan(plans[0], named_parts) + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], named_parts + ) - assert len(result) == 1 - assert torch.allclose(result[0].rename(None), full_tensor) + assert len(unsharder_result.tensors) == 1 + assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor) def test_tp4_reduce(self) -> None: """4 partial tensors sum to full tensor.""" @@ -677,11 +684,12 @@ class TestReduceSum: assert len(plans) == 1 named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs) - with warning_sink.context(): - result = execute_unsharder_plan(plans[0], named_parts) + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], named_parts + ) - assert len(result) == 1 - assert torch.allclose(result[0].rename(None), full_tensor) + assert len(unsharder_result.tensors) == 1 + assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor) def test_multi_axis_concat_then_reduce(self) -> None: """CP concat + TP reduce end-to-end.""" @@ -707,9 +715,9 @@ class TestReduceSum: assert len(plans) == 2 current: list[torch.Tensor] = _name_tensors(tensors, dim_specs) - with warning_sink.context(): - for plan in plans: - current = execute_unsharder_plan(plan, current) + for plan in plans: + unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current) + current = unsharder_result.tensors assert len(current) == 1 assert torch.allclose(current[0].rename(None), full_tensor) @@ -735,11 +743,12 @@ class TestReduceSum: plans = compute_unsharder_plan(dim_specs, parallel_infos) named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs) - with warning_sink.context(): - result = execute_unsharder_plan(plans[0], named_parts) + unsharder_result: UnsharderResult = execute_unsharder_plan( + plans[0], named_parts + ) - assert len(result) == 1 - assert torch.allclose(result[0].rename(None), full_tensor) + assert len(unsharder_result.tensors) == 1 + assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor) def test_reduce_preserves_named_dims(self) -> None: """Named tensor dimensions are preserved through reduce_sum.""" @@ -752,13 +761,16 @@ class TestReduceSum: params=ReduceSumParams(), groups=[[0, 1]], ) - with warning_sink.context(): - result = execute_unsharder_plan(plan, [part_a, part_b]) + unsharder_result: UnsharderResult = execute_unsharder_plan( + plan, [part_a, part_b] + ) - assert len(result) == 1 - assert result[0].names == ("h", "d") + assert len(unsharder_result.tensors) == 1 + assert unsharder_result.tensors[0].names == ("h", "d") expected = (part_a.rename(None) + part_b.rename(None)).refine_names("h", "d") - assert torch.allclose(result[0].rename(None), expected.rename(None)) + assert torch.allclose( + unsharder_result.tensors[0].rename(None), expected.rename(None) + ) if __name__ == "__main__": diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py index 9851b064d..426336121 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py @@ -6,9 +6,9 @@ import torch from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( QUANTILE_NUMEL_THRESHOLD, SAMPLE_DIFF_THRESHOLD, - _compute_diff, _compute_tensor_stats, compare_tensor_pair, + compute_diff, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import DiffInfo from sglang.test.ci.ci_register import register_cpu_ci @@ -55,7 +55,7 @@ class TestComputeTensorStats: class TestComputeDiff: def test_identical_tensors(self): x = torch.ones(10, 10) - diff = _compute_diff(x_baseline=x, x_target=x) + diff = compute_diff(x_baseline=x, x_target=x) assert diff.rel_diff == pytest.approx(0.0, abs=1e-5) assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5) @@ -70,7 +70,7 @@ class TestComputeDiff: y = x.clone() y[3, 7] = 1.5 - diff = _compute_diff(x_baseline=x, x_target=y) + diff = compute_diff(x_baseline=x, x_target=y) assert diff.max_abs_diff == pytest.approx(0.5, abs=1e-4) assert diff.max_diff_coord == [3, 7] @@ -85,14 +85,14 @@ class TestComputeDiff: def test_large_tensor_skips_diff_quantiles(self): x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1) y = x + 0.001 - diff = _compute_diff(x_baseline=x, x_target=y) + diff = compute_diff(x_baseline=x, x_target=y) assert diff.abs_diff_percentiles == {} def test_rel_diff_value(self): x = torch.tensor([1.0, 0.0]) y = torch.tensor([0.0, 1.0]) - diff = _compute_diff(x_baseline=x, x_target=y) + diff = compute_diff(x_baseline=x, x_target=y) assert diff.rel_diff == pytest.approx(1.0, abs=1e-5) assert diff.passed is False @@ -103,7 +103,7 @@ class TestComputeDiff: x: torch.Tensor = torch.randn(8, 16) y: torch.Tensor = x + torch.randn_like(x) * 0.01 - diff: DiffInfo = _compute_diff( + diff: DiffInfo = compute_diff( x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0 ) @@ -117,7 +117,7 @@ class TestComputeDiff: x: torch.Tensor = torch.randn(8, 16) y: torch.Tensor = x + torch.randn_like(x) * 0.01 - diff: DiffInfo = _compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3) + diff: DiffInfo = compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3) assert diff.per_token_rel_diff is None @@ -127,7 +127,7 @@ class TestComputeDiff: x: torch.Tensor = torch.randn(4, 8) y: torch.Tensor = x + torch.randn_like(x) * 0.01 - diff: DiffInfo = _compute_diff( + diff: DiffInfo = compute_diff( x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0 ) diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py index db73fe9c1..7dddc1783 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py @@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.output_types import ( ComparisonRecord, ConfigRecord, GeneralWarning, - ReplicatedMismatchWarning, + ReplicatedCheckResult, SkipRecord, SummaryRecord, WarningRecord, @@ -126,16 +126,24 @@ class TestRecordTypes: assert restored == record -def _make_warning(**overrides) -> ReplicatedMismatchWarning: +def _make_replicated_check(**overrides) -> ReplicatedCheckResult: defaults: dict = dict( axis="tp", group_index=0, - differing_index=1, + compared_index=1, baseline_index=0, - max_abs_diff=0.1, + passed=False, + atol=1e-6, + diff=_make_diff( + rel_diff=0.1, + max_abs_diff=0.1, + mean_abs_diff=0.05, + diff_threshold=1e-6, + passed=False, + ), ) defaults.update(overrides) - return ReplicatedMismatchWarning(**defaults) + return ReplicatedCheckResult(**defaults) class TestWarnings: @@ -148,7 +156,7 @@ class TestWarnings: unified_shape=[4, 8], shape_mismatch=False, diff=_make_diff(passed=True), - warnings=[_make_warning()], + warnings=[GeneralWarning(category="test", message="some warning")], ) assert record.category == "failed" @@ -157,18 +165,44 @@ class TestWarnings: record = SkipRecord( name="x", reason="no_baseline", - warnings=[_make_warning()], + warnings=[GeneralWarning(category="test", message="some warning")], ) assert record.category == "failed" - def test_warnings_json_round_trip(self): - """warnings survive model_dump_json → parse_record_json round-trip.""" - warning = _make_warning( + def test_replicated_checks_all_passed(self): + """ComparisonRecord with all replicated_checks passed → category=='passed'.""" + record = ComparisonRecord( + name="hidden", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(passed=True), + replicated_checks=[_make_replicated_check(passed=True)], + ) + assert record.category == "passed" + + def test_replicated_checks_failed_means_record_failed(self): + """ComparisonRecord with any replicated_check.passed=False → 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), + replicated_checks=[_make_replicated_check(passed=False)], + ) + assert record.category == "failed" + + def test_replicated_check_json_round_trip(self): + """ReplicatedCheckResult survives JSON round-trip via ComparisonRecord.""" + check = _make_replicated_check( axis="cp", group_index=2, - differing_index=3, + compared_index=3, baseline_index=0, - max_abs_diff=0.42, + passed=False, ) record = ComparisonRecord( name="mlp", @@ -177,30 +211,23 @@ class TestWarnings: unified_shape=[4, 8], shape_mismatch=False, diff=_make_diff(), - warnings=[warning], + replicated_checks=[check], ) restored = parse_record_json(record.model_dump_json()) assert isinstance(restored, ComparisonRecord) - assert len(restored.warnings) == 1 + assert len(restored.replicated_checks) == 1 - restored_warning = restored.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) + restored_check: ReplicatedCheckResult = restored.replicated_checks[0] + assert restored_check.axis == "cp" + assert restored_check.group_index == 2 + assert restored_check.compared_index == 3 + assert restored_check.baseline_index == 0 + assert not restored_check.passed def test_any_warning_discriminated_union_round_trip(self): """All AnyWarning variants survive JSON round-trip via a WarningRecord.""" all_warnings = [ - ReplicatedMismatchWarning( - axis="tp", - group_index=0, - differing_index=1, - baseline_index=0, - max_abs_diff=0.1, - ), GeneralWarning( category="aux_tensors_missing", message="Aux tensors missing, skipping token alignment", diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index d54b57b01..1fc285377 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -13,7 +13,7 @@ from sglang.srt.debug_utils.comparator.output_types import ( ConfigRecord, GeneralWarning, NonTensorRecord, - ReplicatedMismatchWarning, + ReplicatedCheckResult, SkipRecord, SummaryRecord, WarningRecord, @@ -910,7 +910,7 @@ class TestEntrypointGroupingLogical: assert comp.name == "hidden" def test_recompute_pseudo_mismatch_warning(self, tmp_path, capsys): - """Recompute pseudo-axis with differing original/recompute → ReplicatedMismatchWarning.""" + """Recompute pseudo-axis with differing original/recompute → failed replicated_checks.""" torch.manual_seed(42) tensor = torch.randn(4, 8) mismatched_tensor = tensor + torch.randn(4, 8) * 10.0 @@ -937,12 +937,112 @@ class TestEntrypointGroupingLogical: comparisons = _get_comparisons(records) assert len(comparisons) == 1 - recompute_warnings = [ - w - for w in comparisons[0].warnings - if isinstance(w, ReplicatedMismatchWarning) and w.axis == "recompute_pseudo" + recompute_checks: list[ReplicatedCheckResult] = [ + c for c in comparisons[0].replicated_checks if c.axis == "recompute_pseudo" ] - assert len(recompute_warnings) > 0 + assert len(recompute_checks) > 0 + assert any(not c.passed for c in recompute_checks) + + def test_tp_partial_reduction_unshard(self, tmp_path, capsys): + """TP=2 with partial reduction: element-wise sum reconstructs full tensor.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8) + full_target = full_baseline + torch.randn(4, 8) * 0.001 + + baseline_dir = tmp_path / "baseline" + target_dir = tmp_path / "target" + + baseline_path = _create_tp_partial_dumps( + baseline_dir, + full_tensor=full_baseline, + name="attn_out", + tp_size=2, + dims_str="b h(tp,partial)", + ) + target_path = _create_tp_partial_dumps( + target_dir, + full_tensor=full_target, + name="attn_out", + tp_size=2, + dims_str="b h(tp,partial)", + ) + + args = _make_args(baseline_path, target_path, diff_threshold=0.01) + + records = _run_and_parse(args, capsys) + comp = _assert_single_comparison_passed(records) + assert comp.name == "attn_out" + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.total == 1 + assert summary.passed == 1 + + def test_tp_partial_vs_single_rank(self, tmp_path, capsys): + """Baseline single rank vs target TP=2 partial: unshard target then compare.""" + torch.manual_seed(42) + full_tensor = torch.randn(4, 8) + target_full = full_tensor + torch.randn(4, 8) * 0.001 + + baseline_dir = tmp_path / "baseline" + target_dir = tmp_path / "target" + + baseline_path = _create_rank_dump( + baseline_dir, rank=0, name="attn_out", tensor=full_tensor + ) + target_path = _create_tp_partial_dumps( + target_dir, + full_tensor=target_full, + name="attn_out", + tp_size=2, + dims_str="b h(tp,partial)", + ) + + args = _make_args(baseline_path, target_path, diff_threshold=0.01) + + records = _run_and_parse(args, capsys) + comp = _assert_single_comparison_passed(records) + assert comp.name == "attn_out" + + def test_cp_concat_tp_partial_reduction(self, tmp_path, capsys): + """CP=2 concat + TP=2 partial reduction: multi-axis unshard.""" + torch.manual_seed(42) + full_baseline = torch.randn(4, 8, 16) + full_target = full_baseline + torch.randn(4, 8, 16) * 0.001 + + for side_dir, full_tensor in [ + (tmp_path / "baseline", full_baseline), + (tmp_path / "target", full_target), + ]: + side_dir.mkdir() + cp_chunks = list(full_tensor.chunk(2, dim=1)) + rank = 0 + for cp_rank in range(2): + for tp_rank in range(2): + _create_rank_dump( + side_dir, + rank=rank, + name="hidden", + tensor=cp_chunks[cp_rank] / 2, + dims="b s(cp) h(tp,partial)", + parallel_info={ + "cp_rank": cp_rank, + "cp_size": 2, + "tp_rank": tp_rank, + "tp_size": 2, + }, + ) + rank += 1 + + args = _make_args( + tmp_path / "baseline" / _FIXED_EXP_NAME, + tmp_path / "target" / _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_tp_partial_reduction_unshard(self, tmp_path, capsys): """TP=2 with partial reduction: element-wise sum reconstructs full tensor.""" @@ -1233,7 +1333,7 @@ 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 warnings.""" + """CP2 TP2, TP replicated and identical → passed, replicated_checks all passed.""" torch.manual_seed(42) full_baseline = torch.randn(4, 8, 6) full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001 @@ -1264,13 +1364,14 @@ class TestEntrypointReplicatedAxis: records = _run_and_parse(args, capsys) comp = _assert_single_comparison_passed(records) assert comp.warnings == [] + assert all(c.passed for c in comp.replicated_checks) 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 warnings.""" + """CP2 TP2, TP replicas differ (> atol) → failed with replicated_checks.""" torch.manual_seed(42) full_baseline = torch.randn(4, 8, 6) full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001 @@ -1303,14 +1404,14 @@ class TestEntrypointReplicatedAxis: comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].category == "failed" - assert len(comparisons[0].warnings) > 0 + assert any(not c.passed for c in comparisons[0].replicated_checks) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.failed == 1 - def test_summary_counts_failed_from_warnings_only(self, tmp_path, capsys): - """Diff itself passes but TP replicas differ → summary.failed=1 from warnings.""" + def test_summary_counts_failed_from_replicated_checks_only(self, tmp_path, capsys): + """Diff itself passes but TP replicas differ → summary.failed=1 from replicated_checks.""" torch.manual_seed(42) full_baseline = torch.randn(4, 8, 6) full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001 @@ -1352,7 +1453,7 @@ class TestEntrypointReplicatedAxis: comp = comparisons[0] assert comp.diff is not None assert comp.diff.passed - assert len(comp.warnings) > 0 + assert any(not c.passed for c in comp.replicated_checks) assert comp.category == "failed" summary = records[-1] diff --git a/test/registered/debug_utils/comparator/test_warning_sink.py b/test/registered/debug_utils/comparator/test_warning_sink.py index 8f9a5ea82..7a41dcf4c 100644 --- a/test/registered/debug_utils/comparator/test_warning_sink.py +++ b/test/registered/debug_utils/comparator/test_warning_sink.py @@ -3,23 +3,20 @@ import sys import pytest -from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning +from sglang.srt.debug_utils.comparator.output_types import GeneralWarning from sglang.srt.debug_utils.comparator.warning_sink import WarningSink from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) -def _make_warning(**overrides) -> ReplicatedMismatchWarning: +def _make_warning(**overrides) -> GeneralWarning: defaults: dict = dict( - axis="tp", - group_index=0, - differing_index=1, - baseline_index=0, - max_abs_diff=0.1, + category="test", + message="test warning", ) defaults.update(overrides) - return ReplicatedMismatchWarning(**defaults) + return GeneralWarning(**defaults) class TestWarningSink: @@ -35,8 +32,8 @@ class TestWarningSink: def test_nested_contexts(self) -> None: sink = WarningSink() - outer_warning = _make_warning(group_index=0) - inner_warning = _make_warning(group_index=1) + outer_warning = _make_warning(message="outer") + inner_warning = _make_warning(message="inner") with sink.context() as outer: sink.add(outer_warning) @@ -61,7 +58,7 @@ class TestWarningSink: sink.add(_make_warning()) captured = capsys.readouterr() - assert "Replicated along tp" in captured.out + assert "test warning" in captured.out def test_context_captures_instead_of_printing(self, capsys) -> None: sink = WarningSink() @@ -96,9 +93,9 @@ class TestWarningSink: assert len(collected) == 1 - sink.add(_make_warning(group_index=99)) + sink.add(_make_warning(message="after exception")) captured = capsys.readouterr() - assert "Replicated along tp" in captured.out + assert "after exception" in captured.out if __name__ == "__main__":