Trace execution information in dump comparator (#19682)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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] = []
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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__]))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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__]))
|
||||
|
||||
84
test/registered/debug_utils/comparator/testing_helpers.py
Normal file
84
test/registered/debug_utils/comparator/testing_helpers.py
Normal 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,
|
||||
)
|
||||
Reference in New Issue
Block a user