From 5b28465eb974055072bbf62b63f723d44e438ef3 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:19:51 +0800 Subject: [PATCH] Cap diagnostic detail computation for failing tensors (#30656) --- .../comparator/bundle_comparator.py | 7 + .../srt/debug_utils/comparator/entrypoint.py | 4 + .../tensor_comparator/comparator.py | 66 +++++++-- .../tensor_comparator/test_comparator.py | 132 ++++++++++++++++++ .../debug_utils/comparator/test_entrypoint.py | 20 +++ 5 files changed, 221 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 8523d262a..4ae852d50 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -41,6 +41,7 @@ from sglang.srt.debug_utils.comparator.output_types import ( _split_logs, ) from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( + FailureDisplayBudget, compare_tensor_pair, compute_tensor_info, ) @@ -132,6 +133,7 @@ def compare_bundle_pair( token_aligner_mode: Optional[str], token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold_rules: Optional[list[DiffThresholdRule]] = None, + failure_display_budget: Optional[FailureDisplayBudget] = None, thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair( x=None, y=None ), @@ -147,6 +149,7 @@ def compare_bundle_pair( token_aligner_mode=token_aligner_mode, token_aligner_plan=token_aligner_plan, diff_threshold_rules=diff_threshold_rules, + failure_display_budget=failure_display_budget, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, compute_per_token=compute_per_token, @@ -165,6 +168,7 @@ def _compare_bundle_pair_inner( token_aligner_mode: Optional[str], token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold_rules: Optional[list[DiffThresholdRule]] = None, + failure_display_budget: Optional[FailureDisplayBudget] = None, thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair( x=None, y=None ), @@ -221,6 +225,7 @@ def _compare_bundle_pair_inner( token_aligner_mode=token_aligner_mode, token_aligner_plan=token_aligner_plan, diff_threshold_rules=diff_threshold_rules, + failure_display_budget=failure_display_budget, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, compute_per_token=compute_per_token, @@ -244,6 +249,7 @@ def _compare_bundle_pair_tensor_type( token_aligner_mode: Optional[str], token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold_rules: Optional[list[DiffThresholdRule]] = None, + failure_display_budget: Optional[FailureDisplayBudget] = None, thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair( x=None, y=None ), @@ -309,6 +315,7 @@ def _compare_bundle_pair_tensor_type( x_target=aligned_target, name=name, diff_threshold_rules=diff_threshold_rules, + failure_display_budget=failure_display_budget, seq_dim=seq_dim, ) record = ComparisonTensorRecord( diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 1ceb464b9..941803814 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -41,6 +41,7 @@ from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset from sglang.srt.debug_utils.comparator.report_sink import report_sink from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( DEFAULT_PREDICATE, + FailureDisplayBudget, ) from sglang.srt.debug_utils.comparator.threshold_dsl import ( DiffThresholdRule, @@ -150,6 +151,7 @@ def run(args: argparse.Namespace) -> int: diff_threshold_rules=parse_diff_threshold_rules( args.diff_threshold, default_predicate=DEFAULT_PREDICATE ), + failure_display_budget=FailureDisplayBudget(), thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, compute_per_token=visualize_per_token is not None, @@ -230,6 +232,7 @@ def _compare_bundle_pairs( token_aligner_mode: Optional[str], token_aligner_plan: Optional[TokenAlignerPlan], diff_threshold_rules: Optional[list[DiffThresholdRule]] = None, + failure_display_budget: Optional[FailureDisplayBudget] = None, thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]], viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, @@ -265,6 +268,7 @@ def _compare_bundle_pairs( token_aligner_mode=token_aligner_mode, token_aligner_plan=token_aligner_plan, diff_threshold_rules=diff_threshold_rules, + failure_display_budget=failure_display_budget, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, compute_per_token=compute_per_token, 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 ec9fc40b2..5fa5b6772 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from typing import Optional import torch @@ -30,11 +31,30 @@ SAMPLE_DIFF_THRESHOLD = 1e-3 DEFAULT_PREDICATE: str = "rel <= 0.001" +@dataclass +class FailureDisplayBudget: + max_detail: int = 50 + num_emitted: int = 0 + + def take(self) -> bool: + if self.max_detail < 0: + return True + if self.num_emitted >= self.max_detail: + return False + self.num_emitted += 1 + return True + + def compute_tensor_info( - tensor: torch.Tensor, *, include_sample: bool = False + tensor: torch.Tensor, + *, + include_sample: bool = False, + include_percentiles: bool = True, ) -> TensorInfo: """Compute TensorInfo (shape, dtype, stats, optional sample) for a single tensor.""" - stats: TensorStats = _compute_tensor_stats(tensor.float()) + stats: TensorStats = _compute_tensor_stats( + tensor.float(), include_percentiles=include_percentiles + ) sample: Optional[str] = ( str(get_truncated_value(tensor.float())) if include_sample else None ) @@ -52,14 +72,13 @@ def compare_tensor_pair( name: str = "", diff_threshold_rules: Optional[list[DiffThresholdRule]] = None, seq_dim: Optional[int] = None, + failure_display_budget: Optional[FailureDisplayBudget] = None, ) -> TensorComparisonInfo: predicate = resolve_predicate( name, diff_threshold_rules, default_predicate=DEFAULT_PREDICATE ) - baseline_info: TensorInfo = compute_tensor_info(x_baseline) - target_info: TensorInfo = compute_tensor_info(x_target) - + x_baseline_original = x_baseline x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape) unified_shape = list(x_baseline.shape) @@ -81,8 +100,31 @@ def compare_tensor_pair( x_target=x_target_f, predicate=predicate, seq_dim=seq_dim, + include_percentiles=False, ) + is_failure = shape_mismatch or (diff is not None and not diff.passed) + needs_detail = is_failure and ( + failure_display_budget is None or failure_display_budget.take() + ) + + baseline_info: TensorInfo = compute_tensor_info( + x_baseline_original, include_percentiles=needs_detail + ) + target_info: TensorInfo = compute_tensor_info( + x_target, include_percentiles=needs_detail + ) + + if not shape_mismatch and needs_detail: + diff = compute_diff( + x_baseline=x_baseline_f, + x_target=x_target_f, + predicate=predicate, + seq_dim=seq_dim, + include_percentiles=True, + ) + + if diff is not None: needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD if needs_sample: baseline_info.sample = str(get_truncated_value(x_baseline_f)) @@ -97,6 +139,7 @@ def compare_tensor_pair( x_baseline=x_baseline_f.to(downcast_dtype), x_target=x_target_f.to(downcast_dtype), predicate=predicate, + include_percentiles=needs_detail, ) return TensorComparisonInfo( @@ -111,7 +154,9 @@ def compare_tensor_pair( ) -def _compute_tensor_stats(x: torch.Tensor) -> TensorStats: +def _compute_tensor_stats( + x: torch.Tensor, *, include_percentiles: bool = True +) -> TensorStats: if x.numel() == 0: return TensorStats( mean=0.0, @@ -122,7 +167,9 @@ def _compute_tensor_stats(x: torch.Tensor) -> TensorStats: percentiles={}, ) - include_quantiles: bool = x.numel() < QUANTILE_NUMEL_THRESHOLD + include_quantiles: bool = ( + include_percentiles and x.numel() < QUANTILE_NUMEL_THRESHOLD + ) return TensorStats( mean=torch.mean(x).item(), abs_mean=torch.mean(x.abs()).item(), @@ -148,6 +195,7 @@ def compute_diff( x_target: torch.Tensor, predicate: str = DEFAULT_PREDICATE, seq_dim: Optional[int] = None, + include_percentiles: bool = True, ) -> DiffInfo: if x_baseline.numel() == 0: return DiffInfo( @@ -171,7 +219,9 @@ def compute_diff( ) mean_abs_diff = raw_abs_diff.mean().item() - include_quantiles: bool = raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD + include_quantiles: bool = ( + include_percentiles and raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD + ) per_token_rel_diff: Optional[list[float]] = None if seq_dim is not None and x_baseline.dim() > seq_dim: 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 e9e8083b8..dd7aff671 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py @@ -6,6 +6,7 @@ import torch from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( QUANTILE_NUMEL_THRESHOLD, SAMPLE_DIFF_THRESHOLD, + FailureDisplayBudget, _compute_tensor_stats, compare_tensor_pair, compute_diff, @@ -542,5 +543,136 @@ class TestCompareTensorPairPredicate: assert info.diff_downcast.predicate == "rel <= 0.0085 or max_abs <= 1e-4" +class TestFailureDisplayBudget: + @staticmethod + def _failing_pair() -> tuple[torch.Tensor, torch.Tensor]: + return torch.zeros(4, 4), torch.ones(4, 4) + + def test_take_grants_exactly_max_detail_units(self) -> None: + """take() grants exactly max_detail units, then denies.""" + budget = FailureDisplayBudget(max_detail=2) + assert budget.take() is True + assert budget.take() is True + assert budget.take() is False + assert budget.num_emitted == 2 + + def test_negative_max_detail_disables_limit(self) -> None: + """A negative max_detail always grants detail and never counts.""" + budget = FailureDisplayBudget(max_detail=-1) + assert all(budget.take() for _ in range(100)) + assert budget.num_emitted == 0 + + def test_no_budget_always_emits_detail(self) -> None: + """Without a budget object, every failing comparison gets full detail.""" + x, y = self._failing_pair() + for _ in range(3): + info = compare_tensor_pair(x_baseline=x, x_target=y, name="fail") + assert info.diff is not None and 50 in info.diff.abs_diff_percentiles + + def test_passing_tensor_skips_percentiles_and_budget(self) -> None: + """A passing comparison emits no percentile detail and consumes no budget.""" + budget = FailureDisplayBudget() + x = torch.randn(4, 4) + info = compare_tensor_pair( + x_baseline=x, x_target=x.clone(), name="pass", failure_display_budget=budget + ) + assert info.diff is not None and info.diff.passed is True + assert info.diff.abs_diff_percentiles == {} + assert info.baseline.stats.percentiles == {} + assert info.target.stats.percentiles == {} + assert budget.num_emitted == 0 + + def test_failing_tensor_within_budget_has_percentiles(self) -> None: + """A failing comparison within budget emits stats and diff percentiles.""" + budget = FailureDisplayBudget() + x, y = self._failing_pair() + info = compare_tensor_pair( + x_baseline=x, x_target=y, name="fail", failure_display_budget=budget + ) + assert info.diff is not None and info.diff.passed is False + assert 50 in info.diff.abs_diff_percentiles + assert 50 in info.baseline.stats.percentiles + assert 50 in info.target.stats.percentiles + assert budget.num_emitted == 1 + + def test_failing_tensor_beyond_budget_keeps_verdict_drops_detail(self) -> None: + """A failing comparison over budget keeps its verdict and metrics but drops percentiles.""" + budget = FailureDisplayBudget(max_detail=1) + x, y = self._failing_pair() + first = compare_tensor_pair( + x_baseline=x, x_target=y, name="first", failure_display_budget=budget + ) + second = compare_tensor_pair( + x_baseline=x, x_target=y, name="second", failure_display_budget=budget + ) + assert first.diff is not None and 50 in first.diff.abs_diff_percentiles + assert second.diff is not None and second.diff.passed is False + assert second.diff.max_abs_diff == pytest.approx(1.0) + assert second.diff.abs_diff_percentiles == {} + assert second.baseline.stats.percentiles == {} + assert second.target.stats.percentiles == {} + + def test_shape_mismatch_consumes_budget_and_emits_stats_detail(self) -> None: + """A shape mismatch counts as a failure and gets full stats percentile detail.""" + budget = FailureDisplayBudget(max_detail=1) + info = compare_tensor_pair( + x_baseline=torch.randn(3, 4), + x_target=torch.randn(5, 6), + name="mismatch", + failure_display_budget=budget, + ) + assert info.shape_mismatch is True and info.diff is None + assert 50 in info.baseline.stats.percentiles + assert 50 in info.target.stats.percentiles + assert budget.num_emitted == 1 + + def test_sample_still_emitted_beyond_budget(self) -> None: + """Sample emission for large diffs is independent of the detail budget.""" + budget = FailureDisplayBudget(max_detail=0) + x, y = self._failing_pair() + info = compare_tensor_pair( + x_baseline=x, x_target=y, name="big", failure_display_budget=budget + ) + assert info.baseline.sample is not None + assert info.target.sample is not None + + def test_downcast_diff_detail_follows_budget(self) -> None: + """The downcast diff carries percentiles only when the failure is within budget.""" + x = torch.zeros(4, 4, dtype=torch.float32) + y = torch.ones(4, 4, dtype=torch.bfloat16) + within = compare_tensor_pair( + x_baseline=x, + x_target=y, + name="within", + failure_display_budget=FailureDisplayBudget(), + ) + assert within.diff_downcast is not None + assert 50 in within.diff_downcast.abs_diff_percentiles + beyond = compare_tensor_pair( + x_baseline=x, + x_target=y, + name="beyond", + failure_display_budget=FailureDisplayBudget(max_detail=0), + ) + assert beyond.diff_downcast is not None + assert beyond.diff_downcast.abs_diff_percentiles == {} + + def test_compute_diff_include_percentiles_flag(self) -> None: + """compute_diff with include_percentiles=False omits abs_diff_percentiles.""" + x, y = self._failing_pair() + without = compute_diff(x_baseline=x, x_target=y, include_percentiles=False) + with_detail = compute_diff(x_baseline=x, x_target=y, include_percentiles=True) + assert without.abs_diff_percentiles == {} + assert 50 in with_detail.abs_diff_percentiles + assert without.passed == with_detail.passed + assert without.rel_diff == with_detail.rel_diff + + def test_compute_tensor_info_include_percentiles_flag(self) -> None: + """compute_tensor_info with include_percentiles=False omits stats percentiles.""" + t = torch.randn(16) + assert compute_tensor_info(t, include_percentiles=False).stats.percentiles == {} + assert 50 in compute_tensor_info(t, include_percentiles=True).stats.percentiles + + 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 9db9874c5..02c7de002 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -5241,5 +5241,25 @@ class TestDiffThresholdPredicateExitCode: assert exit_code == 1 +class TestFailureDisplayBudgetWiring: + def test_run_emits_full_detail_for_failing_tensor(self, tmp_path, capsys) -> None: + """run() builds a fresh per-run budget, so failing tensors carry percentile detail.""" + ones = torch.ones(4, 4) + baseline = _create_rank_dump( + tmp_path / "baseline", rank=0, name="g", tensor=ones + ) + target = _create_rank_dump( + tmp_path / "target", rank=0, name="g", tensor=ones * 2 + ) + + records, exit_code = _run_and_parse(_make_argv(baseline, target), capsys) + + tensors = [r for r in records if isinstance(r, ComparisonTensorRecord)] + assert len(tensors) == 1 + assert tensors[0].diff is not None and tensors[0].diff.passed is False + assert len(tensors[0].diff.abs_diff_percentiles) > 0 + assert exit_code == 1 + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))