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

Enhance output formatting in dump comparator (#21029)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-20 22:04:50 +08:00
committed by GitHub
parent fdbcb8156e
commit 6703cc4484
16 changed files with 815 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
View File

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
from sglang.srt.debug_utils.comparator.dims_spec import (
SEQ_DIM_NAME,
TOKEN_DIM_NAME,
ParallelAxis,
apply_dim_names,
parse_dims,
resolve_dim_names,
@@ -39,11 +40,53 @@ from sglang.srt.debug_utils.comparator.output_types import (
)
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compare_tensor_pair,
compute_tensor_info,
)
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
_FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"}
def _build_skip_from_one_empty_side(
*, name: str, pair: Pair[list[ValueWithMeta]]
) -> ComparisonSkipRecord:
"""Build a skip record when one side of *pair* is empty.
The non-empty side's tensor info is attached to the record.
"""
assert not pair.x or not pair.y
if not pair.x:
reason, available_side, available_items = (
"baseline_load_failed",
"target",
pair.y,
)
else:
reason, available_side, available_items = (
"target_load_failed",
"baseline",
pair.x,
)
tensor_items: list[ValueWithMeta] = [
it for it in available_items if isinstance(it.value, torch.Tensor)
]
if not tensor_items:
return ComparisonSkipRecord(name=name, reason=reason)
first_tensor: torch.Tensor = tensor_items[0].value
tensor_info = compute_tensor_info(first_tensor, include_sample=True)
metas: list[dict[str, Any]] = [it.meta for it in tensor_items]
bundle_info: BundleSideInfo = _collect_bundle_side_info(
items=tensor_items, metas=metas
)
return ComparisonSkipRecord(
name=name,
reason=reason,
available_side=available_side, # type: ignore[arg-type]
available_tensor_info=tensor_info,
available_bundle_info=bundle_info,
)
def _collect_bundle_side_info(
@@ -51,8 +94,8 @@ def _collect_bundle_side_info(
metas: list[dict[str, Any]],
) -> BundleSideInfo:
from sglang.srt.debug_utils.comparator.display import (
_PARALLEL_INFO_KEYS,
extract_parallel_info,
PARALLEL_INFO_KEYS,
_extract_parallel_info,
)
files: list[BundleFileInfo] = []
@@ -61,8 +104,8 @@ def _collect_bundle_side_info(
tensor: torch.Tensor = item.value
parallel_info: dict[str, str] = {}
for key in _PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
for key in PARALLEL_INFO_KEYS:
_extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
files.append(
BundleFileInfo(
@@ -70,6 +113,7 @@ def _collect_bundle_side_info(
dtype=str(tensor.dtype),
rank=meta.get("rank"),
parallel_info=parallel_info if parallel_info else None,
filename=meta.get("filename"),
)
)
@@ -132,8 +176,7 @@ def _compare_bundle_pair_inner(
)
if not all_pair.x or not all_pair.y:
reason = "baseline_load_failed" if not all_pair.x else "target_load_failed"
return ComparisonSkipRecord(name=name, reason=reason)
return _build_skip_from_one_empty_side(name=name, pair=all_pair)
# 1b. Dims override: patch meta["dims"] before DP filter reads it
# (--override-dims may add ``# dp:=moe_dp``, so it must run first)
@@ -157,7 +200,7 @@ def _compare_bundle_pair_inner(
# 1c. DP filter: keep only the non-empty dp_rank
all_pair = all_pair.map(
lambda items: filter_to_non_empty_dp_rank(
items, dp_group_alias=_extract_dp_alias_from_items(items)
items, dp_axis=_extract_dp_axis_from_items(items)
)
)
@@ -181,14 +224,14 @@ def _compare_bundle_pair_inner(
)
def _extract_dp_alias_from_items(items: list[ValueWithMeta]) -> Optional[str]:
"""Extract dp group alias from the first item's ``meta["dims"]``."""
def _extract_dp_axis_from_items(items: list[ValueWithMeta]) -> ParallelAxis:
"""Extract dp axis from the first item's ``meta["dims"]``."""
if not items:
return None
return ParallelAxis.DP
dims_str: Optional[str] = items[0].meta.get("dims")
if dims_str is None:
return None
return parse_dims(dims_str).dp_group_alias
return ParallelAxis.DP
return parse_dims(dims_str).dp_axis
def _compare_bundle_pair_tensor_type(
@@ -205,8 +248,7 @@ def _compare_bundle_pair_tensor_type(
compute_per_token: bool = False,
) -> Union[ComparisonTensorRecord, ComparisonSkipRecord]:
if not valid_pair.x or not valid_pair.y:
reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
return ComparisonSkipRecord(name=name, reason=reason)
return _build_skip_from_one_empty_side(name=name, pair=valid_pair)
# Plan (meta only, no tensor)
metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
@@ -243,9 +285,12 @@ def _compare_bundle_pair_tensor_type(
if aligner_result.tensors is None:
assert aligner_result.failed_side_xy is not None
side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
reason: str = f"{side_name}_load_failed"
return ComparisonSkipRecord(name=name, reason=reason)
failed_xy: str = aligner_result.failed_side_xy
pair_with_failed_emptied: Pair[list[ValueWithMeta]] = Pair(
x=[] if failed_xy == "x" else valid_pair.x,
y=[] if failed_xy == "y" else valid_pair.y,
)
return _build_skip_from_one_empty_side(name=name, pair=pair_with_failed_emptied)
# Resolve seq_dim for per-token computation
seq_dim: Optional[int] = (

13
python/sglang/srt/debug_utils/comparator/display.py
View File

@@ -52,6 +52,19 @@ def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) ->
return buf.getvalue().rstrip("\n")
def _render_polars_as_rich_table(
df: pl.DataFrame, *, title: Optional[str] = None
) -> Any:
from rich.table import Table
table = Table(title=title)
for col in df.columns:
table.add_column(col)
for row in df.iter_rows():
table.add_row(*[str(v) for v in row])
return table
def _collect_rank_info(
df: pl.DataFrame, dump_dir: Path
) -> Optional[list[dict[str, Any]]]:

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

@@ -48,6 +48,14 @@ from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path
_DEFAULT_SKIP_KEYS: set[str] = {"dump_index", "filename"}
_DIMS_DEBUG_HINT: str = (
"\nHint: If this is a dims annotation issue, do NOT re-run expensive dumps.\n"
"Use --override-dims at comparison time, e.g.:\n"
' python -m sglang.srt.debug_utils.comparator --override-dims "tensor_name:b s h[tp] d"\n'
"(Use --override-baseline-dims / --override-target-dims for per-side overrides.\n"
" Use --override-config for bulk overrides via YAML file.)"
)
def main() -> None:
args = parse_args(sys.argv[1:])
@@ -254,10 +262,12 @@ def _compare_bundle_pairs(
meta_overrider=meta_overrider,
)
except Exception as exc:
tb = _traceback_module.format_exc()
record = ComparisonErrorRecord(
name=name,
exception_type=type(exc).__name__,
traceback_str=_traceback_module.format_exc(),
exception_message=str(exc),
traceback_str=f"{_DIMS_DEBUG_HINT}\n\n{tb}",
)
target_steps: set[int] = {info.step for info in bundle_info_pair.y}

75
python/sglang/srt/debug_utils/comparator/output_formatter.py
View File

@@ -116,17 +116,72 @@ def _format_config_rich_body(
def _format_skip_body(record: ComparisonSkipRecord) -> str:
return f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
text: str = (
f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
)
if record.available_side is not None and record.available_tensor_info is not None:
info = record.available_tensor_info
text += f"\n {record.available_side}: shape={info.shape} dtype={info.dtype}"
text += (
f" mean={info.stats.mean:.4f} std={info.stats.std:.4f}"
f" range=[{info.stats.min:.4f}, {info.stats.max:.4f}]"
)
if info.sample is not None:
text += f"\n sample: {info.sample}"
return text
def _format_skip_rich_body(
record: ComparisonSkipRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
suffix: str = record._format_location_suffix()
return (
header: str = (
f"[dim]⊘ {escape(record.name)}{suffix} ── skipped ({escape(record.reason)})[/]"
)
if (
verbosity == "minimal"
or record.available_side is None
or record.available_tensor_info is None
):
return header
info = record.available_tensor_info
side: str = record.available_side
dtype_str: str = info.dtype.replace("torch.", "")
lines: list[str] = [header]
# Bundle info line
if record.available_bundle_info is not None:
bi = record.available_bundle_info
shapes: list[list[int]] = [f.shape for f in bi.files]
unique_shapes: set[str] = {str(s) for s in shapes}
shape_desc: str = (
escape(str(shapes[0])) if len(unique_shapes) == 1 else "mixed shapes"
)
dims_part: str = f" [dim]dims: {bi.dims}[/]" if bi.dims else ""
lines.append(
f" {side:8s} [cyan]{bi.num_files} files[/]"
f" × {shape_desc} {dtype_str}{dims_part}"
)
else:
lines.append(f" {side:8s} {escape(str(info.shape))} {dtype_str}")
# Stats line (compact single-side)
stats = info.stats
range_str: str = escape(f"[{stats.min:.4f}, {stats.max:.4f}]")
lines.append(
f" [dim]stats[/] mean={stats.mean:.4f} std={stats.std:.4f}"
f" range={range_str}"
)
# Sample
if info.sample is not None:
lines.append(f" [dim]sample[/] {escape(info.sample)}")
return "\n".join(lines)
# ── ComparisonErrorRecord ────────────────────────────────────────────
@@ -135,6 +190,7 @@ def _format_error_body(record: ComparisonErrorRecord) -> str:
prefix: str = record._format_location_prefix()
return (
f"{prefix}Error: {record.name} ({record.exception_type})\n"
f"{record.exception_message}\n"
f"{record.traceback_str}"
)
@@ -145,7 +201,8 @@ def _format_error_rich_body(
prefix: str = record._format_location_prefix_rich()
name: str = escape(record.name)
header: str = (
f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)})[/]"
f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)}): "
f"{escape(record.exception_message)}[/]"
)
if verbosity == "minimal":
return header
@@ -300,7 +357,17 @@ def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
sub_desc: str = f"{traced_sub.plan.type}"
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
sub = traced_sub.plan
qualifier: str = ""
if isinstance(sub, UnsharderPlan):
qualifier = f"({sub.axis.value})"
elif isinstance(sub, ReordererPlan):
qualifier = f"({sub.params.op})"
sub_desc: str = f"{sub.type}{qualifier}"
if traced_sub.snapshot is not None:
snap = traced_sub.snapshot

14
python/sglang/srt/debug_utils/comparator/output_types.py
View File

@@ -4,7 +4,7 @@ from abc import abstractmethod
from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
from rich.console import RenderableType
from rich.console import Group, RenderableType
from rich.markup import escape
from sglang.srt.debug_utils.comparator.output_formatter import ( # noqa: F401 — re-export
@@ -32,6 +32,7 @@ from sglang.srt.debug_utils.comparator.output_formatter import (
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorInfo,
)
from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
@@ -84,6 +85,7 @@ class BundleFileInfo(_StrictBase):
dtype: str
rank: Optional[int] = None
parallel_info: Optional[dict[str, str]] = None # e.g. {"tp": "0/4", "ep": "1/2"}
filename: Optional[str] = None
class BundleSideInfo(_StrictBase):
@@ -121,6 +123,12 @@ class RecordLocation(_StrictBase):
class _BaseComparisonRecord(_OutputRecord):
location: RecordLocation = Field(default_factory=RecordLocation)
def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
result = _render_record_rich(self, verbosity=verbosity)
if isinstance(result, str):
return result + "\n"
return Group(result, "")
def _format_location_prefix(self) -> str:
if self.location.step is not None:
return f"[step={self.location.step}] "
@@ -152,6 +160,9 @@ class ComparisonSkipRecord(_BaseComparisonRecord):
type: Literal["comparison_skip"] = "comparison_skip"
name: str
reason: str
available_side: Optional[Literal["baseline", "target"]] = None
available_tensor_info: Optional[TensorInfo] = None
available_bundle_info: Optional[BundleSideInfo] = None
@property
def category(self) -> str:
@@ -170,6 +181,7 @@ class ComparisonErrorRecord(_BaseComparisonRecord):
type: Literal["comparison_error"] = "comparison_error"
name: str
exception_type: str
exception_message: str
traceback_str: str
@property

8
python/sglang/srt/debug_utils/comparator/report_sink.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import IO, Literal, Optional
@@ -72,7 +73,11 @@ class ReportSink:
def _get_console(self) -> Console:
if self._console is None:
self._console = Console()
try:
width = os.get_terminal_size().columns
except OSError:
width = 200
self._console = Console(force_terminal=True, width=width)
return self._console
def _print_to_stdout(self, record: _OutputRecord) -> None:
@@ -81,7 +86,6 @@ class ReportSink:
else:
console: Console = self._get_console()
console.print(record.to_rich(verbosity=self._verbosity))
console.print() # blank line between records
report_sink = ReportSink()

28
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
View File

@@ -23,6 +23,22 @@ QUANTILE_NUMEL_THRESHOLD = 10_000_000
SAMPLE_DIFF_THRESHOLD = 1e-3
def compute_tensor_info(
tensor: torch.Tensor, *, include_sample: bool = False
) -> TensorInfo:
"""Compute TensorInfo (shape, dtype, stats, optional sample) for a single tensor."""
stats: TensorStats = _compute_tensor_stats(tensor.float())
sample: Optional[str] = (
str(get_truncated_value(tensor.float())) if include_sample else None
)
return TensorInfo(
shape=list(tensor.shape),
dtype=str(tensor.dtype),
stats=stats,
sample=sample,
)
def compare_tensor_pair(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
@@ -30,16 +46,8 @@ def compare_tensor_pair(
diff_threshold: float = 1e-3,
seq_dim: Optional[int] = None,
) -> TensorComparisonInfo:
baseline_info = TensorInfo(
shape=list(x_baseline.shape),
dtype=str(x_baseline.dtype),
stats=_compute_tensor_stats(x_baseline.float()),
)
target_info = TensorInfo(
shape=list(x_target.shape),
dtype=str(x_target.dtype),
stats=_compute_tensor_stats(x_target.float()),
)
baseline_info: TensorInfo = compute_tensor_info(x_baseline)
target_info: TensorInfo = compute_tensor_info(x_target)
x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape)
unified_shape = list(x_baseline.shape)

46
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
View File

@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Literal, Optional
from rich.markup import escape
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
@@ -70,11 +71,13 @@ def _category_marker(category: str) -> tuple[bool, str, str]:
# ---------------------------------------------------------------------------
_STAT_HEADER = (
f" [dim]{'':10s} {'baseline':>10s} {'target':>10s} {'Δ':s}[/]"
)
def _format_stat_line(stat_name: str, val_b: float, val_t: float, diff: float) -> str:
return (
f" [blue]{stat_name:10s}[/] {val_b:>10.4f} vs {val_t:>10.4f}"
f" Δ {_fmt_diff_colored(diff)}"
)
return f" [blue]{stat_name:10s}[/] {val_b:>10.4f} {val_t:>10.4f} {_fmt_diff_colored(diff)}"
# ---------------------------------------------------------------------------
@@ -352,7 +355,7 @@ def _format_bundle_section(
for label, side in [("baseline", bundle_info.x), ("target", bundle_info.y)]:
if not side.files:
lines.append(f" {label} [dim](no files)[/]")
lines.append(f" {label:8s} [dim](no files)[/]")
continue
dtype_desc: str = _strip_torch_prefix(side.files[0].dtype)
@@ -360,7 +363,7 @@ def _format_bundle_section(
if verbose:
dims_part: str = f" dims: {side.dims}" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" {label:8s} [cyan]{side.num_files} files[/]"
f" {dtype_desc}{dims_part}"
)
@@ -371,8 +374,9 @@ def _format_bundle_section(
par_part = " " + " ".join(
f"{k}={v}" for k, v in f.parallel_info.items()
)
file_part: str = f" [dim]{escape(f.filename)}[/]" if f.filename else ""
lines.append(
f" [{idx}] {_esc_shape(f.shape)} {rank_part}{par_part}"
f" [{idx}] {_esc_shape(f.shape)} {rank_part}{par_part}{file_part}"
)
else:
shapes: list[list[int]] = [f.shape for f in side.files]
@@ -385,7 +389,7 @@ def _format_bundle_section(
dims_part = f" [dim]dims: {side.dims}[/]" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" {label:8s} [cyan]{side.num_files} files[/]"
f" × {shape_desc} {dtype_desc}{dims_part}"
)
@@ -404,7 +408,7 @@ def _format_plan_section_rich(
("target", traced_plan.per_side.y),
]:
if not traced_side.step_plans:
lines.append(f" {side_label} [dim](passthrough)[/]")
lines.append(f" {side_label:8s} [dim](passthrough)[/]")
continue
parts: list[str] = [
@@ -412,7 +416,7 @@ def _format_plan_section_rich(
for traced_step in traced_side.step_plans
for traced_sub in traced_step.sub_plans
]
lines.append(f" {side_label} " + " → ".join(parts))
lines.append(f" {side_label:8s} " + " → ".join(parts))
lines.extend(_format_cross_side_plan_rich(traced_plan.plan))
return lines
@@ -423,9 +427,11 @@ def _format_sub_plan_rich(traced_sub: TracedSubPlan) -> str:
snapshot: Optional[ShapeSnapshot] = traced_sub.snapshot
op_name: str = sub.type
axis_str: str = ""
qualifier: str = ""
if isinstance(sub, UnsharderPlan):
axis_str = f"({sub.axis})"
qualifier = f"({sub.axis.value})"
elif isinstance(sub, ReordererPlan):
qualifier = f"({sub.params.op})"
shape_change: str = ""
if snapshot:
@@ -437,9 +443,9 @@ def _format_sub_plan_rich(traced_sub: TracedSubPlan) -> str:
out_shape: str = (
_esc_shape(snapshot.output_shapes[0]) if snapshot.output_shapes else "?"
)
shape_change = f" {in_count}×{in_shape} → {out_count}×{out_shape}"
shape_change = f" ({in_count}×{in_shape} → {out_count}×{out_shape})"
return f"[magenta]{op_name}{axis_str}[/]{shape_change}"
return f"[magenta]{op_name}{qualifier}[/]{shape_change}"
def _format_cross_side_plan_rich(plan: AlignerPlan) -> list[str]:
@@ -469,7 +475,7 @@ def _format_stats_rich(
target: TensorStats,
verbose: bool = False,
) -> list[str]:
lines: list[str] = []
lines: list[str] = [_STAT_HEADER]
if verbose:
# All stat fields
@@ -486,7 +492,7 @@ def _format_stats_rich(
val_t = target.percentiles[p]
lines.append(_format_stat_line(f"p{p}", val_b, val_t, val_t - val_b))
else:
# Compact: mean, std, range (min/max combined)
# Compact: mean, std, range, then percentiles
for stat_name in ("mean", "std"):
val_b = getattr(baseline, stat_name)
val_t = getattr(target, stat_name)
@@ -495,7 +501,13 @@ def _format_stats_rich(
# Range line: combine min/max (escape brackets to avoid Rich markup)
range_baseline: str = escape(f"[{baseline.min:.4f}, {baseline.max:.4f}]")
range_target: str = escape(f"[{target.min:.4f}, {target.max:.4f}]")
lines.append(f" [blue]{'range':10s}[/] {range_baseline} vs {range_target}")
lines.append(f" [blue]{'range':10s}[/] {range_baseline} {range_target}")
# Percentiles (compact: same as verbose)
for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
val_b = baseline.percentiles[p]
val_t = target.percentiles[p]
lines.append(_format_stat_line(f"p{p}", val_b, val_t, val_t - val_b))
return lines