From eb0e905fc3149bc6119716bbbbab082cf4cb9a57 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 27 Feb 2026 08:09:15 +0800 Subject: [PATCH] Update token layout and cleanup printer in dump comparator (#19457) --- .../aligner/token_aligner/aux_loader.py | 4 +- .../aligner/token_aligner/aux_plugins.py | 19 +- .../comparator/aligner/token_aligner/types.py | 5 +- .../sglang/srt/debug_utils/comparator/dims.py | 5 + .../tensor_comparator/comparator.py | 1 + .../comparator/tensor_comparator/formatter.py | 5 +- .../comparator/tensor_comparator/printer.py | 86 ------ .../comparator/tensor_comparator/types.py | 1 + .../aligner/token_aligner/test_aux_plugins.py | 15 +- .../aligner/token_aligner/test_executor.py | 3 +- .../aligner/token_aligner/test_planner.py | 23 +- .../tensor_comparator/test_formatter.py | 14 +- .../tensor_comparator/test_printer.py | 265 ------------------ .../tensor_comparator/test_types.py | 2 + .../comparator/test_model_validation.py | 1 + 15 files changed, 56 insertions(+), 393 deletions(-) delete mode 100644 python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py delete mode 100644 test/registered/debug_utils/comparator/tensor_comparator/test_printer.py 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 bc1aee1a3..90b9b8548 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 @@ -24,7 +24,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import ( normalize_parallel_info, ) -from sglang.srt.debug_utils.comparator.dims import ParallelAxis +from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout from sglang.srt.debug_utils.comparator.output_types import GeneralWarning from sglang.srt.debug_utils.comparator.warning_sink import warning_sink from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows @@ -61,7 +61,7 @@ def load_and_normalize_aux( if step_data: steps_data[step] = step_data - layout: str = plugin.detect_layout(steps_data) + layout: TokenLayout = plugin.detect_layout(steps_data) step_auxs: dict[int, TokenAlignerStepAux] = { step: plugin.compute_step_aux(step_data, layout=layout, step=step) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py index 656412807..ef959810f 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py @@ -10,6 +10,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( SGLangSeqId, TokenAlignerStepAux, ) +from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.output_types import GeneralWarning from sglang.srt.debug_utils.comparator.warning_sink import warning_sink @@ -45,11 +46,11 @@ class _AuxFrameworkPlugin(ABC): return frozenset() @abstractmethod - def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: ... + def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: ... @abstractmethod def compute_step_aux( - self, step_data: dict[str, object], *, layout: str, step: int + self, step_data: dict[str, object], *, layout: TokenLayout, step: int ) -> TokenAlignerStepAux: ... @abstractmethod @@ -89,11 +90,11 @@ class _SGLangPlugin(_AuxFrameworkPlugin): def has_required_names(self, names: set[str]) -> bool: return "input_ids" in names and "seq_lens" in names - def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: - return "thd" + def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: + return TokenLayout.T def compute_step_aux( - self, step_data: dict[str, object], *, layout: str, step: int + self, step_data: dict[str, object], *, layout: TokenLayout, step: int ) -> TokenAlignerStepAux: input_ids = step_data["input_ids"] positions = step_data["positions"] @@ -154,13 +155,13 @@ class _MegatronPlugin(_AuxFrameworkPlugin): def has_required_names(self, names: set[str]) -> bool: return "input_ids" in names and "cu_seqlens_q" in names - def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: + def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: for step_data in raw.values(): if (qkv_format := step_data.get("qkv_format")) is not None: fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format) if "bshd" in fmt.lower(): raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG) - return "thd" + return TokenLayout.T input_ids = step_data.get("input_ids") if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2: @@ -175,10 +176,10 @@ class _MegatronPlugin(_AuxFrameworkPlugin): ), ) ) - return "thd" + return TokenLayout.T def compute_step_aux( - self, step_data: dict[str, object], *, layout: str, step: int + self, step_data: dict[str, object], *, layout: TokenLayout, step: int ) -> TokenAlignerStepAux: input_ids: torch.Tensor = step_data["input_ids"] diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py index 33bf5dd88..3a6a2243f 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py @@ -5,6 +5,7 @@ from typing import NamedTuple, Union from pydantic import model_validator +from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.utils import ( Pair, _check_equal_lengths, @@ -50,7 +51,7 @@ class TokenAlignerGlobalAux: step_auxs: dict[int, TokenAlignerStepAux] framework: str # "sglang" | "megatron" - layout: str # "thd" + layout: TokenLayout class TokenLocator(_FrozenBase): @@ -106,7 +107,7 @@ class TokenAlignerSeqsInfo(_FrozenBase): """All sequences for one side across all steps.""" sequences: dict[SeqId, TokenAlignerSeqInfo] - layout: str + layout: TokenLayout class TokenAlignerPlan(_FrozenBase): diff --git a/python/sglang/srt/debug_utils/comparator/dims.py b/python/sglang/srt/debug_utils/comparator/dims.py index 3f2519e2f..925627201 100644 --- a/python/sglang/srt/debug_utils/comparator/dims.py +++ b/python/sglang/srt/debug_utils/comparator/dims.py @@ -8,6 +8,11 @@ BATCH_DIM_NAME: str = "b" SEQ_DIM_NAME: str = "s" +class TokenLayout(Enum): + T = "t" # single flat token dim + BS = "bs" # separate batch + seq dims, need collapse + + class ParallelAxis(Enum): TP = "tp" CP = "cp" 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 c26db9aa2..f4a17918c 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py @@ -125,6 +125,7 @@ def _compute_diff( max_diff_coord=list(max_diff_coord), baseline_at_max=x_baseline[max_diff_coord].item(), target_at_max=x_target[max_diff_coord].item(), + diff_threshold=diff_threshold, passed=( rel_diff <= diff_threshold and max_abs_diff <= diff_threshold 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 c300f85b3..e888bb352 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py @@ -69,13 +69,10 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]: - marker = "✅" if diff.passed else "❌" return [ prefix_text - + marker - + " " + "\t".join( - f"{name}={value}" + f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}" for name, value in [ ("rel_diff", diff.rel_diff), ("max_abs_diff", diff.max_abs_diff), diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py deleted file mode 100644 index 716474085..000000000 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py +++ /dev/null @@ -1,86 +0,0 @@ -from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( - DiffInfo, - TensorComparisonInfo, - TensorStats, -) - - -def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None: - baseline = info.baseline - target = info.target - - dtype_marker = "" if baseline.dtype == target.dtype else "🟠" - print( - f"Raw " - f"[shape] {baseline.shape} vs {target.shape}\t" - f"[{dtype_marker}dtype] {baseline.dtype} vs {target.dtype}" - ) - - if info.unified_shape != baseline.shape: - print( - f"Unify shape: {baseline.shape} -> {info.unified_shape} " - f"(to match {target.shape})" - ) - - print( - f"After unify " - f"[shape] {info.unified_shape} vs {target.shape}\t" - f"[dtype] {baseline.dtype} vs {target.dtype}" - ) - - _print_stats_comparison(baseline=baseline.stats, target=target.stats) - - if info.shape_mismatch: - print("⚠️ Shape mismatch") - return - - if info.diff is not None: - _print_diff( - diff=info.diff, - diff_threshold=diff_threshold, - ) - - if info.diff_downcast is not None and info.downcast_dtype is not None: - _print_diff( - diff=info.diff_downcast, - diff_threshold=diff_threshold, - prefix_text=f"When downcast to {info.downcast_dtype}: ", - ) - - if baseline.sample is not None: - print(f"x_baseline(sample)={baseline.sample}") - if target.sample is not None: - print(f"x_target(sample)={target.sample}") - - -def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None: - stat_names = list(TensorStats.model_fields.keys()) - for stat_name in stat_names: - value_baseline = getattr(baseline, stat_name) - value_target = getattr(target, stat_name) - if value_baseline is None or value_target is None: - continue - print( - f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} " - f"(diff: {value_target - value_baseline:.4f})" - ) - - -def _print_diff(diff: DiffInfo, diff_threshold: float, prefix_text: str = "") -> None: - print( - prefix_text - + "\t".join( - f"{'❌' if value > diff_threshold else '✅'} {name}={value}" - for name, value in [ - ("rel_diff", diff.rel_diff), - ("max_abs_diff", diff.max_abs_diff), - ("mean_abs_diff", diff.mean_abs_diff), - ] - ) - ) - - print( - f"max_abs_diff happens at coord={diff.max_diff_coord} with " - f"baseline={diff.baseline_at_max} " - f"target={diff.target_at_max}" - ) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py index a71402030..00982b6c6 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py @@ -28,6 +28,7 @@ class DiffInfo(_StrictBase): max_diff_coord: list[int] baseline_at_max: float target_at_max: float + diff_threshold: float passed: bool diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py index 6a8c01646..7c30e724c 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py @@ -13,6 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( SGLangSeqId, TokenAlignerStepAux, ) +from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=15, suite="default", nightly=True) @@ -34,7 +35,7 @@ class TestNormalizeSGLang: } result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux( - step_data, layout="thd", step=0 + step_data, layout=TokenLayout.T, step=0 ) assert result.input_ids == [10, 20, 30] @@ -51,7 +52,7 @@ class TestNormalizeSGLang: } result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux( - step_data, layout="thd", step=3 + step_data, layout=TokenLayout.T, step=3 ) assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)] @@ -65,7 +66,7 @@ class TestNormalizeSGLang: } result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux( - step_data, layout="thd", step=0 + step_data, layout=TokenLayout.T, step=0 ) assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")] @@ -81,7 +82,7 @@ class TestNormalizeMegatron: } result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( - step_data, layout="thd", step=0 + step_data, layout=TokenLayout.T, step=0 ) assert result.seq_lens == [3, 2] @@ -94,7 +95,7 @@ class TestNormalizeMegatron: } result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( - step_data, layout="thd", step=0 + step_data, layout=TokenLayout.T, step=0 ) assert result.positions == [0, 1, 2, 0, 1] @@ -108,7 +109,7 @@ class TestNormalizeMegatron: } result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( - step_data, layout="thd", step=0 + step_data, layout=TokenLayout.T, step=0 ) assert result.positions == [5, 6, 7, 8, 9] @@ -121,7 +122,7 @@ class TestNormalizeMegatron: } result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux( - step_data, layout="thd", step=5 + step_data, layout=TokenLayout.T, step=5 ) assert result.seq_ids == [ PositionalSeqId(step=5, seq_index=0), diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py index 072a8f391..ac7be83d9 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py @@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenAlignerStepAux, TokenLocator, ) +from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.utils import Pair from sglang.test.ci.ci_register import register_cpu_ci @@ -52,7 +53,7 @@ class TestExecuteAlignment: side_aux = TokenAlignerGlobalAux( step_auxs={0: aux, 1: aux_step1}, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py index b90ba46ee..b8d3a1c1c 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py @@ -19,6 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenAlignerStepAux, TokenLocator, ) +from sglang.srt.debug_utils.comparator.dims import TokenLayout from sglang.srt.debug_utils.comparator.utils import Pair from sglang.test.ci.ci_register import register_cpu_ci @@ -40,7 +41,7 @@ class TestBuildTokenIndexSGLangThd: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) @@ -75,7 +76,7 @@ class TestBuildTokenIndexSGLangThd: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) @@ -108,7 +109,7 @@ class TestBuildTokenIndexSGLangThd: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) @@ -132,7 +133,7 @@ class TestBuildTokenIndexSGLangThd: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) @@ -163,7 +164,7 @@ class TestBuildTokenIndexMegatronThd: ), }, framework="megatron", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) @@ -204,7 +205,7 @@ class TestBuildTokenIndexMegatronThd: ), }, framework="megatron", - layout="thd", + layout=TokenLayout.T, ) index = build_seqs_info(side_aux) @@ -389,7 +390,7 @@ class TestComputeAlignmentPlanCrossLayout: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) side_aux_b = TokenAlignerGlobalAux( step_auxs={ @@ -401,7 +402,7 @@ class TestComputeAlignmentPlanCrossLayout: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) index_a = build_seqs_info(side_aux_a) @@ -428,7 +429,7 @@ class TestComputeAlignmentPlanCrossLayout: ), }, framework="sglang", - layout="thd", + layout=TokenLayout.T, ) side_aux_b = TokenAlignerGlobalAux( step_auxs={ @@ -443,7 +444,7 @@ class TestComputeAlignmentPlanCrossLayout: ), }, framework="megatron", - layout="thd", + layout=TokenLayout.T, ) index_a = build_seqs_info(side_aux_a) @@ -467,7 +468,7 @@ def _int_to_seq_id(k: int) -> SeqId: def _make_index( *, sequences: dict[int, tuple[int, ...]], - layout: str = "thd", + layout: TokenLayout = TokenLayout.T, ) -> TokenAlignerSeqsInfo: """Create a TokenAlignerSeqsInfo from simplified input_ids-only specification.""" records: dict[SeqId, TokenAlignerSeqInfo] = {} diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py index ac3d6ab06..94fcb1035 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py @@ -35,6 +35,7 @@ def _make_diff( rel_diff: float = 0.0001, max_abs_diff: float = 0.0005, mean_abs_diff: float = 0.0002, + diff_threshold: float = 1e-3, passed: bool = True, ) -> DiffInfo: return DiffInfo( @@ -44,6 +45,7 @@ def _make_diff( max_diff_coord=[2, 3], baseline_at_max=1.0, target_at_max=1.0005, + diff_threshold=diff_threshold, passed=passed, ) @@ -94,7 +96,7 @@ class TestFormatComparison: "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005" ) @@ -153,11 +155,11 @@ class TestFormatComparison: "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "❌ rel_diff=0.002\tmax_abs_diff=0.005\tmean_abs_diff=0.001\n" + "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n" "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" "When downcast to torch.bfloat16: " - "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005" ) @@ -187,7 +189,7 @@ class TestFormatComparison: "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005" ) @@ -215,7 +217,7 @@ class TestFormatComparison: "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005\n" "x_baseline(sample)=tensor([0.1, 0.2, ...])\n" @@ -243,7 +245,7 @@ class TestFormatComparison: "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" "max_abs_diff happens at coord=[2, 3] with " "baseline=1.0 target=1.0005" ) diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_printer.py b/test/registered/debug_utils/comparator/tensor_comparator/test_printer.py deleted file mode 100644 index f03a83462..000000000 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_printer.py +++ /dev/null @@ -1,265 +0,0 @@ -import sys - -import pytest - -from sglang.srt.debug_utils.comparator.tensor_comparator.printer import ( - print_comparison, -) -from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( - DiffInfo, - TensorComparisonInfo, - TensorInfo, - TensorStats, -) -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=10, suite="default", nightly=True) - - -def _make_stats( - mean: float = 0.0, - std: float = 1.0, - min: float = -2.0, - max: float = 2.0, - p1: float | None = -1.8, - p5: float | None = -1.5, - p95: float | None = 1.5, - p99: float | None = 1.8, -) -> TensorStats: - return TensorStats( - mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99 - ) - - -def _make_diff( - rel_diff: float = 0.0001, - max_abs_diff: float = 0.0005, - mean_abs_diff: float = 0.0002, - passed: bool = True, -) -> DiffInfo: - return DiffInfo( - rel_diff=rel_diff, - max_abs_diff=max_abs_diff, - mean_abs_diff=mean_abs_diff, - max_diff_coord=[2, 3], - baseline_at_max=1.0, - target_at_max=1.0005, - passed=passed, - ) - - -def _make_tensor_info( - shape: list[int] | None = None, - dtype: str = "torch.float32", - stats: TensorStats | None = None, - sample: str | None = None, -) -> TensorInfo: - return TensorInfo( - shape=shape if shape is not None else [4, 8], - dtype=dtype, - stats=stats if stats is not None else _make_stats(), - sample=sample, - ) - - -# Snapshot strings below are intentionally spelled out in full per test. -# The shared skeleton (stats block, diff block) looks duplicated, but keeping -# each test self-contained makes failures immediately readable without chasing -# helper functions. Do not extract common fragments. -class TestPrintComparison: - def test_normal(self, capsys): - info = TensorComparisonInfo( - name="test", - baseline=_make_tensor_info( - stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0), - ), - target=_make_tensor_info( - stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001), - ), - unified_shape=[4, 8], - shape_mismatch=False, - diff=_make_diff(), - ) - - print_comparison(info=info, diff_threshold=1e-3) - - assert capsys.readouterr().out == ( - "Raw [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n" - "[std] 1.0000 vs 1.0001 (diff: 0.0001)\n" - "[min] -2.0000 vs -2.0001 (diff: -0.0001)\n" - "[max] 2.0000 vs 2.0001 (diff: 0.0001)\n" - "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" - "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" - "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" - "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=[2, 3] with " - "baseline=1.0 target=1.0005\n" - ) - - def test_shape_mismatch(self, capsys): - info = TensorComparisonInfo( - name="mismatch", - baseline=_make_tensor_info(shape=[3, 4]), - target=_make_tensor_info(shape=[5, 6]), - unified_shape=[3, 4], - shape_mismatch=True, - ) - - print_comparison(info=info, diff_threshold=1e-3) - - assert capsys.readouterr().out == ( - "Raw [shape] [3, 4] vs [5, 6]\t" - "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] [3, 4] vs [5, 6]\t" - "[dtype] torch.float32 vs torch.float32\n" - "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" - "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" - "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" - "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" - "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" - "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" - "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" - "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "⚠️ Shape mismatch\n" - ) - - def test_with_downcast(self, capsys): - info = TensorComparisonInfo( - name="downcast", - baseline=_make_tensor_info(), - target=_make_tensor_info(dtype="torch.bfloat16"), - unified_shape=[4, 8], - shape_mismatch=False, - diff=_make_diff( - rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001, passed=False - ), - diff_downcast=_make_diff( - rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002, passed=True - ), - downcast_dtype="torch.bfloat16", - ) - - print_comparison(info=info, diff_threshold=1e-3) - - assert capsys.readouterr().out == ( - "Raw [shape] [4, 8] vs [4, 8]\t" - "[🟠dtype] torch.float32 vs torch.bfloat16\n" - "After unify [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.bfloat16\n" - "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" - "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" - "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" - "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" - "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" - "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" - "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" - "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n" - "max_abs_diff happens at coord=[2, 3] with " - "baseline=1.0 target=1.0005\n" - "When downcast to torch.bfloat16: " - "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=[2, 3] with " - "baseline=1.0 target=1.0005\n" - ) - - def test_with_shape_unification(self, capsys): - info = TensorComparisonInfo( - name="unify", - baseline=_make_tensor_info(shape=[1, 1, 4, 8]), - target=_make_tensor_info(), - unified_shape=[4, 8], - shape_mismatch=False, - diff=_make_diff(), - ) - - print_comparison(info=info, diff_threshold=1e-3) - - assert capsys.readouterr().out == ( - "Raw [shape] [1, 1, 4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "Unify shape: [1, 1, 4, 8] -> [4, 8] " - "(to match [4, 8])\n" - "After unify [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" - "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" - "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" - "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" - "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" - "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" - "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" - "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=[2, 3] with " - "baseline=1.0 target=1.0005\n" - ) - - def test_with_samples(self, capsys): - info = TensorComparisonInfo( - name="samples", - baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"), - target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"), - unified_shape=[4, 8], - shape_mismatch=False, - diff=_make_diff(), - ) - - print_comparison(info=info, diff_threshold=1e-3) - - assert capsys.readouterr().out == ( - "Raw [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" - "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" - "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" - "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" - "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" - "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" - "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" - "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=[2, 3] with " - "baseline=1.0 target=1.0005\n" - "x_baseline(sample)=tensor([0.1, 0.2, ...])\n" - "x_target(sample)=tensor([0.1, 0.3, ...])\n" - ) - - def test_none_quantiles(self, capsys): - stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None) - - info = TensorComparisonInfo( - name="no_quantiles", - baseline=_make_tensor_info(stats=stats_no_quantiles), - target=_make_tensor_info(stats=stats_no_quantiles), - unified_shape=[4, 8], - shape_mismatch=False, - diff=_make_diff(), - ) - - print_comparison(info=info, diff_threshold=1e-3) - - assert capsys.readouterr().out == ( - "Raw [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "After unify [shape] [4, 8] vs [4, 8]\t" - "[dtype] torch.float32 vs torch.float32\n" - "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" - "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" - "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" - "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" - "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" - "max_abs_diff happens at coord=[2, 3] with " - "baseline=1.0 target=1.0005\n" - ) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__])) 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 d4ab7d985..5afe27053 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py @@ -46,6 +46,7 @@ def _make_diff(**overrides) -> DiffInfo: max_diff_coord=[2, 3], baseline_at_max=1.0, target_at_max=1.0005, + diff_threshold=1e-3, passed=True, ) defaults.update(overrides) @@ -76,6 +77,7 @@ class TestStrictBase: max_diff_coord=[0], baseline_at_max=0.0, target_at_max=0.0, + diff_threshold=1e-3, passed=True, extra_field=123, ) diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index 5dc896999..7cb1e1648 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -185,6 +185,7 @@ def _make_diff_info(*, passed: bool) -> DiffInfo: max_diff_coord=[0, 0], baseline_at_max=1.0, target_at_max=1.01, + diff_threshold=1e-3, passed=passed, )