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

Support presets and arbitrary skipping keys in dump comparator (#19676)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:41:49 +08:00
committed by GitHub
parent 2e15c015c0
commit ec44bc82ab
13 changed files with 802 additions and 391 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
View File

@@ -4,12 +4,12 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
ConfigRecord,
GeneralWarning,
ReplicatedCheckResult,
SkipRecord,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
WarningRecord,
parse_record_json,
)
@@ -83,7 +83,7 @@ class TestStrictBase:
class TestRecordTypes:
def test_comparison_record_inherits_tensor_fields(self):
record = ComparisonRecord(
record = TensorComparisonRecord(
name="hidden_states",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -108,8 +108,8 @@ class TestRecordTypes:
"end_step": 100,
}
),
SkipRecord(name="attn", reason="no_baseline"),
ComparisonRecord(
SkipComparisonRecord(name="attn", reason="no_baseline"),
TensorComparisonRecord(
name="mlp",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -148,8 +148,8 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
class TestWarnings:
def test_comparison_record_failed_when_diff_passed_but_warnings(self):
"""ComparisonRecord with diff.passed=True but warnings → category=='failed'."""
record = ComparisonRecord(
"""TensorComparisonRecord with diff.passed=True but warnings → category=='failed'."""
record = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -161,8 +161,8 @@ class TestWarnings:
assert record.category == "failed"
def test_skip_record_failed_when_warnings(self):
"""SkipRecord with warnings → category=='failed' instead of 'skipped'."""
record = SkipRecord(
"""SkipComparisonRecord with warnings → category=='failed' instead of 'skipped'."""
record = SkipComparisonRecord(
name="x",
reason="no_baseline",
warnings=[GeneralWarning(category="test", message="some warning")],
@@ -170,8 +170,8 @@ class TestWarnings:
assert record.category == "failed"
def test_replicated_checks_all_passed(self):
"""ComparisonRecord with all replicated_checks passed → category=='passed'."""
record = ComparisonRecord(
"""TensorComparisonRecord with all replicated_checks passed → category=='passed'."""
record = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -183,8 +183,8 @@ class TestWarnings:
assert record.category == "passed"
def test_replicated_checks_failed_means_record_failed(self):
"""ComparisonRecord with any replicated_check.passed=False → category=='failed'."""
record = ComparisonRecord(
"""TensorComparisonRecord with any replicated_check.passed=False → category=='failed'."""
record = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -196,7 +196,7 @@ class TestWarnings:
assert record.category == "failed"
def test_replicated_check_json_round_trip(self):
"""ReplicatedCheckResult survives JSON round-trip via ComparisonRecord."""
"""ReplicatedCheckResult survives JSON round-trip via TensorComparisonRecord."""
check = _make_replicated_check(
axis="cp",
group_index=2,
@@ -204,7 +204,7 @@ class TestWarnings:
baseline_index=0,
passed=False,
)
record = ComparisonRecord(
record = TensorComparisonRecord(
name="mlp",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -215,7 +215,7 @@ class TestWarnings:
)
restored = parse_record_json(record.model_dump_json())
assert isinstance(restored, ComparisonRecord)
assert isinstance(restored, TensorComparisonRecord)
assert len(restored.replicated_checks) == 1
restored_check: ReplicatedCheckResult = restored.replicated_checks[0]

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

File diff suppressed because it is too large Load Diff

16
test/registered/debug_utils/comparator/test_manually_verify.py
View File

@@ -207,7 +207,9 @@ class TestPerTokenHeatmapManualVerify:
the left (small diff), bright/hot on the right (large diff). Multiple
rows for different tensor names. Colorbar shows log10 scale.
"""
from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord
from sglang.srt.debug_utils.comparator.output_types import (
TensorComparisonRecord,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
)
@@ -220,7 +222,7 @@ class TestPerTokenHeatmapManualVerify:
hidden_dim: int = 128
num_tensors: int = 5
records: list[ComparisonRecord] = []
records: list[TensorComparisonRecord] = []
for i in range(num_tensors):
baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
noise_scale: torch.Tensor = torch.linspace(
@@ -235,7 +237,7 @@ class TestPerTokenHeatmapManualVerify:
diff_threshold=1e-3,
seq_dim=0,
)
records.append(ComparisonRecord(**info.model_dump()))
records.append(TensorComparisonRecord(**info.model_dump()))
output_path: Path = tmp_path / "per_token_increasing_diff.png"
result = generate_per_token_heatmap(records=records, output_path=output_path)
@@ -250,7 +252,9 @@ class TestPerTokenHeatmapManualVerify:
Expected: Heatmap shows one bright vertical stripe at the spike position,
rest is dark/cold.
"""
from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord
from sglang.srt.debug_utils.comparator.output_types import (
TensorComparisonRecord,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
)
@@ -264,7 +268,7 @@ class TestPerTokenHeatmapManualVerify:
spike_pos: int = 32
num_tensors: int = 4
records: list[ComparisonRecord] = []
records: list[TensorComparisonRecord] = []
for i in range(num_tensors):
baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
target: torch.Tensor = baseline.clone()
@@ -277,7 +281,7 @@ class TestPerTokenHeatmapManualVerify:
diff_threshold=1e-3,
seq_dim=0,
)
records.append(ComparisonRecord(**info.model_dump()))
records.append(TensorComparisonRecord(**info.model_dump()))
output_path: Path = tmp_path / "per_token_single_spike.png"
result = generate_per_token_heatmap(records=records, output_path=output_path)

48
test/registered/debug_utils/comparator/test_model_validation.py
View File

@@ -22,11 +22,11 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
)
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
GeneralWarning,
NonTensorRecord,
SkipRecord,
NonTensorComparisonRecord,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
parse_record_json,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -208,9 +208,9 @@ def _make_comparison_record(
*,
diff: DiffInfo | None,
warnings: list | None = None,
) -> ComparisonRecord:
) -> TensorComparisonRecord:
ti: TensorInfo = _make_tensor_info()
return ComparisonRecord(
return TensorComparisonRecord(
name="t",
baseline=ti,
target=ti,
@@ -223,7 +223,7 @@ def _make_comparison_record(
class TestOutputRecordCategories:
def test_skip_record_with_warnings_is_failed(self) -> None:
record = SkipRecord(
record = SkipComparisonRecord(
name="t",
reason="test",
warnings=[GeneralWarning(category="c", message="m")],
@@ -231,28 +231,28 @@ class TestOutputRecordCategories:
assert record.category == "failed"
def test_skip_record_no_warnings_is_skipped(self) -> None:
record = SkipRecord(name="t", reason="test")
record = SkipComparisonRecord(name="t", reason="test")
assert record.category == "skipped"
def test_comparison_record_diff_none_is_failed(self) -> None:
record: ComparisonRecord = _make_comparison_record(diff=None)
record: TensorComparisonRecord = _make_comparison_record(diff=None)
assert record.category == "failed"
def test_comparison_record_passed_with_warnings_is_failed(self) -> None:
record: ComparisonRecord = _make_comparison_record(
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
warnings=[GeneralWarning(category="c", message="m")],
)
assert record.category == "failed"
def test_comparison_record_passed_no_warnings_is_passed(self) -> None:
record: ComparisonRecord = _make_comparison_record(
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
assert record.category == "passed"
def test_non_tensor_record_equal_is_passed(self) -> None:
record = NonTensorRecord(
record = NonTensorComparisonRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.125",
@@ -263,7 +263,7 @@ class TestOutputRecordCategories:
assert record.category == "passed"
def test_non_tensor_record_different_is_failed(self) -> None:
record = NonTensorRecord(
record = NonTensorComparisonRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.25",
@@ -274,7 +274,7 @@ class TestOutputRecordCategories:
assert record.category == "failed"
def test_non_tensor_record_with_warnings_is_failed(self) -> None:
record = NonTensorRecord(
record = NonTensorComparisonRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.125",
@@ -286,7 +286,7 @@ class TestOutputRecordCategories:
assert record.category == "failed"
def test_non_tensor_record_json_roundtrip(self) -> None:
record = NonTensorRecord(
record = NonTensorComparisonRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.25",
@@ -296,14 +296,14 @@ class TestOutputRecordCategories:
)
json_str: str = record.model_dump_json()
roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, NonTensorRecord)
assert isinstance(roundtripped, NonTensorComparisonRecord)
assert roundtripped.name == "sm_scale"
assert roundtripped.values_equal is False
assert roundtripped.baseline_value == "0.125"
assert roundtripped.target_value == "0.25"
def test_non_tensor_record_text_format_equal(self) -> None:
record = NonTensorRecord(
record = NonTensorComparisonRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.125",
@@ -316,7 +316,7 @@ class TestOutputRecordCategories:
assert "[equal]" in text
def test_non_tensor_record_text_format_different(self) -> None:
record = NonTensorRecord(
record = NonTensorComparisonRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.25",
@@ -351,10 +351,10 @@ def _make_aligner_plan() -> AlignerPlan:
)
class TestAlignerPlanInComparisonRecord:
class TestAlignerPlanInTensorComparisonRecord:
def test_comparison_record_with_aligner_plan(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: ComparisonRecord = _make_comparison_record(
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
@@ -363,7 +363,7 @@ class TestAlignerPlanInComparisonRecord:
def test_aligner_plan_json_roundtrip(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: ComparisonRecord = _make_comparison_record(
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
@@ -376,7 +376,7 @@ class TestAlignerPlanInComparisonRecord:
== "unsharder"
)
roundtripped: ComparisonRecord = parse_record_json(json_str)
roundtripped: TensorComparisonRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is not None
assert (
roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
@@ -384,16 +384,16 @@ class TestAlignerPlanInComparisonRecord:
)
def test_comparison_record_without_aligner_plan(self) -> None:
record: ComparisonRecord = _make_comparison_record(
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
json_str: str = record.model_dump_json()
roundtripped: ComparisonRecord = parse_record_json(json_str)
roundtripped: TensorComparisonRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is None
def test_aligner_plan_text_format(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: ComparisonRecord = _make_comparison_record(
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})

14
test/registered/debug_utils/comparator/test_per_token_visualizer.py
View File

@@ -9,7 +9,7 @@ from pathlib import Path
import pytest
import torch
from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord
from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compare_tensor_pair,
)
@@ -31,8 +31,8 @@ def _make_comparison_record(
baseline: torch.Tensor,
target: torch.Tensor,
seq_dim: int = 0,
) -> ComparisonRecord:
"""Build a ComparisonRecord with per-token data from raw tensors."""
) -> TensorComparisonRecord:
"""Build a TensorComparisonRecord with per-token data from raw tensors."""
info = compare_tensor_pair(
x_baseline=baseline,
x_target=target,
@@ -40,7 +40,7 @@ def _make_comparison_record(
diff_threshold=1e-3,
seq_dim=seq_dim,
)
return ComparisonRecord(**info.model_dump())
return TensorComparisonRecord(**info.model_dump())
class TestPerTokenVisualizer:
@@ -68,7 +68,7 @@ class TestPerTokenVisualizer:
name="no_per_token",
diff_threshold=1e-3,
)
record = ComparisonRecord(**info.model_dump())
record = TensorComparisonRecord(**info.model_dump())
output_path: Path = tmp_path / "no_data.png"
result = generate_per_token_heatmap(records=[record], output_path=output_path)
@@ -82,7 +82,7 @@ class TestPerTokenVisualizer:
)
torch.manual_seed(42)
records: list[ComparisonRecord] = [
records: list[TensorComparisonRecord] = [
_make_comparison_record(
name=f"tensor_{i}",
baseline=torch.randn(16, 32),
@@ -108,7 +108,7 @@ class TestPerTokenVisualizer:
)
torch.manual_seed(42)
records: list[ComparisonRecord] = [
records: list[TensorComparisonRecord] = [
_make_comparison_record(
name="short",
baseline=torch.randn(4, 8),

44
test/registered/debug_utils/comparator/test_preset.py Normal file
View File

@@ -0,0 +1,44 @@
import pytest
from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="default", nightly=True)
class TestExpandPreset:
"""Test preset expansion logic."""
def test_explicit_preset(self):
"""--preset sglang_megatron expands into its argv."""
argv = [
"--baseline-path",
"/a",
"--preset",
"sglang_megatron",
"--diff-threshold",
"0.01",
]
result = expand_preset(argv, presets=PRESETS)
assert "--preset" not in result
assert "--grouping-skip-keys" in result
assert "concat_steps" in result
assert "--baseline-path" in result
assert "--diff-threshold" in result
def test_default_preset_applied(self):
"""No --preset and no --grouping-skip-keys triggers default preset."""
argv = ["--baseline-path", "/a"]
result = expand_preset(argv, presets=PRESETS)
assert "--grouping-skip-keys" in result
def test_explicit_skip_keys_prevents_default(self):
"""Explicit --grouping-skip-keys prevents default preset injection."""
argv = ["--grouping-skip-keys", "rank", "--baseline-path", "/a"]
result = expand_preset(argv, presets=PRESETS)
assert result == argv
def test_unknown_preset_raises(self):
"""Unknown preset name raises ValueError."""
with pytest.raises(ValueError, match="Unknown value for --preset"):
expand_preset(["--preset", "nonexistent"], presets=PRESETS)