minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support per-regex diff-threshold predicates in the tensor comparator (#30654)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-09 20:15:36 +08:00
committed by GitHub
parent 4153410477
commit 0d7e8cfb85
17 changed files with 799 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py
+2 -2
View File
@@ -133,9 +133,9 @@ def _check_replicated_pair(
diff_info = compute_diff(
x_baseline=baseline,
x_target=other_float,
diff_threshold=_REPLICATED_ATOL,
predicate=f"max_abs <= {_REPLICATED_ATOL}",
)
passed = diff_info.max_abs_diff <= _REPLICATED_ATOL
passed = diff_info.passed
return ReplicatedCheckResult(
axis=axis.value,
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
+7 -6
View File
@@ -44,6 +44,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compare_tensor_pair,
compute_tensor_info,
)
from sglang.srt.debug_utils.comparator.threshold_dsl import DiffThresholdRule
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
@@ -130,7 +131,7 @@ def compare_bundle_pair(
dir_pair: Pair[Path],
token_aligner_mode: Optional[str],
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=None, y=None
),
@@ -145,7 +146,7 @@ def compare_bundle_pair(
dir_pair=dir_pair,
token_aligner_mode=token_aligner_mode,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
diff_threshold_rules=diff_threshold_rules,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=compute_per_token,
@@ -163,7 +164,7 @@ def _compare_bundle_pair_inner(
dir_pair: Pair[Path],
token_aligner_mode: Optional[str],
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=None, y=None
),
@@ -219,7 +220,7 @@ def _compare_bundle_pair_inner(
valid_pair=all_pair,
token_aligner_mode=token_aligner_mode,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
diff_threshold_rules=diff_threshold_rules,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=compute_per_token,
@@ -242,7 +243,7 @@ def _compare_bundle_pair_tensor_type(
valid_pair: Pair[list[ValueWithMeta]],
token_aligner_mode: Optional[str],
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=None, y=None
),
@@ -307,7 +308,7 @@ def _compare_bundle_pair_tensor_type(
x_baseline=aligned_baseline,
x_target=aligned_target,
name=name,
diff_threshold=diff_threshold,
diff_threshold_rules=diff_threshold_rules,
seq_dim=seq_dim,
)
record = ComparisonTensorRecord(
python/sglang/srt/debug_utils/comparator/entrypoint.py
+24 -4
View File
@@ -39,6 +39,13 @@ from sglang.srt.debug_utils.comparator.per_token_visualizer import (
)
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,
)
from sglang.srt.debug_utils.comparator.threshold_dsl import (
DiffThresholdRule,
parse_diff_threshold_rules,
)
from sglang.srt.debug_utils.comparator.utils import (
Pair,
auto_descend_dir,
@@ -140,7 +147,9 @@ def run(args: argparse.Namespace) -> int:
dir_pair=dir_pair,
token_aligner_mode=ta_result.mode,
token_aligner_plan=ta_result.plan,
diff_threshold=args.diff_threshold,
diff_threshold_rules=parse_diff_threshold_rules(
args.diff_threshold, default_predicate=DEFAULT_PREDICATE
),
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,
@@ -220,7 +229,7 @@ def _compare_bundle_pairs(
dir_pair: Pair[Path],
token_aligner_mode: Optional[str],
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]],
viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False,
@@ -255,7 +264,7 @@ def _compare_bundle_pairs(
dir_pair=dir_pair,
token_aligner_mode=token_aligner_mode,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
diff_threshold_rules=diff_threshold_rules,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=compute_per_token,
@@ -331,7 +340,18 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
parser.add_argument("--target-path", type=str)
parser.add_argument("--start-step", type=int, default=0)
parser.add_argument("--end-step", type=int, default=1000000)
parser.add_argument("--diff-threshold", type=float, default=1e-3)
parser.add_argument(
"--diff-threshold",
nargs="*",
default=None,
metavar="REGEX PREDICATE",
help="Per-tensor pass criterion. Either a single float shorthand "
"(0.0085 == '.*' 'rel <= 0.0085'), or (regex predicate) pairs, e.g. "
"--diff-threshold '.*expert.*' 'rel <= 0.0085 or max_abs <= 1e-3' '.*' 'rel <= 0.0085'. "
"A tensor uses the first fullmatching regex's predicate -- a boolean expression "
"over rel/max_abs/mean_abs with < <= > >= and and/or. A tensor matching no "
"pattern is an error. Default: 'rel <= 1e-3' for every tensor.",
)
parser.add_argument(
"--filter", type=str, default=None, help="Regex to filter filenames (include)"
)
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
+23 -7
View File
@@ -9,6 +9,12 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
TensorInfo,
TensorStats,
)
from sglang.srt.debug_utils.comparator.threshold_dsl import (
DiffThresholdRule,
evaluate_predicate,
parse_predicate,
resolve_predicate,
)
from sglang.srt.debug_utils.comparator.utils import (
Pair,
argmax_coord,
@@ -21,6 +27,7 @@ from sglang.srt.debug_utils.dumper import get_truncated_value
QUANTILE_NUMEL_THRESHOLD = 10_000_000
SAMPLE_DIFF_THRESHOLD = 1e-3
DEFAULT_PREDICATE: str = "rel <= 0.001"
def compute_tensor_info(
@@ -43,9 +50,13 @@ def compare_tensor_pair(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
name: str = "",
diff_threshold: float = 1e-3,
diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
seq_dim: Optional[int] = 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)
@@ -68,7 +79,7 @@ def compare_tensor_pair(
diff = compute_diff(
x_baseline=x_baseline_f,
x_target=x_target_f,
diff_threshold=diff_threshold,
predicate=predicate,
seq_dim=seq_dim,
)
@@ -85,7 +96,7 @@ def compare_tensor_pair(
diff_downcast = compute_diff(
x_baseline=x_baseline_f.to(downcast_dtype),
x_target=x_target_f.to(downcast_dtype),
diff_threshold=diff_threshold,
predicate=predicate,
)
return TensorComparisonInfo(
@@ -135,7 +146,7 @@ def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
def compute_diff(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
diff_threshold: float = 1e-3,
predicate: str = DEFAULT_PREDICATE,
seq_dim: Optional[int] = None,
) -> DiffInfo:
if x_baseline.numel() == 0:
@@ -147,7 +158,7 @@ def compute_diff(
max_diff_coord=[],
baseline_at_max=0.0,
target_at_max=0.0,
diff_threshold=diff_threshold,
predicate=predicate,
passed=True,
)
@@ -176,7 +187,12 @@ 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,
predicate=predicate,
passed=evaluate_predicate(
parse_predicate(predicate),
rel=rel_diff,
max_abs=max_abs_diff,
mean_abs=mean_abs_diff,
),
per_token_rel_diff=per_token_rel_diff,
)
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
+2 -2
View File
@@ -183,10 +183,10 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list
def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
rel_diff_marker: str = "❌" if diff.rel_diff > diff.diff_threshold else "✅"
marker: str = "✅" if diff.passed else "❌"
lines: list[str] = [
prefix_text
+ f"{rel_diff_marker} rel_diff={diff.rel_diff}\t"
+ f"{marker} rel_diff={diff.rel_diff}\t"
+ f"max_abs_diff={diff.max_abs_diff}\t"
+ f"mean_abs_diff={diff.mean_abs_diff}",
f"max_abs_diff happens at coord={diff.max_diff_coord} with "
python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py
+1 -1
View File
@@ -29,7 +29,7 @@ class DiffInfo(_StrictBase):
max_diff_coord: list[int]
baseline_at_max: float
target_at_max: float
diff_threshold: float
predicate: str = ""
passed: bool
per_token_rel_diff: Optional[list[float]] = None
python/sglang/srt/debug_utils/comparator/threshold_dsl.py
+84
View File
@@ -0,0 +1,84 @@
import re
from dataclasses import dataclass
from functools import lru_cache
from types import CodeType
from typing import Optional
ALLOWED_NAMES: tuple[str, ...] = ("rel", "max_abs", "mean_abs")
_EVAL_GLOBALS: dict = {"__builtins__": {}}
_DUMMY_ENV: dict[str, float] = {name: 1.0 for name in ALLOWED_NAMES}
@dataclass(frozen=True)
class DiffThresholdRule:
pattern: str
predicate: str
def parse_diff_threshold_rules(
raw: Optional[list[str]], *, default_predicate: str
) -> list[DiffThresholdRule]:
if not raw:
return [DiffThresholdRule(".*", default_predicate)]
if len(raw) == 1:
try:
value = float(raw[0])
except ValueError as e:
raise ValueError(
f"--diff-threshold with a single argument must be a float shorthand "
f"(e.g. 0.0085); got {raw[0]!r}. For per-regex predicates pass "
f"(regex predicate) pairs."
) from e
return [DiffThresholdRule(".*", f"rel <= {value}")]
if len(raw) % 2 != 0:
raise ValueError(
f"--diff-threshold expects a single float shorthand or (regex predicate) "
f"pairs; got an odd number of arguments: {raw}"
)
rules = [DiffThresholdRule(raw[i], raw[i + 1]) for i in range(0, len(raw), 2)]
for rule in rules:
parse_predicate(rule.predicate)
return rules
def resolve_predicate(
name: str,
diff_threshold_rules: Optional[list[DiffThresholdRule]],
*,
default_predicate: str,
) -> str:
if not diff_threshold_rules:
return default_predicate
for rule in diff_threshold_rules:
if re.fullmatch(rule.pattern, name):
return rule.predicate
raise ValueError(
f"tensor {name!r} matched no --diff-threshold pattern "
f"({[rule.pattern for rule in diff_threshold_rules]}); add a catch-all '.*' rule or a matching pattern."
)
@lru_cache(maxsize=None)
def parse_predicate(expr: str) -> CodeType:
try:
code = compile(expr, "<predicate>", "eval")
except SyntaxError as e:
raise ValueError(f"invalid predicate {expr!r}: {e}") from e
try:
eval(code, _EVAL_GLOBALS, dict(_DUMMY_ENV))
except Exception as e:
raise ValueError(
f"invalid predicate {expr!r}: {e}; allowed names are {ALLOWED_NAMES}."
) from e
return code
def evaluate_predicate(
code: CodeType, *, rel: float, max_abs: float, mean_abs: float
) -> bool:
return bool(
eval(
code, _EVAL_GLOBALS, {"rel": rel, "max_abs": max_abs, "mean_abs": mean_abs}
)
)
test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
+159 -7
View File
@@ -12,6 +12,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compute_tensor_info,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import DiffInfo
from sglang.srt.debug_utils.comparator.threshold_dsl import DiffThresholdRule
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=20, suite="base-a-test-cpu", nightly=True)
@@ -269,9 +270,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, seq_dim=0
)
diff: DiffInfo = compute_diff(x_baseline=x, x_target=y, seq_dim=0)
assert diff.per_token_rel_diff is not None
assert isinstance(diff.per_token_rel_diff, list)
@@ -283,7 +282,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)
assert diff.per_token_rel_diff is None
@@ -293,9 +292,7 @@ class TestComputeDiff:
x: torch.Tensor = torch.randn(4, 8)
y: torch.Tensor = x + torch.randn_like(x) * 0.01
diff: DiffInfo = compute_diff(
x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0
)
diff: DiffInfo = compute_diff(x_baseline=x, x_target=y, seq_dim=0)
json_str: str = diff.model_dump_json()
assert "per_token_rel_diff" in json_str
@@ -376,5 +373,160 @@ class TestCompareTensors:
assert info.target.sample is None
class TestComputeDiffPredicate:
@staticmethod
def _near_zero_pair() -> tuple[torch.Tensor, torch.Tensor]:
"""Sign-flipped near-zero pair: rel_diff == 2.0 but max_abs/mean_abs == 2e-5."""
x = torch.tensor([1e-5, -1e-5, 1e-5, -1e-5])
return x, -x
def test_default_predicate(self) -> None:
"""No predicate → the default 'rel <= 0.001'; near-zero pair fails and the string is recorded."""
x, y = self._near_zero_pair()
diff = compute_diff(x_baseline=x, x_target=y)
assert diff.rel_diff == pytest.approx(2.0, abs=1e-4)
assert diff.max_abs_diff == pytest.approx(2e-5, abs=1e-7)
assert diff.predicate == "rel <= 0.001"
assert diff.passed is False
def test_predicate_rescues_near_zero_via_max_abs(self) -> None:
"""A 'rel or max_abs' predicate passes the near-zero pair despite a failing rel."""
x, y = self._near_zero_pair()
diff = compute_diff(
x_baseline=x, x_target=y, predicate="rel <= 0.0085 or max_abs <= 1e-4"
)
assert diff.rel_diff > 1.0 # relative term still fails
assert diff.passed is True
assert diff.predicate == "rel <= 0.0085 or max_abs <= 1e-4"
def test_predicate_does_not_rescue_real_magnitude_diff(self) -> None:
"""A real-magnitude diff fails both terms of a 'rel or max_abs' predicate."""
x = torch.ones(10)
y = x.clone()
y[0] = 2.0 # max_abs_diff == 1.0, rel_diff ~0.043
diff = compute_diff(
x_baseline=x, x_target=y, predicate="rel <= 0.0085 or max_abs <= 1e-3"
)
assert diff.max_abs_diff == pytest.approx(1.0, abs=1e-4)
assert diff.passed is False
def test_and_predicate_requires_both(self) -> None:
"""'rel and max_abs' fails the near-zero pair (rel huge) but passes a small-both diff."""
x, y = self._near_zero_pair()
assert (
compute_diff(
x_baseline=x, x_target=y, predicate="rel <= 0.0085 and max_abs <= 1e-4"
).passed
is False
)
small = torch.ones(10)
small_y = small + 1e-4 # rel ~5e-9, max_abs 1e-4
assert (
compute_diff(
x_baseline=small,
x_target=small_y,
predicate="rel <= 0.0085 and max_abs <= 1e-3",
).passed
is True
)
def test_mean_abs_variable(self) -> None:
"""A mean_abs predicate uses the mean absolute diff (2e-5 for the near-zero pair)."""
x, y = self._near_zero_pair()
assert (
compute_diff(x_baseline=x, x_target=y, predicate="mean_abs <= 1e-4").passed
is True
)
assert (
compute_diff(x_baseline=x, x_target=y, predicate="mean_abs <= 1e-6").passed
is False
)
def test_boundary_le_inclusive(self) -> None:
"""<= includes the boundary, < excludes it (max_abs_diff == 0.5 exactly)."""
x = torch.tensor([1.0, 1.0])
y = torch.tensor([1.5, 1.5])
assert (
compute_diff(x_baseline=x, x_target=y, predicate="max_abs <= 0.5").passed
is True
)
assert (
compute_diff(x_baseline=x, x_target=y, predicate="max_abs < 0.5").passed
is False
)
def test_predicate_recorded_for_empty_tensor(self) -> None:
"""Empty tensors short-circuit to passed=True and still record the predicate."""
empty = torch.empty(0)
diff = compute_diff(x_baseline=empty, x_target=empty, predicate="rel <= 0")
assert diff.passed is True
assert diff.predicate == "rel <= 0"
def test_predicate_json_roundtrip(self) -> None:
"""DiffInfo.predicate survives JSON serialization."""
x, y = self._near_zero_pair()
diff = compute_diff(
x_baseline=x, x_target=y, predicate="rel <= 0.0085 or max_abs <= 1e-4"
)
roundtripped = DiffInfo.model_validate_json(diff.model_dump_json())
assert roundtripped.predicate == "rel <= 0.0085 or max_abs <= 1e-4"
assert roundtripped.passed is True
class TestCompareTensorPairPredicate:
def test_predicate_resolved_per_name(self) -> None:
"""compare_tensor_pair resolves the per-regex predicate by tensor name into the verdict."""
x = torch.tensor([1e-5, -1e-5, 1e-5, -1e-5])
y = -x
without = compare_tensor_pair(x_baseline=x, x_target=y, name="g.expert.0")
assert without.diff is not None
assert without.diff.passed is False
with_pred = compare_tensor_pair(
x_baseline=x,
x_target=y,
name="g.expert.0",
diff_threshold_rules=[
DiffThresholdRule(".*expert.*", "rel <= 0.0085 or max_abs <= 1e-4")
],
)
assert with_pred.diff is not None
assert with_pred.diff.passed is True
assert with_pred.diff.predicate == "rel <= 0.0085 or max_abs <= 1e-4"
def test_unmatched_name_raises(self) -> None:
"""A tensor matching no pattern raises (fail-closed)."""
x = torch.tensor([1e-5, -1e-5, 1e-5, -1e-5])
with pytest.raises(ValueError, match="matched no --diff-threshold pattern"):
compare_tensor_pair(
x_baseline=x,
x_target=-x,
name="g.attn.qkv",
diff_threshold_rules=[
DiffThresholdRule(".*expert.*", "rel <= 0.0085 or max_abs <= 1e-4")
],
)
def test_predicate_propagates_to_downcast_diff(self) -> None:
"""When baseline/target dtypes differ, the resolved predicate also drives the downcast diff."""
x = torch.tensor([1e-5, -1e-5, 1e-5, -1e-5])
info = compare_tensor_pair(
x_baseline=x,
x_target=(-x).to(torch.bfloat16),
name="g.expert.0",
diff_threshold_rules=[
DiffThresholdRule(".*expert.*", "rel <= 0.0085 or max_abs <= 1e-4")
],
)
assert info.diff_downcast is not None
assert info.diff_downcast.predicate == "rel <= 0.0085 or max_abs <= 1e-4"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
+18
View File
@@ -1,4 +1,9 @@
import sys
from pathlib import Path
_TEST_ROOT: Path = Path(__file__).resolve().parents[4]
if str(_TEST_ROOT) not in sys.path:
sys.path.insert(0, str(_TEST_ROOT))
import pytest
from registered.debug_utils.comparator.testing_helpers import (
@@ -117,6 +122,19 @@ class TestFormatComparison:
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
def test_marker_follows_passed_not_rel_threshold(self):
"""The ✅/❌ marker reflects diff.passed (the predicate verdict), not rel_diff vs a threshold."""
info = TensorComparisonInfo(
name="rescued",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(rel_diff=2.0, passed=True),
)
assert "✅ rel_diff=2.0\t" in format_comparison(info)
def test_shape_mismatch(self):
info = TensorComparisonInfo(
name="mismatch",
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
-3
View File
@@ -46,7 +46,6 @@ 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)
@@ -77,7 +76,6 @@ 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,
)
@@ -140,7 +138,6 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
rel_diff=0.1,
max_abs_diff=0.1,
mean_abs_diff=0.05,
diff_threshold=1e-6,
passed=False,
),
)
test/registered/debug_utils/comparator/test_entrypoint.py
+253 -3
View File
@@ -2709,7 +2709,8 @@ def _make_argv(
preset: str | None = None,
grouping_skip_keys: list[str] | None = None,
token_aligner: str | None = None,
diff_threshold: float = 1e-3,
diff_threshold: float | None = 1e-3,
diff_thresholds: list[tuple[str, str]] | None = None,
output_format: str = "json",
start_step: int | None = None,
end_step: int | None = None,
@@ -2730,8 +2731,6 @@ def _make_argv(
str(baseline_path),
"--target-path",
str(target_path),
"--diff-threshold",
str(diff_threshold),
"--output-format",
output_format,
]
@@ -2768,6 +2767,12 @@ def _make_argv(
argv += ["--viz-output-dir", viz_output_dir]
if visualize_per_token is not None:
argv += ["--visualize-per-token", visualize_per_token]
if diff_thresholds is not None:
argv.append("--diff-threshold")
for pattern, predicate in diff_thresholds:
argv += [pattern, predicate]
elif diff_threshold is not None:
argv += ["--diff-threshold", str(diff_threshold)]
return argv
@@ -4991,5 +4996,250 @@ class TestErrorResilience:
assert hint_pos < traceback_pos
class TestDiffThresholdCliParsing:
def test_parse_args_collects_diff_threshold_tokens(self) -> None:
"""parse_args leaves --diff-threshold as raw tokens (None when the flag is absent)."""
argv = ["--baseline-path", "b", "--target-path", "t"]
assert parse_args(argv).diff_threshold is None
argv += ["--diff-threshold", ".*expert.*", "rel <= 0.0085 or max_abs <= 1e-3"]
assert parse_args(argv).diff_threshold == [
".*expert.*",
"rel <= 0.0085 or max_abs <= 1e-3",
]
class TestDiffThresholdPredicateExitCode:
"""End-to-end: a per-regex --diff-threshold predicate drives the per-tensor verdict and exit code."""
@staticmethod
def _dump_near_zero_pair(tmp_path: Path) -> tuple[Path, Path]:
baseline_t = torch.tensor([[1e-5, -1e-5], [1e-5, -1e-5]])
baseline = _create_rank_dump(
tmp_path / "baseline", rank=0, name="g", tensor=baseline_t
)
target = _create_rank_dump(
tmp_path / "target", rank=0, name="g", tensor=-baseline_t
)
return baseline, target
def test_default_predicate_fails_near_zero(self, tmp_path, capsys) -> None:
"""The default 'rel <= X' predicate fails the near-zero pair (exit code 1)."""
baseline, target = self._dump_near_zero_pair(tmp_path)
argv = _make_argv(baseline, target, diff_threshold=0.0085)
records, exit_code = _run_and_parse(argv, 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 tensors[0].diff.predicate == "rel <= 0.0085"
assert exit_code == 1
def test_predicate_passes_near_zero(self, tmp_path, capsys) -> None:
"""A 'rel or max_abs' predicate passes the near-zero pair (exit code 0)."""
baseline, target = self._dump_near_zero_pair(tmp_path)
argv = _make_argv(
baseline,
target,
diff_thresholds=[(".*", "rel <= 0.0085 or max_abs <= 1e-4")],
)
records, exit_code = _run_and_parse(argv, 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 True
assert tensors[0].diff.predicate == "rel <= 0.0085 or max_abs <= 1e-4"
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.failed == 0
assert exit_code == 0
def test_predicate_does_not_rescue_real_magnitude_failure(
self, tmp_path, capsys
) -> None:
"""With a co-located passing tensor, a real-magnitude divergence still fails (exit code 1)."""
ones = torch.ones(4, 4)
real_target = ones.clone()
real_target[0, 0] = 5.0 # max_abs_diff == 4.0, far above the floor
baseline = _create_rank_dump(
tmp_path / "baseline",
rank=0,
name="g_real",
tensor=ones,
extra_dumps=[("g_pass", ones)],
)
target = _create_rank_dump(
tmp_path / "target",
rank=0,
name="g_real",
tensor=real_target,
extra_dumps=[("g_pass", ones)],
)
argv = _make_argv(
baseline,
target,
diff_thresholds=[(".*", "rel <= 0.0085 or max_abs <= 1e-3")],
)
records, exit_code = _run_and_parse(argv, capsys)
by_name = {r.name: r for r in records if isinstance(r, ComparisonTensorRecord)}
assert (
by_name["g_pass"].diff is not None and by_name["g_pass"].diff.passed is True
)
assert (
by_name["g_real"].diff is not None
and by_name["g_real"].diff.passed is False
)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 1 and summary.failed == 1
assert exit_code == 1
def test_per_regex_selectivity_with_catch_all(self, tmp_path, capsys) -> None:
"""A specific predicate rescues only its name; the catch-all keeps the rest strict (exit code 1)."""
near_zero = torch.tensor([[1e-5, -1e-5], [1e-5, -1e-5]])
baseline = _create_rank_dump(
tmp_path / "baseline",
rank=0,
name="g_apple",
tensor=near_zero,
extra_dumps=[("g_orange", near_zero)],
)
target = _create_rank_dump(
tmp_path / "target",
rank=0,
name="g_apple",
tensor=-near_zero,
extra_dumps=[("g_orange", -near_zero)],
)
argv = _make_argv(
baseline,
target,
diff_thresholds=[
(".*apple.*", "rel <= 0.0085 or max_abs <= 1e-4"),
(".*", "rel <= 0.0085"),
],
)
records, exit_code = _run_and_parse(argv, capsys)
by_name = {r.name: r for r in records if isinstance(r, ComparisonTensorRecord)}
assert by_name["g_apple"].diff is not None
assert by_name["g_apple"].diff.passed is True # rescued by max_abs
assert by_name["g_orange"].diff is not None
assert by_name["g_orange"].diff.passed is False # catch-all strict rel
assert exit_code == 1
def test_unmatched_tensor_errors(self, tmp_path, capsys) -> None:
"""A tensor matching no pattern (no catch-all) becomes an error -> exit code 1 (fail-closed)."""
near_zero = torch.tensor([[1e-5, -1e-5], [1e-5, -1e-5]])
baseline = _create_rank_dump(
tmp_path / "baseline", rank=0, name="g_orange", tensor=near_zero
)
target = _create_rank_dump(
tmp_path / "target", rank=0, name="g_orange", tensor=-near_zero
)
argv = _make_argv(
baseline,
target,
diff_thresholds=[(".*apple.*", "rel <= 0.0085 or max_abs <= 1e-4")],
)
records, exit_code = _run_and_parse(argv, capsys)
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert "matched no --diff-threshold pattern" in errors[0].exception_message
assert exit_code == 1
def test_diff_threshold_flag_omitted_uses_default_predicate(
self, tmp_path, capsys
) -> None:
"""With no --diff-threshold flag, run() falls back to the default 'rel <= 0.001' rule."""
baseline, target = self._dump_near_zero_pair(tmp_path)
argv = _make_argv(baseline, target, diff_threshold=None)
assert "--diff-threshold" not in argv
records, exit_code = _run_and_parse(argv, capsys)
tensors = [r for r in records if isinstance(r, ComparisonTensorRecord)]
assert len(tensors) == 1
assert tensors[0].diff is not None
assert tensors[0].diff.predicate == "rel <= 0.001"
assert tensors[0].diff.passed is False
assert exit_code == 1
def test_miles_per_regex_selectivity_mixed_verdicts(self, tmp_path, capsys) -> None:
"""'.*apple.*' tensors get a max_abs rescue, the strict catch-all does not: an identical near-zero pair passes as apple.weight yet fails as banana.bias."""
near_zero = torch.tensor([[1e-5, -1e-5], [1e-5, -1e-5]])
ones = torch.ones(4, 4)
ones_spike = ones.clone()
ones_spike[0, 0] = 5.0
baseline = _create_rank_dump(
tmp_path / "baseline",
rank=0,
name="layer.apple.weight",
tensor=near_zero,
extra_dumps=[
("layer.apple.bias", ones),
("layer.banana.weight", ones),
("layer.banana.bias", near_zero),
],
)
target = _create_rank_dump(
tmp_path / "target",
rank=0,
name="layer.apple.weight",
tensor=-near_zero,
extra_dumps=[
("layer.apple.bias", ones_spike),
("layer.banana.weight", ones + 1e-5),
("layer.banana.bias", -near_zero),
],
)
argv = _make_argv(
baseline,
target,
diff_thresholds=[
(".*apple.*", "rel < 1e-3 or max_abs < 0.01"),
(".*", "rel < 1e-3"),
],
)
records, exit_code = _run_and_parse(argv, capsys)
by_name = {r.name: r for r in records if isinstance(r, ComparisonTensorRecord)}
assert {*by_name} == {
"layer.apple.weight",
"layer.apple.bias",
"layer.banana.weight",
"layer.banana.bias",
}
assert all(r.diff is not None for r in by_name.values())
apple_predicate = "rel < 1e-3 or max_abs < 0.01"
assert by_name["layer.apple.weight"].diff.predicate == apple_predicate
assert by_name["layer.apple.bias"].diff.predicate == apple_predicate
assert by_name["layer.banana.weight"].diff.predicate == "rel < 1e-3"
assert by_name["layer.banana.bias"].diff.predicate == "rel < 1e-3"
assert by_name["layer.apple.weight"].diff.passed is True
assert by_name["layer.apple.bias"].diff.passed is False
assert by_name["layer.banana.weight"].diff.passed is True
assert by_name["layer.banana.bias"].diff.passed is False
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 2 and summary.failed == 2
assert exit_code == 1
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/test_manually_verify.py
-2
View File
@@ -234,7 +234,6 @@ class TestPerTokenHeatmapManualVerify:
x_baseline=baseline,
x_target=target,
name=f"layer_{i}_hidden_states",
diff_threshold=1e-3,
seq_dim=0,
)
records.append(ComparisonTensorRecord(**info.model_dump()))
@@ -278,7 +277,6 @@ class TestPerTokenHeatmapManualVerify:
x_baseline=baseline,
x_target=target,
name=f"layer_{i}_attn_output",
diff_threshold=1e-3,
seq_dim=0,
)
records.append(ComparisonTensorRecord(**info.model_dump()))
test/registered/debug_utils/comparator/test_model_validation.py
-1
View File
@@ -215,7 +215,6 @@ 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,
)
test/registered/debug_utils/comparator/test_output_types.py
+5
View File
@@ -1,5 +1,10 @@
import sys
from io import StringIO
from pathlib import Path
_TEST_ROOT: Path = Path(__file__).resolve().parents[3]
if str(_TEST_ROOT) not in sys.path:
sys.path.insert(0, str(_TEST_ROOT))
import pytest
from registered.debug_utils.comparator.testing_helpers import (
test/registered/debug_utils/comparator/test_per_token_visualizer.py
-2
View File
@@ -37,7 +37,6 @@ def _make_comparison_record(
x_baseline=baseline,
x_target=target,
name=name,
diff_threshold=1e-3,
seq_dim=seq_dim,
)
return ComparisonTensorRecord(**info.model_dump())
@@ -66,7 +65,6 @@ class TestPerTokenVisualizer:
x_baseline=torch.randn(4, 8),
x_target=torch.randn(4, 8),
name="no_per_token",
diff_threshold=1e-3,
)
record = ComparisonTensorRecord(**info.model_dump())
test/registered/debug_utils/comparator/test_threshold_dsl.py
+221
View File
@@ -0,0 +1,221 @@
import sys
import pytest
from sglang.srt.debug_utils.comparator.threshold_dsl import (
DiffThresholdRule,
evaluate_predicate,
parse_diff_threshold_rules,
parse_predicate,
resolve_predicate,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-a-test-cpu", nightly=True)
def _ev(
expr: str, *, rel: float = 0.0, max_abs: float = 0.0, mean_abs: float = 0.0
) -> bool:
return evaluate_predicate(
parse_predicate(expr), rel=rel, max_abs=max_abs, mean_abs=mean_abs
)
class TestParsePredicate:
@pytest.mark.parametrize(
"expr",
[
"rel <= 0.0085",
"rel < 1",
"max_abs > 0",
"mean_abs >= 1e-5",
"rel <= 0.01 or max_abs <= 1e-4",
"rel <= 0.01 and max_abs <= 1e-4",
"(rel <= 0.01 and max_abs <= 1e-4) or mean_abs <= 1e-5",
"0 <= rel < 1",
"rel <= -0.0",
"rel <= 0",
],
)
def test_valid_predicates_parse(self, expr: str) -> None:
"""All supported forms parse without error."""
parse_predicate(expr)
@pytest.mark.parametrize(
"expr",
[
"abs(rel) < 1",
"rel.x < 1",
"foo < 1",
"rel < 'x'",
"",
"rel <",
],
)
def test_invalid_predicates_raise(self, expr: str) -> None:
"""Unknown names, attribute access, bad types, and syntax errors raise ValueError."""
with pytest.raises(ValueError):
parse_predicate(expr)
def test_unknown_name_message_lists_allowed(self) -> None:
"""The error for an unknown variable names the allowed variables."""
with pytest.raises(ValueError, match="rel.*max_abs.*mean_abs"):
parse_predicate("foo < 1")
class TestEvaluatePredicate:
def test_rel_only_true_and_false(self) -> None:
"""A pure rel predicate uses only rel."""
assert _ev("rel <= 0.01", rel=0.005) is True
assert _ev("rel <= 0.01", rel=0.02) is False
def test_le_boundary_inclusive(self) -> None:
"""<= includes the boundary; < excludes it."""
assert _ev("rel <= 0.01", rel=0.01) is True
assert _ev("rel < 0.01", rel=0.01) is False
def test_or_short_circuit_semantics(self) -> None:
"""or passes if either side holds (near-zero rescue pattern)."""
assert _ev("rel <= 0.0085 or max_abs <= 1e-3", rel=2.0, max_abs=2e-5) is True
assert _ev("rel <= 0.0085 or max_abs <= 1e-3", rel=2.0, max_abs=0.5) is False
def test_and_requires_both(self) -> None:
"""and passes only if both sides hold."""
assert _ev("rel <= 0.01 and max_abs <= 1e-3", rel=0.005, max_abs=1e-4) is True
assert _ev("rel <= 0.01 and max_abs <= 1e-3", rel=0.005, max_abs=0.5) is False
def test_mean_abs_variable(self) -> None:
"""mean_abs is a usable variable."""
assert _ev("mean_abs <= 1e-5", mean_abs=1e-6) is True
assert _ev("mean_abs <= 1e-5", mean_abs=1e-4) is False
def test_parentheses_grouping(self) -> None:
"""Parentheses override and/or precedence."""
assert (
_ev(
"(rel <= 0.01 and max_abs <= 1e-4) or mean_abs <= 1e-5",
rel=2.0,
max_abs=2.0,
mean_abs=1e-6,
)
is True
)
def test_chained_comparison(self) -> None:
"""Chained comparison follows Python all-must-hold semantics."""
assert _ev("0 <= rel < 1", rel=0.5) is True
assert _ev("0 <= rel < 1", rel=1.5) is False
def test_bitwise_zero_predicate(self) -> None:
"""rel <= 0 passes only for an exactly-zero rel (bitwise)."""
assert _ev("rel <= 0", rel=0.0) is True
assert _ev("rel <= 0", rel=1e-12) is False
class TestDiffThresholdRule:
def test_is_frozen_value_object(self) -> None:
"""DiffThresholdRule is a frozen, value-equal dataclass."""
assert DiffThresholdRule(".*", "rel <= 1e-3") == DiffThresholdRule(
".*", "rel <= 1e-3"
)
with pytest.raises(Exception):
DiffThresholdRule(".*", "rel <= 1e-3").pattern = "x"
class TestParseDiffThresholdRules:
def test_none_and_empty_return_default(self) -> None:
"""Missing flag (None) and bare flag (empty list) both yield the caller's default."""
assert parse_diff_threshold_rules(None, default_predicate="rel <= 0.001") == [
DiffThresholdRule(".*", "rel <= 0.001")
]
assert parse_diff_threshold_rules([], default_predicate="rel <= 0.001") == [
DiffThresholdRule(".*", "rel <= 0.001")
]
def test_single_float_shorthand(self) -> None:
"""A single float token expands to a global 'rel <= X' rule."""
assert parse_diff_threshold_rules(
["0.0085"], default_predicate="rel <= 0.001"
) == [DiffThresholdRule(".*", "rel <= 0.0085")]
def test_single_non_float_raises(self) -> None:
"""A single non-float token is a usage error (not a valid shorthand)."""
with pytest.raises(ValueError):
parse_diff_threshold_rules([".*expert.*"], default_predicate="rel <= 0.001")
def test_pairs_parsed_in_order(self) -> None:
"""Flat regex/predicate tokens parse into ordered rules."""
assert parse_diff_threshold_rules(
[".*apple.*", "rel <= 0.01 or max_abs <= 1e-4", ".*", "rel <= 0.0085"],
default_predicate="rel <= 0.001",
) == [
DiffThresholdRule(".*apple.*", "rel <= 0.01 or max_abs <= 1e-4"),
DiffThresholdRule(".*", "rel <= 0.0085"),
]
def test_odd_number_of_tokens_raises(self) -> None:
"""An unpaired token is a usage error."""
with pytest.raises(ValueError):
parse_diff_threshold_rules(
[".*apple.*", "rel <= 0.01", ".*orange.*"],
default_predicate="rel <= 0.001",
)
def test_bad_predicate_raises(self) -> None:
"""A malformed predicate fails fast at parse time."""
with pytest.raises(ValueError):
parse_diff_threshold_rules(
[".*", "rel <= "], default_predicate="rel <= 0.001"
)
class TestResolvePredicate:
def test_none_or_empty_returns_explicit_default(self) -> None:
"""No rules → the supplied default predicate."""
assert (
resolve_predicate("x", None, default_predicate="rel <= 1e-3")
== "rel <= 1e-3"
)
assert (
resolve_predicate("x", [], default_predicate="rel <= 1e-3") == "rel <= 1e-3"
)
def test_first_matching_pattern_wins(self) -> None:
"""Rules are tried in order; the first fullmatch wins (specific before general)."""
rules = [DiffThresholdRule(".*expert.*", "P1"), DiffThresholdRule(".*", "P2")]
assert (
resolve_predicate("layer.expert.weight", rules, default_predicate="D")
== "P1"
)
assert (
resolve_predicate("layer.attn.weight", rules, default_predicate="D") == "P2"
)
def test_unmatched_name_raises(self) -> None:
"""A name matching no pattern is a fail-closed error."""
with pytest.raises(ValueError, match="matched no --diff-threshold pattern"):
resolve_predicate(
"k_layernorm",
[DiffThresholdRule(".*expert.*", "P")],
default_predicate="D",
)
def test_fullmatch_semantics(self) -> None:
"""Matching is fullmatch: a partial pattern does not match a longer name (→ raises)."""
assert (
resolve_predicate(
"expert", [DiffThresholdRule("expert", "P")], default_predicate="D"
)
== "P"
)
with pytest.raises(ValueError):
resolve_predicate(
"layer.expert.weight",
[DiffThresholdRule("expert", "P")],
default_predicate="D",
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/testing_helpers.py
-2
View File
@@ -63,7 +63,6 @@ def make_diff(
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
abs_diff_percentiles: Optional[dict[int, float]] = None,
diff_threshold: float = 1e-3,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
@@ -78,7 +77,6 @@ def make_diff(
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=diff_threshold,
passed=passed,
)