Visualize comparison detailed results in dump comparator (#19565)

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

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__]))

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)

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__]))

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__]))