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

Visualize comparison detailed results in dump comparator (#19565)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-28 18:08:16 +08:00
committed by GitHub
parent 40facdb28c
commit 80bbd30909
11 changed files with 867 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
from sglang.srt.debug_utils.comparator.dims import apply_dim_names, parse_dim_names
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
GeneralWarning,
NonTensorRecord,
SkipRecord,
)
@@ -45,6 +46,7 @@ def compare_bundle_pair(
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=None, y=None
),
viz_output_dir: Optional[Path] = None,
) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
with warning_sink.context() as collected_warnings:
result = _compare_bundle_pair_inner(
@@ -55,6 +57,7 @@ def compare_bundle_pair(
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
)
return result.model_copy(update={"warnings": collected_warnings})
@@ -71,6 +74,7 @@ def _compare_bundle_pair_inner(
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=None, y=None
),
viz_output_dir: Optional[Path] = None,
) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
# 1. Load all successfully loaded values
all_pair: Pair[list[ValueWithMeta]] = Pair(
@@ -96,6 +100,7 @@ def _compare_bundle_pair_inner(
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
)
@@ -108,6 +113,7 @@ def _compare_bundle_pair_tensor_type(
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=None, y=None
),
viz_output_dir: Optional[Path] = None,
) -> Union[ComparisonRecord, SkipRecord]:
if not valid_pair.x or not valid_pair.y:
reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
@@ -145,13 +151,59 @@ def _compare_bundle_pair_tensor_type(
return SkipRecord(name=name, reason=reason)
# Compare
aligned_baseline: torch.Tensor = aligner_result.tensors.x.rename(None)
aligned_target: torch.Tensor = aligner_result.tensors.y.rename(None)
info = compare_tensor_pair(
x_baseline=aligner_result.tensors.x.rename(None),
x_target=aligner_result.tensors.y.rename(None),
x_baseline=aligned_baseline,
x_target=aligned_target,
name=name,
diff_threshold=diff_threshold,
)
return ComparisonRecord(**info.model_dump(), aligner_plan=plan)
record = ComparisonRecord(**info.model_dump(), aligner_plan=plan)
if viz_output_dir is not None:
_try_generate_viz(
baseline=aligned_baseline,
target=aligned_target,
name=name,
viz_output_dir=viz_output_dir,
)
return record
def _try_generate_viz(
*,
baseline: torch.Tensor,
target: torch.Tensor,
name: str,
viz_output_dir: Path,
) -> None:
from sglang.srt.debug_utils.comparator.visualizer import (
generate_comparison_figure,
)
from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
_sanitize_filename,
)
filename: str = _sanitize_filename(name) + ".png"
output_path: Path = viz_output_dir / filename
try:
generate_comparison_figure(
baseline=baseline,
target=target,
name=name,
output_path=output_path,
)
except Exception as exc:
warning_sink.add(
GeneralWarning(
category="visualizer",
message=f"Visualization failed for {name}: {exc}",
)
)
def _compare_bundle_pair_non_tensor_type(

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

@@ -74,6 +74,10 @@ def run(args: argparse.Namespace) -> None:
),
)
viz_output_dir: Optional[Path] = (
Path(args.viz_output_dir) if args.viz_bundle_details else None
)
comparison_records = _compare_bundle_pairs(
bundle_info_pairs=bundle_info_pairs,
baseline_path=Path(args.baseline_path),
@@ -81,6 +85,7 @@ def run(args: argparse.Namespace) -> None:
token_aligner_plan=ta_result.plan,
diff_threshold=args.diff_threshold,
thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
)
_consume_comparison_records(
comparison_records=comparison_records, output_format=args.output_format
@@ -138,6 +143,7 @@ def _compare_bundle_pairs(
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]],
viz_output_dir: Optional[Path] = None,
) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]:
for bundle_info_pair in bundle_info_pairs:
if not bundle_info_pair.y:
@@ -155,6 +161,7 @@ def _compare_bundle_pairs(
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
)
@@ -205,4 +212,16 @@ def _parse_args() -> argparse.Namespace:
default=None,
help="Tokenizer path for decoding input_ids (auto-discovered from dump metadata if not set)",
)
parser.add_argument(
"--viz-bundle-details",
action="store_true",
default=False,
help="Generate comparison heatmap/histogram PNG for each compared tensor",
)
parser.add_argument(
"--viz-output-dir",
type=str,
default="/tmp/comparator_viz/",
help="Output directory for visualization PNGs (default: /tmp/comparator_viz/)",
)
return parser.parse_args()

3
python/sglang/srt/debug_utils/comparator/visualizer/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
from sglang.srt.debug_utils.comparator.visualizer.figure import ( # noqa: F401
generate_comparison_figure,
)

116
python/sglang/srt/debug_utils/comparator/visualizer/figure.py Normal file
View File

@@ -0,0 +1,116 @@
"""Main orchestration logic for comparison figure generation."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional
import numpy as np
import torch
from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
_preprocess_tensor,
)
@dataclass(frozen=True)
class _PanelContext:
baseline_2d: torch.Tensor
target_2d: torch.Tensor
diff: Optional[torch.Tensor] # None when shapes differ
name: str
@dataclass(frozen=True)
class _Panel:
label: str
requires_diff: bool
draw: Callable[[np.ndarray, int, _PanelContext], Optional[str]]
def _build_panels() -> list[_Panel]:
from sglang.srt.debug_utils.comparator.visualizer.panels import (
_draw_baseline_heatmap,
_draw_diff_heatmap,
_draw_diff_histogram,
_draw_hist2d,
_draw_sampled,
_draw_target_heatmap,
)
return [
_Panel(
label="Baseline Heatmap", requires_diff=False, draw=_draw_baseline_heatmap
),
_Panel(label="Target Heatmap", requires_diff=False, draw=_draw_target_heatmap),
_Panel(label="Abs Diff Heatmap", requires_diff=True, draw=_draw_diff_heatmap),
_Panel(label="Abs Diff Hist", requires_diff=True, draw=_draw_diff_histogram),
_Panel(label="Hist2D", requires_diff=True, draw=_draw_hist2d),
_Panel(label="Sampled", requires_diff=True, draw=_draw_sampled),
]
def generate_comparison_figure(
*,
baseline: torch.Tensor,
target: torch.Tensor,
name: str,
output_path: Path,
) -> None:
"""Generate a multi-panel comparison PNG for a baseline/target tensor pair.
Panels (6 rows x 2 cols, left=normal, right=log10):
Row 0: Baseline heatmap
Row 1: Target heatmap
Row 2: Abs Diff heatmap
Row 3: Abs Diff histogram
Row 4: Hist2D scatter (baseline vs target density)
Row 5: Sampled scatter (10k sampled mini-heatmap)
"""
import matplotlib.pyplot as plt
baseline_f: torch.Tensor = baseline.detach().cpu().float()
target_f: torch.Tensor = target.detach().cpu().float()
can_diff: bool = baseline_f.shape == target_f.shape
baseline_2d: torch.Tensor = _preprocess_tensor(baseline_f)
target_2d: torch.Tensor = _preprocess_tensor(target_f)
diff: Optional[torch.Tensor] = (baseline_2d - target_2d).abs() if can_diff else None
ctx = _PanelContext(
baseline_2d=baseline_2d,
target_2d=target_2d,
diff=diff,
name=name,
)
panels: list[_Panel] = _build_panels()
active: list[_Panel] = [p for p in panels if not p.requires_diff or can_diff]
nrows: int = len(active)
ncols: int = 2
fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 3.5 * nrows))
if nrows == 1:
axes = axes.reshape(1, -1)
stats_lines: list[str] = []
for i, panel in enumerate(active):
stats_line: Optional[str] = panel.draw(axes, i, ctx)
if stats_line is not None:
stats_lines.append(stats_line)
num_stats: int = len(stats_lines)
title_height: float = 0.015 * num_stats + 0.015
fig.suptitle(
"\n".join(stats_lines),
fontsize=9,
family="monospace",
y=1 - title_height / 2,
)
plt.tight_layout(rect=[0, 0, 1, 1 - title_height])
output_path.parent.mkdir(parents=True, exist_ok=True)
plt.savefig(str(output_path), dpi=150, bbox_inches="tight")
plt.close(fig)

226
python/sglang/srt/debug_utils/comparator/visualizer/panels.py Normal file
View File

@@ -0,0 +1,226 @@
"""Panel draw functions for tensor comparison visualization."""
from __future__ import annotations
from typing import Optional
import numpy as np
import torch
from sglang.srt.debug_utils.comparator.visualizer.figure import _PanelContext
from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
_SCATTER_SAMPLE_SIZE,
_format_log_ticks,
_format_stats,
_maybe_downsample_numpy,
_safe_hist,
_to_log10,
)
def _draw_baseline_heatmap(
axes: np.ndarray, row_idx: int, ctx: _PanelContext
) -> Optional[str]:
_draw_heatmap_pair(
axes, row_idx=row_idx, t=ctx.baseline_2d, title=f"{ctx.name} Baseline"
)
return _format_stats("Baseline", ctx.baseline_2d)
def _draw_target_heatmap(
axes: np.ndarray, row_idx: int, ctx: _PanelContext
) -> Optional[str]:
_draw_heatmap_pair(
axes, row_idx=row_idx, t=ctx.target_2d, title=f"{ctx.name} Target"
)
return _format_stats("Target", ctx.target_2d)
def _draw_diff_heatmap(
axes: np.ndarray, row_idx: int, ctx: _PanelContext
) -> Optional[str]:
assert ctx.diff is not None
_draw_heatmap_pair(axes, row_idx=row_idx, t=ctx.diff, title=f"{ctx.name} Abs Diff")
return _format_stats("Abs Diff", ctx.diff)
def _draw_diff_histogram(
axes: np.ndarray, row_idx: int, ctx: _PanelContext
) -> Optional[str]:
assert ctx.diff is not None
_draw_histogram_pair(
axes, row_idx=row_idx, diff=ctx.diff, label=f"{ctx.name} Abs Diff"
)
return None
def _draw_hist2d(axes: np.ndarray, row_idx: int, ctx: _PanelContext) -> Optional[str]:
_draw_scatter_hist2d(
axes,
row_idx=row_idx,
baseline=ctx.baseline_2d,
target=ctx.target_2d,
label=ctx.name,
)
return None
def _draw_sampled(axes: np.ndarray, row_idx: int, ctx: _PanelContext) -> Optional[str]:
_draw_scatter_sampled(
axes,
row_idx=row_idx,
baseline=ctx.baseline_2d,
target=ctx.target_2d,
label=ctx.name,
)
return None
# ────────────────────── internal drawing helpers ──────────────────────
def _draw_heatmap_pair(
axes: np.ndarray,
*,
row_idx: int,
t: torch.Tensor,
title: str,
) -> None:
import matplotlib.pyplot as plt
ax_normal = axes[row_idx, 0]
ax_log = axes[row_idx, 1]
im = ax_normal.imshow(t.numpy(), aspect="auto", cmap="viridis")
ax_normal.set_title(title)
plt.colorbar(im, ax=ax_normal)
im_log = ax_log.imshow(_to_log10(t).numpy(), aspect="auto", cmap="viridis")
ax_log.set_title(f"{title} (Log10)")
cbar = plt.colorbar(im_log, ax=ax_log)
_format_log_ticks(cbar.ax, axis="y")
def _draw_histogram_pair(
axes: np.ndarray,
*,
row_idx: int,
diff: torch.Tensor,
label: str,
) -> None:
ax_normal = axes[row_idx, 0]
ax_log = axes[row_idx, 1]
diff_flat: np.ndarray = _maybe_downsample_numpy(diff.flatten())
_safe_hist(ax_normal, diff_flat, bins=100, edgecolor="none")
ax_normal.set_title(f"{label} Histogram")
ax_normal.set_xlabel("Abs Diff")
ax_normal.set_ylabel("Count")
log_flat: np.ndarray = np.log10(np.abs(diff_flat) + 1e-10)
_safe_hist(ax_log, log_flat, bins=100, edgecolor="none")
ax_log.set_title(f"{label} Histogram (Log10)")
ax_log.set_xlabel("Abs Diff")
ax_log.set_ylabel("Count")
_format_log_ticks(ax_log, axis="x")
def _draw_scatter_hist2d(
axes: np.ndarray,
*,
row_idx: int,
baseline: torch.Tensor,
target: torch.Tensor,
label: str,
) -> None:
import matplotlib.pyplot as plt
ax_normal = axes[row_idx, 0]
ax_log = axes[row_idx, 1]
b_flat: np.ndarray = _maybe_downsample_numpy(baseline.flatten())
t_flat: np.ndarray = _maybe_downsample_numpy(target.flatten())
min_len: int = min(len(b_flat), len(t_flat))
b_flat = b_flat[:min_len]
t_flat = t_flat[:min_len]
# Normal scale
lim: float = float(max(np.abs(b_flat).max(), np.abs(t_flat).max())) * 1.05
if lim == 0:
lim = 1.0
_h, _xe, _ye, im = ax_normal.hist2d(
b_flat,
t_flat,
bins=200,
range=[[-lim, lim], [-lim, lim]],
cmap="viridis",
norm="log",
)
ax_normal.plot([-lim, lim], [-lim, lim], "r--", linewidth=0.5)
ax_normal.set_title(f"{label} Hist2D")
ax_normal.set_xlabel("Baseline")
ax_normal.set_ylabel("Target")
ax_normal.set_aspect("equal")
plt.colorbar(im, ax=ax_normal)
# Log scale
b_log: np.ndarray = np.log10(np.abs(b_flat) + 1e-10)
t_log: np.ndarray = np.log10(np.abs(t_flat) + 1e-10)
vmin: float = float(min(b_log.min(), t_log.min())) - 0.5
vmax: float = float(max(b_log.max(), t_log.max())) + 0.5
_h2, _xe2, _ye2, im2 = ax_log.hist2d(
b_log,
t_log,
bins=200,
range=[[vmin, vmax], [vmin, vmax]],
cmap="viridis",
norm="log",
)
ax_log.plot([vmin, vmax], [vmin, vmax], "r--", linewidth=0.5)
ax_log.set_title(f"{label} Hist2D (Log10 Abs)")
ax_log.set_xlabel("Baseline")
ax_log.set_ylabel("Target")
ax_log.set_aspect("equal")
plt.colorbar(im2, ax=ax_log)
_format_log_ticks(ax_log, axis="both")
def _draw_scatter_sampled(
axes: np.ndarray,
*,
row_idx: int,
baseline: torch.Tensor,
target: torch.Tensor,
label: str,
) -> None:
import matplotlib.pyplot as plt
ax_baseline = axes[row_idx, 0]
ax_target = axes[row_idx, 1]
b_flat: np.ndarray = baseline.flatten().numpy()
t_flat: np.ndarray = target.flatten().numpy()
n_samples: int = min(_SCATTER_SAMPLE_SIZE, len(b_flat))
rng: np.random.Generator = np.random.default_rng(seed=42)
indices: np.ndarray = np.sort(rng.choice(len(b_flat), n_samples, replace=False))
b_sampled: np.ndarray = b_flat[indices]
t_sampled: np.ndarray = t_flat[indices]
side: int = int(np.sqrt(n_samples))
n_use: int = side * side
b_2d: np.ndarray = b_sampled[:n_use].reshape(side, side)
t_2d: np.ndarray = t_sampled[:n_use].reshape(side, side)
vmin: float = float(min(b_2d.min(), t_2d.min()))
vmax: float = float(max(b_2d.max(), t_2d.max()))
im_b = ax_baseline.imshow(b_2d, aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
ax_baseline.set_title(f"{label} Baseline (10k sampled)")
plt.colorbar(im_b, ax=ax_baseline)
im_t = ax_target.imshow(t_2d, aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
ax_target.set_title(f"{label} Target (10k sampled)")
plt.colorbar(im_t, ax=ax_target)

101
python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py Normal file
View File

@@ -0,0 +1,101 @@
"""Tensor preprocessing and utility functions for visualization."""
from __future__ import annotations
import math
import re
import numpy as np
import torch
_DOWNSAMPLE_THRESHOLD: int = 10_000_000
_SCATTER_SAMPLE_SIZE: int = 10_000
def _preprocess_tensor(tensor: torch.Tensor) -> torch.Tensor:
t: torch.Tensor = tensor.squeeze()
while t.ndim < 2:
t = t.unsqueeze(0)
if t.ndim > 2:
t = t.reshape(-1, t.shape[-1])
t = _reshape_to_balanced_aspect(t)
return t
def _reshape_to_balanced_aspect(
t: torch.Tensor, max_ratio: float = 5.0
) -> torch.Tensor:
assert t.ndim == 2
h, w = t.shape
ratio: float = h / w if w > 0 else float("inf")
if 1 / max_ratio <= ratio <= max_ratio:
return t
total: int = h * w
target_side: int = int(math.sqrt(total))
for new_h in range(target_side, 0, -1):
if total % new_h == 0:
new_w: int = total // new_h
new_ratio: float = new_h / new_w
if 1 / max_ratio <= new_ratio <= max_ratio:
return t.reshape(new_h, new_w)
return t.reshape(1, -1)
# ────────────────────── utility ──────────────────────
def _to_log10(t: torch.Tensor) -> torch.Tensor:
return t.abs().clamp(min=1e-10).log10()
def _format_log_ticks(ax: object, axis: str = "both") -> None:
from matplotlib.ticker import FuncFormatter
formatter = FuncFormatter(
lambda x, _: f"1e{int(x)}" if x == int(x) else f"1e{x:.1f}"
)
if axis in ("x", "both"):
ax.xaxis.set_major_formatter(formatter)
if axis in ("y", "both"):
ax.yaxis.set_major_formatter(formatter)
def _format_stats(name: str, t: torch.Tensor) -> str:
return (
f"{name}: shape={tuple(t.shape)}, "
f"min={t.min().item():.4g}, max={t.max().item():.4g}, "
f"mean={t.mean().item():.4g}, std={t.std().item():.4g}"
)
def _safe_hist(
ax: object, data: np.ndarray, *, bins: int = 100, **kwargs: object
) -> None:
data_f64: np.ndarray = data.astype(np.float64)
try:
ax.hist(data_f64, bins=bins, **kwargs)
except ValueError:
ax.hist(data_f64, bins=max(1, len(np.unique(data_f64[:1000]))), **kwargs)
def _maybe_downsample_numpy(
t: torch.Tensor,
max_elements: int = _DOWNSAMPLE_THRESHOLD,
) -> np.ndarray:
if t.numel() <= max_elements:
return t.numpy()
rng: np.random.Generator = np.random.default_rng(seed=0)
indices: np.ndarray = rng.choice(t.numel(), max_elements, replace=False)
return t.numpy()[indices]
def _sanitize_filename(name: str) -> str:
return re.sub(r"[/\.\s]+", "_", name).strip("_")

2
python/sglang/srt/debug_utils/dump_loader.py
View File

@@ -9,8 +9,6 @@ import torch
LOAD_FAILED: object = object()
LOAD_FAILED: object = object()
def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
stem = Path(path).stem

114
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
View File

@@ -639,119 +639,5 @@ class TestThdCpConcat:
)
class TestThdCpConcat:
def test_single_seq(self) -> None:
"""Single seq THD unshard: 2 ranks → per-seq concat."""
rank0 = torch.tensor([1, 2, 3]).refine_names("t")
rank1 = torch.tensor([4, 5, 6]).refine_names("t")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
expected = torch.tensor([1, 2, 3, 4, 5, 6])
assert torch.equal(result[0].rename(None), expected)
def test_multi_seq(self) -> None:
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
seq_a_r0 = torch.arange(0, 50)
seq_b_r0 = torch.arange(100, 132)
pad_r0 = torch.full((46,), -1)
rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
seq_a_r1 = torch.arange(50, 100)
seq_b_r1 = torch.arange(132, 164)
pad_r1 = torch.full((46,), -2)
rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
# seqB: r0(32) + r1(32) = 64 tokens
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
# pad: r0(46) + r1(46) = 92 tokens
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
def test_with_hidden_dim(self) -> None:
"""THD unshard with trailing hidden dim: shape [T, H]."""
torch.manual_seed(42)
hidden: int = 4
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
seq_a_r0 = torch.randn(3, hidden)
seq_b_r0 = torch.randn(2, hidden)
rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
seq_a_r1 = torch.randn(3, hidden)
seq_b_r1 = torch.randn(2, hidden)
rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert unsharded.shape == (10, hidden)
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
def test_with_leading_batch_dim(self) -> None:
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
torch.manual_seed(42)
batch: int = 2
hidden: int = 4
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
seq_a_r0 = torch.randn(batch, 3, hidden)
seq_b_r0 = torch.randn(batch, 2, hidden)
rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
seq_a_r1 = torch.randn(batch, 3, hidden)
seq_b_r1 = torch.randn(batch, 2, hidden)
rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
groups=[[0, 1]],
)
with warning_sink.context():
result = execute_unsharder_plan(plan, [rank0, rank1])
assert len(result) == 1
unsharded: torch.Tensor = result[0].rename(None)
assert unsharded.shape == (batch, 10, hidden)
# seqA: r0(3) + r1(3) = 6 tokens per batch
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
# seqB: r0(2) + r1(2) = 4 tokens per batch
assert torch.equal(
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

122
test/registered/debug_utils/comparator/test_entrypoint.py
View File

@@ -1019,80 +1019,6 @@ class TestEntrypointAxisSwapper:
assert comp.name == "hidden"
class TestEntrypointAxisSwapper:
"""Test cross-framework dim reordering through the full entrypoint pipeline."""
def test_axis_swap_different_dim_order(self, tmp_path, capsys):
"""Baseline dims 'b h d' vs target dims 'b d h': axis swapper rearranges baseline to match."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
_create_rank_dump(
baseline_dir,
rank=0,
name="hidden",
tensor=full_tensor,
dims="b h d",
)
_create_rank_dump(
target_dir,
rank=0,
name="hidden",
tensor=full_tensor.permute(0, 2, 1).contiguous(),
dims="b d h",
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
assert comp.baseline.shape == [4, 16, 8]
assert comp.target.shape == [4, 16, 8]
def test_axis_swap_with_tp_unshard(self, tmp_path, capsys):
"""Baseline TP=2 with dims 'b h(tp) d' vs target TP=2 with dims 'b d h(tp)': unshard + axis swap."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
baseline_dir = tmp_path / "baseline"
target_dir = tmp_path / "target"
_create_tp_sharded_dumps(
baseline_dir,
full_tensor=full_tensor,
name="hidden",
tp_size=2,
shard_dim=1,
dims_str="b h(tp) d",
)
_create_tp_sharded_dumps(
target_dir,
full_tensor=full_tensor.permute(0, 2, 1).contiguous(),
name="hidden",
tp_size=2,
shard_dim=2,
dims_str="b d h(tp)",
)
args = _make_args(
baseline_dir / _FIXED_EXP_NAME,
target_dir / _FIXED_EXP_NAME,
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
class TestEntrypointReplicatedAxis:
"""Test replicated-axis scenarios through the full entrypoint pipeline."""
@@ -1578,6 +1504,52 @@ class TestEntrypointNonTensorValues:
assert roundtripped.values_equal is True
# ───────────────────── Visualization integration tests ─────────────────────
class TestEntrypointVisualize:
"""Test --visualize-bundle-details integration."""
@pytest.fixture(autouse=True)
def _skip_if_no_matplotlib(self) -> None:
pytest.importorskip("matplotlib")
def test_visualize_creates_pngs(self, tmp_path, capsys):
"""--visualize-bundle-details with --filter produces PNG files."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
viz_dir = tmp_path / "viz_out"
args = _make_args(
baseline_path,
target_path,
grouping="raw",
filter="tensor_a",
viz_bundle_details=True,
viz_output_dir=str(viz_dir),
)
records = _run_and_parse(args, capsys)
assert len(_get_comparisons(records)) == 1
png_files = list(viz_dir.glob("*.png"))
assert len(png_files) == 1
assert png_files[0].stat().st_size > 0
def test_no_visualize_no_png(self, tmp_path, capsys):
"""Without --visualize-bundle-details, no PNGs are created."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
viz_dir = tmp_path / "viz_out"
args = _make_args(
baseline_path,
target_path,
grouping="raw",
viz_bundle_details=False,
viz_output_dir=str(viz_dir),
)
_run_and_parse(args, capsys)
assert not viz_dir.exists() or len(list(viz_dir.glob("*.png"))) == 0
# --------------------------- Assertion helpers -------------------
@@ -1702,6 +1674,8 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
filter=None,
output_format="json",
grouping="logical",
viz_bundle_details=False,
viz_output_dir="/tmp/comparator_viz/",
)
defaults.update(overrides)
return Namespace(**defaults)

203
test/registered/debug_utils/comparator/test_manually_verify.py Normal file
View File

@@ -0,0 +1,203 @@
"""Visual comparison figure tests — CI sanity check + human verification.
This file serves two purposes:
1. CI sanity check: ensures generate_comparison_figure() runs without errors
across various tensor scenarios (registered via register_cpu_ci).
2. Human verification: all generated PNGs are copied to /tmp/comparator_manual_verify/
so they can be pulled back to a local machine for visual inspection.
Run:
python -m pytest test/registered/debug_utils/comparator/test_manually_verify.py -x -v
Human verification:
After running, images are at /tmp/comparator_manual_verify/.
Each test's docstring describes the expected visual appearance.
"""
import shutil
import sys
from pathlib import Path
import pytest
import torch
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=60, suite="default", nightly=True)
_PUBLISH_DIR: Path = Path("/tmp/comparator_manual_verify")
_PNG_MAGIC: bytes = b"\x89PNG"
@pytest.fixture(scope="session")
def publish_dir() -> Path:
"""Fixed output dir for human inspection — files are copied here after generation."""
if _PUBLISH_DIR.exists():
shutil.rmtree(_PUBLISH_DIR)
_PUBLISH_DIR.mkdir(parents=True)
return _PUBLISH_DIR
def _assert_valid_png(path: Path) -> None:
assert path.exists(), f"PNG not created: {path}"
assert path.stat().st_size > 0, f"PNG is empty: {path}"
with open(path, "rb") as f:
magic: bytes = f.read(4)
assert magic == _PNG_MAGIC, f"Not a valid PNG: {path}"
def _generate_and_publish(
*,
baseline: torch.Tensor,
target: torch.Tensor,
name: str,
tmp_path: Path,
publish_dir: Path,
) -> Path:
from sglang.srt.debug_utils.comparator.visualizer import (
generate_comparison_figure,
)
output_path: Path = tmp_path / f"{name}.png"
generate_comparison_figure(
baseline=baseline,
target=target,
name=name,
output_path=output_path,
)
_assert_valid_png(output_path)
shutil.copy2(src=output_path, dst=publish_dir / output_path.name)
return output_path
@pytest.fixture(autouse=True)
def _skip_if_no_matplotlib() -> None:
pytest.importorskip("matplotlib")
class TestManuallyVerify:
def test_normal_small_diff(self, tmp_path: Path, publish_dir: Path) -> None:
"""Two nearly-identical tensors (randn + 0.01 noise).
Expected: All 6 panel rows visible. Diff heatmap nearly uniform light color.
Hist2d tightly clustered along the red diagonal line.
"""
baseline: torch.Tensor = torch.randn(32, 64)
target: torch.Tensor = baseline + torch.randn(32, 64) * 0.01
_generate_and_publish(
baseline=baseline,
target=target,
name="normal_small_diff",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
def test_significant_diff(self, tmp_path: Path, publish_dir: Path) -> None:
"""Two tensors with larger differences (randn + 0.5 noise).
Expected: All 6 panel rows visible. Diff heatmap shows noticeable structure.
Hist2d scatter is broader, spread away from the diagonal.
"""
baseline: torch.Tensor = torch.randn(32, 64)
target: torch.Tensor = baseline + torch.randn(32, 64) * 0.5
_generate_and_publish(
baseline=baseline,
target=target,
name="significant_diff",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
def test_shape_mismatch(self, tmp_path: Path, publish_dir: Path) -> None:
"""Baseline 32x64, target 16x32 — shapes do not match.
Expected: Only 2 panel rows (baseline heatmap, target heatmap).
No diff/histogram/hist2d/sampled panels since diff cannot be computed.
"""
baseline: torch.Tensor = torch.randn(32, 64)
target: torch.Tensor = torch.randn(16, 32)
_generate_and_publish(
baseline=baseline,
target=target,
name="shape_mismatch",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
def test_large_tensor(self, tmp_path: Path, publish_dir: Path) -> None:
"""4000x4000 tensor — triggers internal downsampling.
Expected: Figure renders normally without OOM. Downsampled panels
should still look reasonable.
"""
baseline: torch.Tensor = torch.randn(4000, 4000)
target: torch.Tensor = baseline + torch.randn(4000, 4000) * 0.001
_generate_and_publish(
baseline=baseline,
target=target,
name="large_tensor",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
def test_1d_tensor(self, tmp_path: Path, publish_dir: Path) -> None:
"""1D tensor (256,) — internally reshaped to 2D before plotting.
Expected: All 6 panel rows visible. The heatmap shape reflects the
reshaped 2D form, not the original 1D.
"""
baseline: torch.Tensor = torch.randn(256)
target: torch.Tensor = baseline + 0.01
_generate_and_publish(
baseline=baseline,
target=target,
name="1d_tensor",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
def test_constant_tensor(self, tmp_path: Path, publish_dir: Path) -> None:
"""All-zero baseline, tiny-valued target.
Expected: Colorbar range is extremely small. Histogram concentrates in
a single bin. No rendering errors from near-zero variance.
"""
baseline: torch.Tensor = torch.zeros(32, 64)
target: torch.Tensor = torch.ones(32, 64) * 1e-8
_generate_and_publish(
baseline=baseline,
target=target,
name="constant_tensor",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
def test_extreme_values(self, tmp_path: Path, publish_dir: Path) -> None:
"""Tensor containing values spanning 1e-10 to 1e10.
Expected: Log10 panels handle the wide range gracefully. No inf/nan
artifacts in the rendered figure.
"""
baseline: torch.Tensor = torch.randn(32, 64).abs()
baseline[0, 0] = 1e-10
baseline[0, 1] = 1e10
target: torch.Tensor = baseline + torch.randn(32, 64) * 0.01
_generate_and_publish(
baseline=baseline,
target=target,
name="extreme_values",
tmp_path=tmp_path,
publish_dir=publish_dir,
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

96
test/registered/debug_utils/comparator/test_visualizer.py Normal file
View File

@@ -0,0 +1,96 @@
import sys
from pathlib import Path
import pytest
import torch
from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
_preprocess_tensor,
_reshape_to_balanced_aspect,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=30, suite="default", nightly=True)
class TestPreprocessTensor:
def test_1d_becomes_2d(self) -> None:
t: torch.Tensor = torch.randn(100)
result: torch.Tensor = _preprocess_tensor(t)
assert result.ndim == 2
def test_3d_becomes_2d(self) -> None:
t: torch.Tensor = torch.randn(2, 3, 4)
result: torch.Tensor = _preprocess_tensor(t)
assert result.ndim == 2
assert result.numel() == t.numel()
def test_high_dim_becomes_2d(self) -> None:
t: torch.Tensor = torch.randn(2, 3, 4, 5)
result: torch.Tensor = _preprocess_tensor(t)
assert result.ndim == 2
assert result.numel() == t.numel()
def test_scalar_becomes_2d(self) -> None:
t: torch.Tensor = torch.tensor(3.14)
result: torch.Tensor = _preprocess_tensor(t)
assert result.ndim == 2
assert result.numel() == 1
def test_already_2d_preserves_elements(self) -> None:
t: torch.Tensor = torch.randn(10, 20)
result: torch.Tensor = _preprocess_tensor(t)
assert result.ndim == 2
assert result.numel() == 200
class TestReshapeToBalancedAspect:
def test_extreme_wide_gets_fixed(self) -> None:
t: torch.Tensor = torch.randn(1, 10000)
result: torch.Tensor = _reshape_to_balanced_aspect(t)
h, w = result.shape
ratio: float = max(h, w) / max(min(h, w), 1)
assert ratio <= 5.0
def test_extreme_tall_gets_fixed(self) -> None:
t: torch.Tensor = torch.randn(10000, 1)
result: torch.Tensor = _reshape_to_balanced_aspect(t)
h, w = result.shape
ratio: float = max(h, w) / max(min(h, w), 1)
assert ratio <= 5.0
def test_already_balanced_unchanged(self) -> None:
t: torch.Tensor = torch.randn(100, 100)
result: torch.Tensor = _reshape_to_balanced_aspect(t)
assert result.shape == (100, 100)
def test_preserves_numel(self) -> None:
t: torch.Tensor = torch.randn(1, 7919)
result: torch.Tensor = _reshape_to_balanced_aspect(t)
assert result.numel() == t.numel()
class TestGenerateComparisonFigure:
@pytest.fixture(autouse=True)
def _skip_if_no_matplotlib(self) -> None:
pytest.importorskip("matplotlib")
def test_nested_output_dir(self, tmp_path: Path) -> None:
from sglang.srt.debug_utils.comparator.visualizer import (
generate_comparison_figure,
)
output_path: Path = tmp_path / "a" / "b" / "c" / "nested.png"
generate_comparison_figure(
baseline=torch.randn(10, 10),
target=torch.randn(10, 10),
name="nested",
output_path=output_path,
)
assert output_path.exists()
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))