From 80bbd30909becc3f009771e59642af1f45eb6d72 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 28 Feb 2026 18:08:16 +0800 Subject: [PATCH] Visualize comparison detailed results in dump comparator (#19565) --- .../comparator/bundle_comparator.py | 58 ++++- .../srt/debug_utils/comparator/entrypoint.py | 19 ++ .../comparator/visualizer/__init__.py | 3 + .../comparator/visualizer/figure.py | 116 +++++++++ .../comparator/visualizer/panels.py | 226 ++++++++++++++++++ .../comparator/visualizer/preprocessing.py | 101 ++++++++ python/sglang/srt/debug_utils/dump_loader.py | 2 - .../aligner/unsharder/test_executor.py | 114 --------- .../debug_utils/comparator/test_entrypoint.py | 122 ++++------ .../comparator/test_manually_verify.py | 203 ++++++++++++++++ .../debug_utils/comparator/test_visualizer.py | 96 ++++++++ 11 files changed, 867 insertions(+), 193 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/visualizer/__init__.py create mode 100644 python/sglang/srt/debug_utils/comparator/visualizer/figure.py create mode 100644 python/sglang/srt/debug_utils/comparator/visualizer/panels.py create mode 100644 python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py create mode 100644 test/registered/debug_utils/comparator/test_manually_verify.py create mode 100644 test/registered/debug_utils/comparator/test_visualizer.py diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index b89721b06..58a093211 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -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( diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 1d749fc89..2587a9566 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -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() diff --git a/python/sglang/srt/debug_utils/comparator/visualizer/__init__.py b/python/sglang/srt/debug_utils/comparator/visualizer/__init__.py new file mode 100644 index 000000000..476ddce36 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/visualizer/__init__.py @@ -0,0 +1,3 @@ +from sglang.srt.debug_utils.comparator.visualizer.figure import ( # noqa: F401 + generate_comparison_figure, +) diff --git a/python/sglang/srt/debug_utils/comparator/visualizer/figure.py b/python/sglang/srt/debug_utils/comparator/visualizer/figure.py new file mode 100644 index 000000000..08c919282 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/visualizer/figure.py @@ -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) diff --git a/python/sglang/srt/debug_utils/comparator/visualizer/panels.py b/python/sglang/srt/debug_utils/comparator/visualizer/panels.py new file mode 100644 index 000000000..ff9a6d614 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/visualizer/panels.py @@ -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) diff --git a/python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py b/python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py new file mode 100644 index 000000000..67e1b14b8 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py @@ -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("_") diff --git a/python/sglang/srt/debug_utils/dump_loader.py b/python/sglang/srt/debug_utils/dump_loader.py index b2156bbb0..f35a455c2 100644 --- a/python/sglang/srt/debug_utils/dump_loader.py +++ b/python/sglang/srt/debug_utils/dump_loader.py @@ -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 diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py index eca1fe27a..c1cbf5a29 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 65304b7c9..97751663b 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -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) diff --git a/test/registered/debug_utils/comparator/test_manually_verify.py b/test/registered/debug_utils/comparator/test_manually_verify.py new file mode 100644 index 000000000..0c2b822a1 --- /dev/null +++ b/test/registered/debug_utils/comparator/test_manually_verify.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_visualizer.py b/test/registered/debug_utils/comparator/test_visualizer.py new file mode 100644 index 000000000..8a01187bc --- /dev/null +++ b/test/registered/debug_utils/comparator/test_visualizer.py @@ -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__]))