Trace execution information in dump comparator (#19682)

This commit is contained in:
fzyzcjy
2026-03-02 18:46:27 +08:00
committed by GitHub
parent 3ebd85bf1c
commit 5bf3deb4bc
17 changed files with 1180 additions and 148 deletions

View File

@@ -5,6 +5,8 @@ import torch
from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
AlignerResult,
StepPlansResult,
SubPlansResult,
_execute_step_plans,
execute_aligner_plan,
execute_sub_plan,
@@ -31,25 +33,28 @@ register_cpu_ci(est_time=15, suite="default", nightly=True)
class TestExecuteSubPlans:
def test_empty_tensors_returns_none(self) -> None:
result, checks = execute_sub_plans(tensors=[], plans=[])
assert result is None
assert checks == []
r: SubPlansResult = execute_sub_plans(tensors=[], plans=[])
assert r.tensor is None
assert r.checks == []
assert r.snapshots == []
def test_no_plans_single_tensor_passthrough(self) -> None:
tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
result, checks = execute_sub_plans(tensors=[tensor], plans=[])
assert result is not None
assert torch.equal(result, tensor)
assert checks == []
r: SubPlansResult = execute_sub_plans(tensors=[tensor], plans=[])
assert r.tensor is not None
assert torch.equal(r.tensor, tensor)
assert r.checks == []
assert r.snapshots == []
def test_no_plans_multiple_tensors_returns_none(self) -> None:
tensors: list[torch.Tensor] = [
torch.tensor([1.0]),
torch.tensor([2.0]),
]
result, checks = execute_sub_plans(tensors=tensors, plans=[])
assert result is None
assert checks == []
r: SubPlansResult = execute_sub_plans(tensors=tensors, plans=[])
assert r.tensor is None
assert r.checks == []
assert r.snapshots == []
def test_with_unsharder_plan(self) -> None:
t0: torch.Tensor = torch.tensor([[1.0, 2.0]]).refine_names("b", "h")
@@ -61,12 +66,13 @@ class TestExecuteSubPlans:
groups=[[0, 1]],
)
result, checks = execute_sub_plans(tensors=[t0, t1], plans=[plan])
r: SubPlansResult = execute_sub_plans(tensors=[t0, t1], plans=[plan])
assert result is not None
assert r.tensor is not None
expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
assert torch.equal(result.rename(None), expected)
assert checks == []
assert torch.equal(r.tensor.rename(None), expected)
assert r.checks == []
assert len(r.snapshots) == 1
class TestExecuteSubPlan:
@@ -91,10 +97,13 @@ class TestExecuteStepPlans:
sub_plans=[],
)
result, checks = _execute_step_plans(tensors=tensors, step_plans=[step_plan])
r: StepPlansResult = _execute_step_plans(
tensors=tensors, step_plans=[step_plan]
)
assert result == {}
assert checks == []
assert r.tensors == {}
assert r.checks == []
assert len(r.traced_side.step_plans) == 1
def test_single_step_passthrough(self) -> None:
tensor: torch.Tensor = torch.tensor([1.0, 2.0])
@@ -105,11 +114,15 @@ class TestExecuteStepPlans:
sub_plans=[],
)
result, checks = _execute_step_plans(tensors=[tensor], step_plans=[step_plan])
r: StepPlansResult = _execute_step_plans(
tensors=[tensor], step_plans=[step_plan]
)
assert 5 in result
assert torch.equal(result[5], tensor)
assert checks == []
assert 5 in r.tensors
assert torch.equal(r.tensors[5], tensor)
assert r.checks == []
assert len(r.traced_side.step_plans) == 1
assert r.traced_side.step_plans[0].step == 5
class TestExecuteAlignerPlan:

View File

@@ -1,12 +1,37 @@
import sys
import warnings
from pathlib import Path
warnings.filterwarnings(
"ignore", message="builtin type Swig.*", category=DeprecationWarning
)
# Add the test root to sys.path so `registered.debug_utils.comparator.testing_helpers`
# can be imported by test modules.
_TEST_ROOT: Path = Path(__file__).resolve().parents[3]
if str(_TEST_ROOT) not in sys.path:
sys.path.insert(0, str(_TEST_ROOT))
import pytest
from sglang.srt.debug_utils.comparator.output_types import report_sink
from sglang.srt.debug_utils.comparator.report_sink import report_sink
collect_ignore_glob: list[str] = []
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line(
"filterwarnings",
"ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning",
)
config.addinivalue_line(
"filterwarnings",
"ignore:builtin type Swig.*:DeprecationWarning",
)
config.addinivalue_line(
"filterwarnings",
"ignore:Named tensors and all their associated APIs:UserWarning",
)
collect_ignore_glob: list[str] = []

View File

@@ -9,8 +9,8 @@ import torch
from sglang.srt.debug_utils.comparator.display import (
_collect_input_ids_and_positions,
_collect_rank_info,
_extract_parallel_info,
_render_polars_as_text,
extract_parallel_info,
)
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
@@ -359,26 +359,26 @@ class TestExtractParallelInfo:
"pp_size": 2,
}
row_data: dict = {}
_extract_parallel_info(row_data=row_data, info=info)
extract_parallel_info(row_data=row_data, info=info)
assert row_data["tp"] == "1/4"
assert row_data["pp"] == "0/2"
def test_skips_error_info(self) -> None:
row_data: dict = {}
_extract_parallel_info(
extract_parallel_info(
row_data=row_data, info={"error": True, "tp_rank": 0, "tp_size": 1}
)
assert row_data == {}
def test_skips_empty_info(self) -> None:
row_data: dict = {}
_extract_parallel_info(row_data=row_data, info={})
extract_parallel_info(row_data=row_data, info={})
assert row_data == {}
def test_ignores_rank_without_size(self) -> None:
row_data: dict = {}
_extract_parallel_info(row_data=row_data, info={"tp_rank": 0})
extract_parallel_info(row_data=row_data, info={"tp_rank": 0})
assert "tp" not in row_data

View File

@@ -1458,7 +1458,7 @@ class TestEntrypointConcatMode:
assert len(comparisons) == 3
def test_concat_aligner_plan_fields(self, tmp_path, capsys):
"""TensorComparisonRecord.aligner_plan reports mode='concat' with plan=None."""
"""TensorComparisonRecord.traced_plan reports mode='concat' with plan=None."""
torch.manual_seed(42)
records = self._run_concat(
@@ -1470,8 +1470,9 @@ class TestEntrypointConcatMode:
)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
plan = comparisons[0].aligner_plan
assert plan is not None
traced_plan = comparisons[0].traced_plan
assert traced_plan is not None
plan = traced_plan.plan
assert plan.token_aligner_mode == "concat_steps"
assert plan.token_aligner_plan is None
@@ -4250,7 +4251,7 @@ class TestReportOutput:
def test_streaming_flush(self, tmp_path, capsys):
"""Report file is flushed after each record (readable before close)."""
from sglang.srt.debug_utils.comparator.output_types import report_sink
from sglang.srt.debug_utils.comparator.report_sink import report_sink
report_file: Path = tmp_path / "stream_report.jsonl"
report_sink.configure(
@@ -4362,5 +4363,102 @@ class TestEntrypointAutoDescend:
run(parse_args(argv))
class TestEntrypointDpAttentionMissingAlias:
"""Regression: dp-attention without ``# dp:=attn_dp`` → shape mismatch failure.
In dp-attention mode (tp_size=2, attn_dp_size=2), layer_input is dumped
after prepare_attn which DP-distributes tokens. One rank gets 0 tokens
(shape [0, H]), the other gets all tokens (shape [T, H]).
Without ``# dp:=attn_dp`` in dims, the comparator has no dp_rank/dp_size
to filter on, so it picks one rank via TP pick — potentially the empty
one — causing a shape mismatch with the baseline.
"""
@staticmethod
def _sglang_dp_attn_parallel_info(*, tp_rank: int) -> dict:
return {
"tp_rank": tp_rank,
"tp_size": 2,
"pp_rank": 0,
"pp_size": 1,
"moe_ep_rank": 0,
"moe_ep_size": 1,
"moe_tp_rank": tp_rank,
"moe_tp_size": 2,
"moe_dp_rank": 0,
"moe_dp_size": 1,
"enable_dp_attention": True,
"attn_tp_rank": 0,
"attn_tp_size": 1,
"attn_dp_rank": tp_rank,
"attn_dp_size": 2,
"local_attn_dp_rank": tp_rank,
"local_attn_dp_size": 2,
"attn_cp_rank": 0,
"attn_cp_size": 1,
}
def test_missing_dp_alias_causes_shape_mismatch(
self, tmp_path: Path, capsys
) -> None:
"""dims='t h' (no dp:=attn_dp) → comparator picks empty rank → shape_mismatch failure."""
torch.manual_seed(42)
tensor_data: torch.Tensor = torch.randn(5, 8)
target_data: torch.Tensor = tensor_data + torch.randn(5, 8) * 0.001
for side_name, data in [("baseline", tensor_data), ("target", target_data)]:
side_dir: Path = tmp_path / side_name
side_dir.mkdir()
# Baseline: single rank, no DP attention
if side_name == "baseline":
_create_rank_dump(
side_dir,
rank=0,
name="layer_input",
tensor=data,
dims="t h",
parallel_info={"tp_rank": 0, "tp_size": 1},
framework="sglang",
)
else:
# Target: dp-attention, tp_rank=0 gets 0 tokens, tp_rank=1 gets all
_create_rank_dump(
side_dir,
rank=0,
name="layer_input",
tensor=torch.empty(0, 8),
dims="t h",
parallel_info=self._sglang_dp_attn_parallel_info(tp_rank=0),
framework="sglang",
)
_create_rank_dump(
side_dir,
rank=1,
name="layer_input",
tensor=data,
dims="t h",
parallel_info=self._sglang_dp_attn_parallel_info(tp_rank=1),
framework="sglang",
)
argv: list[str] = _make_argv(
tmp_path / "baseline" / _FIXED_EXP_NAME,
tmp_path / "target" / _FIXED_EXP_NAME,
diff_threshold=1e-3,
)
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 1
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
assert len(comparisons) == 1
comparison: TensorComparisonRecord = comparisons[0]
assert comparison.shape_mismatch is True
assert comparison.diff is None
assert comparison.category == "failed"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

View File

@@ -7,8 +7,8 @@ from sglang.srt.debug_utils.comparator.log_sink import LogSink
from sglang.srt.debug_utils.comparator.output_types import (
ErrorLog,
InfoLog,
report_sink,
)
from sglang.srt.debug_utils.comparator.report_sink import report_sink
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)

View File

@@ -1,18 +1,67 @@
import sys
from io import StringIO
import pytest
from registered.debug_utils.comparator.testing_helpers import make_diff as _make_diff
from registered.debug_utils.comparator.testing_helpers import (
make_tensor_info as _make_tensor_info,
)
from rich.console import Console
from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSidePlan,
TracedStepPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
AlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
ReordererPlan,
ZigzagToNaturalParams,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
TokenAlignerPlan,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
ConcatParams,
UnsharderPlan,
)
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import (
ConfigRecord,
ErrorLog,
InfoLog,
LogRecord,
NonTensorComparisonRecord,
RecordLocation,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
_format_aligner_plan,
_split_logs,
)
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _render_rich(renderable: object) -> str:
buf: StringIO = StringIO()
Console(file=buf, force_terminal=False, width=120).print(renderable)
return buf.getvalue().rstrip("\n")
# ---------------------------------------------------------------------------
# Existing tests (preserved)
# ---------------------------------------------------------------------------
def test_split_logs_mixed_list() -> None:
"""_split_logs correctly partitions a mixed list of ErrorLog and InfoLog."""
errors, infos = _split_logs(
@@ -40,5 +89,523 @@ def test_log_record_to_text_format() -> None:
assert " fyi" in text
class TestLogRecord:
def test_format_body_returns_empty(self) -> None:
record: LogRecord = LogRecord()
assert record._format_body() == ""
def test_format_rich_body_returns_empty(self) -> None:
record: LogRecord = LogRecord()
assert record._format_rich_body() == ""
def test_to_text_empty_no_logs(self) -> None:
record: LogRecord = LogRecord()
assert record.to_text() == ""
def test_to_text_with_errors_and_infos(self) -> None:
record: LogRecord = LogRecord(
errors=[ErrorLog(category="a", message="bad thing")],
infos=[InfoLog(category="b", message="fyi")],
)
text: str = record.to_text()
assert text == "\n ✗ bad thing\n fyi"
# ---------------------------------------------------------------------------
# ConfigRecord
# ---------------------------------------------------------------------------
class TestConfigRecord:
def test_format_body(self) -> None:
record: ConfigRecord = ConfigRecord(config={"a": 1, "b": "two"})
assert record._format_body() == "Config: {'a': 1, 'b': 'two'}"
def test_to_text_with_errors(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"x": 1},
errors=[ErrorLog(category="cfg", message="bad config")],
)
text: str = record.to_text()
assert text.startswith("Config: {'x': 1}")
assert "✗ bad config" in text
# ---------------------------------------------------------------------------
# SkipComparisonRecord
# ---------------------------------------------------------------------------
class TestSkipComparisonRecord:
def test_format_body_no_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="layer.weight",
reason="zero-dim tensor",
)
assert record._format_body() == "Skip: layer.weight (zero-dim tensor)"
def test_format_body_with_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="layer.weight",
reason="scalar",
location=RecordLocation(step=3),
)
assert record._format_body() == "Skip: layer.weight (step=3) (scalar)"
def test_category_skipped(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="x",
reason="r",
)
assert record.category == "skipped"
def test_category_failed(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="x",
reason="r",
errors=[ErrorLog(category="e", message="boom")],
)
assert record.category == "failed"
# ---------------------------------------------------------------------------
# NonTensorComparisonRecord
# ---------------------------------------------------------------------------
class TestNonTensorComparisonRecord:
def test_format_body_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.001",
baseline_type="float",
target_type="float",
values_equal=True,
)
assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]"
def test_format_body_not_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.01",
baseline_type="float",
target_type="float",
values_equal=False,
)
assert record._format_body() == (
"NonTensor: config.lr\n"
" baseline = 0.001 (float)\n"
" target = 0.01 (float)"
)
def test_with_step(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="bias",
baseline_value="True",
target_value="True",
baseline_type="bool",
target_type="bool",
values_equal=True,
location=RecordLocation(step=5),
)
assert "(step=5)" in record._format_body()
def test_category(self) -> None:
passed: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="x",
baseline_value="1",
target_value="1",
baseline_type="int",
target_type="int",
values_equal=True,
)
failed: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="x",
baseline_value="1",
target_value="2",
baseline_type="int",
target_type="int",
values_equal=False,
)
assert passed.category == "passed"
assert failed.category == "failed"
# ---------------------------------------------------------------------------
# SummaryRecord
# ---------------------------------------------------------------------------
class TestSummaryRecord:
def test_format_body(self) -> None:
record: SummaryRecord = SummaryRecord(
total=10,
passed=7,
failed=2,
skipped=1,
)
assert record._format_body() == (
"Summary: 7 passed, 2 failed, 1 skipped (total 10)"
)
def test_validation_error(self) -> None:
with pytest.raises(ValueError, match="total=5 !="):
SummaryRecord(total=5, passed=1, failed=1, skipped=1)
# ---------------------------------------------------------------------------
# TensorComparisonRecord._format_body
# ---------------------------------------------------------------------------
class TestTensorComparisonRecordFormatBody:
def test_basic(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
body: str = record._format_body()
assert body == (
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
def test_with_replicated_checks(self) -> None:
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
replicated_checks=[
ReplicatedCheckResult(
axis="tp",
group_index=0,
compared_index=1,
baseline_index=0,
passed=True,
atol=1e-3,
diff=_make_diff(
rel_diff=1e-6, max_abs_diff=1e-5, mean_abs_diff=1e-6
),
),
],
)
body: str = record._format_body()
assert body == (
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n"
"Replicated checks:\n"
" ✅ axis=tp group=0 idx=1 vs 0: "
"rel_diff=1.000000e-06 max_abs_diff=1.000000e-05 mean_abs_diff=1.000000e-06"
)
def test_with_aligner_plan(self) -> None:
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
)
traced: TracedAlignerPlan = TracedAlignerPlan(
plan=plan,
per_side=Pair(
x=TracedSidePlan(step_plans=[]),
y=TracedSidePlan(step_plans=[]),
),
)
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
traced_plan=traced,
)
body: str = record._format_body()
assert body == (
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n"
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: (no steps)"
)
def test_with_step(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
location=RecordLocation(step=2),
)
body: str = record._format_body()
assert body == (
"[step=2] "
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
# ---------------------------------------------------------------------------
# _format_aligner_plan
# ---------------------------------------------------------------------------
def _wrap_plan(plan: AlignerPlan) -> TracedAlignerPlan:
"""Wrap an AlignerPlan into a TracedAlignerPlan with no snapshots."""
baseline_traced_steps: list[TracedStepPlan] = [
TracedStepPlan(
step=sp.step,
input_object_indices=sp.input_object_indices,
sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans],
)
for sp in plan.per_step_plans.x
]
target_traced_steps: list[TracedStepPlan] = [
TracedStepPlan(
step=sp.step,
input_object_indices=sp.input_object_indices,
sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans],
)
for sp in plan.per_step_plans.y
]
return TracedAlignerPlan(
plan=plan,
per_side=Pair(
x=TracedSidePlan(step_plans=baseline_traced_steps),
y=TracedSidePlan(step_plans=target_traced_steps),
),
)
class TestFormatAlignerPlan:
def test_passthrough(self) -> None:
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: (no steps)"
)
def test_unsharder(self) -> None:
unsharder: UnsharderPlan = UnsharderPlan(
axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"),
groups=[[0, 1]],
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(
x=[],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
)
],
),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: unsharder]"
)
def test_reorderer(self) -> None:
reorderer: ReordererPlan = ReordererPlan(
params=ZigzagToNaturalParams(dim_name="s", cp_size=2),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(
x=[],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0], sub_plans=[reorderer]
)
],
),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: reorderer]"
)
def test_multi_step(self) -> None:
unsharder: UnsharderPlan = UnsharderPlan(
axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"),
groups=[[0, 1]],
)
reorderer: ReordererPlan = ReordererPlan(
params=ZigzagToNaturalParams(dim_name="s", cp_size=2),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(
x=[],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
),
AlignerPerStepPlan(
step=1, input_object_indices=[0], sub_plans=[reorderer]
),
],
),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: unsharder; step=1: reorderer]"
)
def test_with_token_aligner(self) -> None:
ta_plan: TokenAlignerPlan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]),
y=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
token_aligner_plan=ta_plan,
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: (no steps)\n"
" token_aligner: 3 tokens aligned"
)
def test_with_axis_aligner(self) -> None:
aa_plan: AxisAlignerPlan = AxisAlignerPlan(
pattern=Pair(x="b s d -> s b d", y=None),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
axis_aligner_plan=aa_plan,
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: (no steps)\n"
" axis_aligner: x: b s d -> s b d"
)
# ---------------------------------------------------------------------------
# _OutputRecord log attachment (to_text)
# ---------------------------------------------------------------------------
class TestOutputRecordLogAttachment:
def test_to_text_no_logs(self) -> None:
record: ConfigRecord = ConfigRecord(config={"a": 1})
text: str = record.to_text()
assert text == "Config: {'a': 1}"
def test_to_text_errors_only(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"a": 1},
errors=[ErrorLog(category="x", message="err1")],
)
text: str = record.to_text()
assert text == "Config: {'a': 1}\n ✗ err1"
def test_to_text_infos_only(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"a": 1},
infos=[InfoLog(category="x", message="note1")],
)
text: str = record.to_text()
assert text == "Config: {'a': 1}\n note1"
def test_to_text_mixed(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"a": 1},
errors=[ErrorLog(category="x", message="err1")],
infos=[InfoLog(category="y", message="note1")],
)
text: str = record.to_text()
assert text == "Config: {'a': 1}\n ✗ err1\n note1"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

View File

@@ -0,0 +1,84 @@
"""Shared test helpers for comparator tests."""
from __future__ import annotations
from typing import Optional
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorInfo,
TensorStats,
)
DEFAULT_PERCENTILES: dict[int, float] = {
1: -1.8,
5: -1.5,
50: 0.0,
95: 1.5,
99: 1.8,
}
DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = {
1: 0.0001,
5: 0.0001,
50: 0.0002,
95: 0.0004,
99: 0.0005,
}
def make_stats(
mean: float = 0.0,
abs_mean: float = 0.8,
std: float = 1.0,
min: float = -2.0,
max: float = 2.0,
percentiles: Optional[dict[int, float]] = None,
) -> TensorStats:
return TensorStats(
mean=mean,
abs_mean=abs_mean,
std=std,
min=min,
max=max,
percentiles=percentiles if percentiles is not None else DEFAULT_PERCENTILES,
)
def make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
abs_diff_percentiles: Optional[dict[int, float]] = None,
diff_threshold: float = 1e-3,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_diff,
abs_diff_percentiles=(
abs_diff_percentiles
if abs_diff_percentiles is not None
else DEFAULT_ABS_DIFF_PERCENTILES
),
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=diff_threshold,
passed=passed,
)
def make_tensor_info(
shape: Optional[list[int]] = None,
dtype: str = "torch.float32",
stats: Optional[TensorStats] = None,
sample: Optional[str] = None,
) -> TensorInfo:
return TensorInfo(
shape=shape if shape is not None else [4, 8],
dtype=dtype,
stats=stats if stats is not None else make_stats(),
sample=sample,
)