chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support agent-friendly output formats in dump comparator (#19275)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-25 09:40:33 +08:00
committed by GitHub
parent d7578ce279
commit 8b1ab4aaf9
11 changed files with 730 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
python/sglang/srt/debug_utils/comparator/entrypoint.py
View File

@@ -3,13 +3,16 @@ from pathlib import Path
import polars as pl
from sglang.srt.debug_utils.comparator.tensor_comparison import (
compare_tensors,
print_comparison,
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
ConfigRecord,
SkipRecord,
SummaryRecord,
print_record,
)
from sglang.srt.debug_utils.comparator.tensor_comparison import compare_tensors
from sglang.srt.debug_utils.comparator.utils import load_object
from sglang.srt.debug_utils.dump_loader import find_row, read_meta
from sglang.srt.debug_utils.dumper import get_truncated_value
def main() -> None:
@@ -27,8 +30,19 @@ def run(args: argparse.Namespace) -> None:
assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
df_baseline = read_meta(args.baseline_path)
print("df_target", df_target)
print("df_baseline", df_baseline)
print_record(
ConfigRecord(
baseline_path=args.baseline_path,
target_path=args.target_path,
diff_threshold=args.diff_threshold,
start_step=args.start_step,
end_step=args.end_step,
),
output_format=args.output_format,
)
counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
for row in df_target.iter_rows(named=True):
path_target = Path(args.target_path) / row["filename"]
@@ -47,26 +61,23 @@ def run(args: argparse.Namespace) -> None:
)
if row_baseline is None:
print(f"Skip: target={str(path_target)} since no baseline")
x_target = load_object(path_target)
if x_target is not None:
print(f"x_target(sample)={get_truncated_value(x_target)}")
counts["skipped"] += 1
print_record(
SkipRecord(name=row["name"], reason="no_baseline"),
output_format=args.output_format,
)
continue
path_baseline = Path(args.baseline_path) / row_baseline["filename"]
print(
f"Check:\n"
f"target={str(path_target)} (duplicate_index={row['duplicate_index']})\n"
f"baseline={str(path_baseline)} (duplicate_index={row_baseline['duplicate_index']})"
)
x_baseline = load_object(path_baseline)
x_target = load_object(path_target)
if x_baseline is None or x_target is None:
print(
f"Skip comparison because of None: "
f"x_baseline={x_baseline}, x_target={x_target}"
counts["skipped"] += 1
print_record(
SkipRecord(name=row["name"], reason="load_failed"),
output_format=args.output_format,
)
continue
@@ -74,9 +85,23 @@ def run(args: argparse.Namespace) -> None:
x_baseline=x_baseline,
x_target=x_target,
name=row["name"],
diff_threshold=args.diff_threshold,
)
print_comparison(info=info, diff_threshold=args.diff_threshold)
print()
if info.diff is not None and info.diff.passed:
counts["passed"] += 1
else:
counts["failed"] += 1
print_record(
ComparisonRecord(**info.model_dump()),
output_format=args.output_format,
)
print_record(
SummaryRecord(total=sum(counts.values()), **counts),
output_format=args.output_format,
)
def _parse_args() -> argparse.Namespace:
@@ -90,4 +115,11 @@ def _parse_args() -> argparse.Namespace:
parser.add_argument(
"--filter", type=str, default=None, help="Regex to filter filenames"
)
parser.add_argument(
"--output-format",
type=str,
choices=["text", "json"],
default="text",
help="Output format: text (default) or json (JSONL, one JSON object per line)",
)
return parser.parse_args()

83
python/sglang/srt/debug_utils/comparator/output_types.py Normal file
View File

@@ -0,0 +1,83 @@
from abc import abstractmethod
from typing import Annotated, Literal, Union
from pydantic import Discriminator, TypeAdapter
from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
format_comparison,
)
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
TensorComparisonInfo,
)
from sglang.srt.debug_utils.comparator.utils import _StrictBase
class _OutputRecord(_StrictBase):
@abstractmethod
def to_text(self) -> str: ...
class ConfigRecord(_OutputRecord):
type: Literal["config"] = "config"
baseline_path: str
target_path: str
diff_threshold: float
start_step: int
end_step: int
def to_text(self) -> str:
return (
f"Config: baseline={self.baseline_path} target={self.target_path}\n"
f"diff_threshold={self.diff_threshold} "
f"steps=[{self.start_step}, {self.end_step}]"
)
class SkipRecord(_OutputRecord):
type: Literal["skip"] = "skip"
name: str
reason: str
def to_text(self) -> str:
return f"Skip: {self.name} ({self.reason})"
class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
type: Literal["comparison"] = "comparison"
def to_text(self) -> str:
return format_comparison(self)
class SummaryRecord(_OutputRecord):
type: Literal["summary"] = "summary"
total: int
passed: int
failed: int
skipped: int
def to_text(self) -> str:
return (
f"Summary: {self.passed} passed, {self.failed} failed, "
f"{self.skipped} skipped (total {self.total})"
)
AnyRecord = Annotated[
Union[ConfigRecord, SkipRecord, ComparisonRecord, SummaryRecord],
Discriminator("type"),
]
_any_record_adapter = TypeAdapter(AnyRecord)
def parse_record_json(json_str: str | bytes) -> AnyRecord:
return _any_record_adapter.validate_json(json_str)
def print_record(record: _OutputRecord, output_format: str) -> None:
if output_format == "json":
print(record.model_dump_json())
else:
print(record.to_text())

3
python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py
View File

@@ -1,4 +1 @@
from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors
from sglang.srt.debug_utils.comparator.tensor_comparison.printer import (
print_comparison,
)

49
python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py
View File

@@ -24,22 +24,21 @@ def compare_tensors(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
name: str = "",
diff_threshold: float = 1e-3,
) -> TensorComparisonInfo:
baseline_info = TensorInfo(
shape=x_baseline.shape,
dtype=x_baseline.dtype,
shape=list(x_baseline.shape),
dtype=str(x_baseline.dtype),
stats=_compute_tensor_stats(x_baseline.float()),
sample=None,
)
target_info = TensorInfo(
shape=x_target.shape,
dtype=x_target.dtype,
shape=list(x_target.shape),
dtype=str(x_target.dtype),
stats=_compute_tensor_stats(x_target.float()),
sample=None,
)
x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape)
unified_shape = x_baseline.shape
unified_shape = list(x_baseline.shape)
baseline_original_dtype = x_baseline.dtype
target_original_dtype = x_target.dtype
@@ -54,12 +53,16 @@ def compare_tensors(
downcast_dtype: Optional[torch.dtype] = None
if not shape_mismatch:
diff = _compute_diff(x_baseline=x_baseline_f, x_target=x_target_f)
diff = _compute_diff(
x_baseline=x_baseline_f,
x_target=x_target_f,
diff_threshold=diff_threshold,
)
needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD
if needs_sample:
baseline_info.sample = get_truncated_value(x_baseline_f)
target_info.sample = get_truncated_value(x_target_f)
baseline_info.sample = str(get_truncated_value(x_baseline_f))
target_info.sample = str(get_truncated_value(x_target_f))
if baseline_original_dtype != target_original_dtype:
downcast_dtype = compute_smaller_dtype(
@@ -69,6 +72,7 @@ def compare_tensors(
diff_downcast = _compute_diff(
x_baseline=x_baseline_f.to(downcast_dtype),
x_target=x_target_f.to(downcast_dtype),
diff_threshold=diff_threshold,
)
return TensorComparisonInfo(
@@ -79,7 +83,7 @@ def compare_tensors(
shape_mismatch=shape_mismatch,
diff=diff,
diff_downcast=diff_downcast,
downcast_dtype=downcast_dtype,
downcast_dtype=str(downcast_dtype) if downcast_dtype is not None else None,
)
@@ -101,15 +105,28 @@ def _quantile_or_none(x: torch.Tensor, *, q: float, include: bool) -> Optional[f
return torch.quantile(x, q).item() if include else None
def _compute_diff(x_baseline: torch.Tensor, x_target: torch.Tensor) -> DiffInfo:
def _compute_diff(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
diff_threshold: float = 1e-3,
) -> DiffInfo:
raw_abs_diff = (x_target - x_baseline).abs()
max_diff_coord = argmax_coord(raw_abs_diff)
rel_diff = calc_rel_diff(x_target, x_baseline).item()
max_abs_diff = raw_abs_diff.max().item()
mean_abs_diff = raw_abs_diff.mean().item()
return DiffInfo(
rel_diff=calc_rel_diff(x_target, x_baseline).item(),
max_abs_diff=raw_abs_diff.max().item(),
mean_abs_diff=raw_abs_diff.mean().item(),
max_diff_coord=max_diff_coord,
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_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(),
passed=(
rel_diff <= diff_threshold
and max_abs_diff <= diff_threshold
and mean_abs_diff <= diff_threshold
),
)

88
python/sglang/srt/debug_utils/comparator/tensor_comparison/formatter.py Normal file
View File

@@ -0,0 +1,88 @@
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
DiffInfo,
TensorComparisonInfo,
TensorStats,
)
def format_comparison(info: TensorComparisonInfo) -> str:
lines: list[str] = []
baseline = info.baseline
target = info.target
dtype_marker = "" if baseline.dtype == target.dtype else "🟠"
lines.append(
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:
lines.append(
f"Unify shape: {baseline.shape} -> {info.unified_shape} "
f"(to match {target.shape})"
)
lines.append(
f"After unify "
f"[shape] {info.unified_shape} vs {target.shape}\t"
f"[dtype] {baseline.dtype} vs {target.dtype}"
)
lines.extend(_format_stats_comparison(baseline=baseline.stats, target=target.stats))
if info.shape_mismatch:
lines.append("⚠️ Shape mismatch")
return "\n".join(lines)
if info.diff is not None:
lines.extend(_format_diff(diff=info.diff))
if info.diff_downcast is not None and info.downcast_dtype is not None:
lines.extend(
_format_diff(
diff=info.diff_downcast,
prefix_text=f"When downcast to {info.downcast_dtype}: ",
)
)
if baseline.sample is not None:
lines.append(f"x_baseline(sample)={baseline.sample}")
if target.sample is not None:
lines.append(f"x_target(sample)={target.sample}")
return "\n".join(lines)
def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]:
lines: list[str] = []
for stat_name in TensorStats.model_fields:
value_baseline = getattr(baseline, stat_name)
value_target = getattr(target, stat_name)
if value_baseline is None or value_target is None:
continue
lines.append(
f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
f"(diff: {value_target - value_baseline:.4f})"
)
return lines
def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
marker = "✅" if diff.passed else "❌"
return [
prefix_text
+ marker
+ " "
+ "\t".join(
f"{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),
]
),
f"max_abs_diff happens at coord={diff.max_diff_coord} with "
f"baseline={diff.baseline_at_max} "
f"target={diff.target_at_max}",
]

28
python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py
View File

@@ -1,11 +1,9 @@
from dataclasses import dataclass
from typing import Optional, Tuple
from typing import Optional
import torch
from sglang.srt.debug_utils.comparator.utils import _StrictBase
@dataclass
class TensorStats:
class TensorStats(_StrictBase):
mean: float
std: float
min: float
@@ -16,31 +14,29 @@ class TensorStats:
p99: Optional[float] = None
@dataclass
class TensorInfo:
shape: torch.Size
dtype: torch.dtype
class TensorInfo(_StrictBase):
shape: list[int]
dtype: str
stats: TensorStats
sample: Optional[str] = None
@dataclass
class DiffInfo:
class DiffInfo(_StrictBase):
rel_diff: float
max_abs_diff: float
mean_abs_diff: float
max_diff_coord: Tuple[int, ...]
max_diff_coord: list[int]
baseline_at_max: float
target_at_max: float
passed: bool
@dataclass
class TensorComparisonInfo:
class TensorComparisonInfo(_StrictBase):
name: str
baseline: TensorInfo
target: TensorInfo
unified_shape: Optional[torch.Size]
unified_shape: Optional[list[int]]
shape_mismatch: bool
diff: Optional[DiffInfo] = None
diff_downcast: Optional[DiffInfo] = None
downcast_dtype: Optional[torch.dtype] = None
downcast_dtype: Optional[str] = None

5
python/sglang/srt/debug_utils/comparator/utils.py
View File

@@ -3,6 +3,11 @@ from pathlib import Path
from typing import Optional, Tuple
import torch
from pydantic import BaseModel, ConfigDict
class _StrictBase(BaseModel):
model_config = ConfigDict(extra="forbid")
def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]: