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}
)
)