Support multiple verbosity in dump comparator (#19684)

This commit is contained in:
fzyzcjy
2026-03-02 18:47:30 +08:00
committed by GitHub
parent 3dd4649b42
commit e5ef845cad
7 changed files with 519 additions and 249 deletions
@@ -409,18 +409,61 @@ def _make_traced_plan(
# ---------------------------------------------------------------------------
# Rich format snapshot tests (normal mode only)
# Rich format snapshot tests
# ---------------------------------------------------------------------------
class TestFormatComparisonRichNormal:
"""format_comparison_rich() snapshot tests."""
class TestFormatComparisonRichMinimal:
"""format_comparison_rich() with verbosity='minimal'."""
def test_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True),
)
result: str = format_comparison_rich(record)
result: str = format_comparison_rich(record, verbosity="minimal")
assert result == (
"[green]✅[/] [bold green]hidden_states [/] "
"rel_diff=1.00e-04"
)
def test_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff(rel_diff=0.5, passed=False),
)
result: str = format_comparison_rich(record, verbosity="minimal")
assert result == (
"[red]❌[/] [bold red]hidden_states [/] "
"rel_diff=5.00e-01"
)
def test_shape_mismatch(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
shape_mismatch=True,
)
result: str = format_comparison_rich(record, verbosity="minimal")
assert result == (
"[red]❌[/] [bold red]hidden_states [/] "
"[yellow]shape mismatch[/]"
)
def test_no_diff(self) -> None:
record: TensorComparisonRecord = _make_comparison_record()
result: str = format_comparison_rich(record, verbosity="minimal")
assert result == ("[red]❌[/] [bold red]hidden_states [/]")
class TestFormatComparisonRichNormal:
"""format_comparison_rich() with verbosity='normal'."""
def test_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True),
)
result: str = format_comparison_rich(record, verbosity="normal")
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -439,7 +482,7 @@ class TestFormatComparisonRichNormal:
rel_diff=0.5, max_abs_diff=1.0, mean_abs_diff=0.3, passed=False
),
)
result: str = format_comparison_rich(record)
result: str = format_comparison_rich(record, verbosity="normal")
assert result == (
"[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -459,7 +502,7 @@ class TestFormatComparisonRichNormal:
record: TensorComparisonRecord = _make_comparison_record(
shape_mismatch=True,
)
result: str = format_comparison_rich(record)
result: str = format_comparison_rich(record, verbosity="normal")
assert result == (
"[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -478,7 +521,7 @@ class TestFormatComparisonRichNormal:
diff_downcast=_make_diff(rel_diff=1e-5, passed=True),
downcast_dtype="torch.bfloat16",
)
result: str = format_comparison_rich(record)
result: str = format_comparison_rich(record, verbosity="normal")
assert result == (
"[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -504,7 +547,7 @@ class TestFormatComparisonRichNormal:
diff=_make_diff(passed=True),
raw_bundle_info=bundle_info,
)
result: str = format_comparison_rich(record)
result: str = format_comparison_rich(record, verbosity="normal")
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -526,7 +569,7 @@ class TestFormatComparisonRichNormal:
diff=_make_diff(passed=True),
traced_plan=_make_traced_plan(plan),
)
result: str = format_comparison_rich(record)
result: str = format_comparison_rich(record, verbosity="normal")
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -543,6 +586,113 @@ class TestFormatComparisonRichNormal:
)
class TestFormatComparisonRichVerbose:
"""format_comparison_rich() with verbosity='verbose'."""
def test_passed_full_detail(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True),
sample="tensor([0.1, 0.2, ...])",
)
result: str = format_comparison_rich(record, verbosity="verbose")
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
" [green]rel_diff=1.00e-04[/] max_abs=5.00e-04 mean_abs=2.00e-04\n"
" [dim]Aligned[/]\n"
" [4, 8] vs [4, 8] torch.float32 vs torch.float32\n"
" [dim]Stats[/]\n"
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]\n"
" [dim]Abs Diff Percentiles[/]\n"
" p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04\n"
" [dim]Samples[/]\n"
" baseline tensor([0.1, 0.2, ...])\n"
" target tensor([0.1, 0.2, ...])"
)
def test_with_bundle_verbose(self) -> None:
bundle_info: Pair[BundleSideInfo] = Pair(
x=_make_bundle_side_info(num_files=2, with_parallel_info=True),
y=_make_bundle_side_info(num_files=2, with_parallel_info=True),
)
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff(passed=True),
raw_bundle_info=bundle_info,
)
result: str = format_comparison_rich(record, verbosity="verbose")
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
" [green]rel_diff=1.00e-04[/] max_abs=5.00e-04 mean_abs=2.00e-04\n"
" [dim]Bundle[/]\n"
" baseline [cyan]2 files[/] float32\n"
" [0] [2, 4096] rank=0 tp=0/2\n"
" [1] [2, 4096] rank=1 tp=1/2\n"
" target [cyan]2 files[/] float32\n"
" [0] [2, 4096] rank=0 tp=0/2\n"
" [1] [2, 4096] rank=1 tp=1/2\n"
" [dim]Aligned[/]\n"
" [4, 8] vs [4, 8] torch.float32 vs torch.float32\n"
" [dim]Stats[/]\n"
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]\n"
" [dim]Abs Diff Percentiles[/]\n"
" p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04"
)
def test_with_plan_and_traces(self) -> None:
plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True)
record: TensorComparisonRecord = _make_comparison_record(
diff=_make_diff(passed=True),
traced_plan=_make_traced_plan(
plan,
target_input_shapes=[[2, 4096], [2, 4096]],
target_output_shapes=[[4, 4096]],
),
)
result: str = format_comparison_rich(record, verbosity="verbose")
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
" [green]rel_diff=1.00e-04[/] max_abs=5.00e-04 mean_abs=2.00e-04\n"
" [dim]Plan[/]\n"
" baseline [dim](passthrough)[/]\n"
" target [magenta]unsharder(ParallelAxis.TP)[/] 2×[2, 4096] → 1×[4, 4096]\n"
" [dim]Aligned[/]\n"
" [4, 8] vs [4, 8] torch.float32 vs torch.float32\n"
" [dim]Stats[/]\n"
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]\n"
" [dim]Abs Diff Percentiles[/]\n"
" p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04"
)
class TestFormatBundleSection:
"""_format_bundle_section() snapshot tests."""
@@ -597,6 +747,36 @@ class TestFormatBundleSection:
]
class TestFormatBundleSectionVerbose:
"""_format_bundle_section(verbose=True) snapshot tests."""
def test_per_file_listing(self) -> None:
bundle: Pair[BundleSideInfo] = Pair(
x=_make_bundle_side_info(num_files=2, with_parallel_info=True),
y=_make_bundle_side_info(num_files=2, with_parallel_info=True),
)
lines: list[str] = _format_bundle_section(bundle, verbose=True)
assert lines == [
" baseline [cyan]2 files[/] float32",
" [0] [2, 4096] rank=0 tp=0/2",
" [1] [2, 4096] rank=1 tp=1/2",
" target [cyan]2 files[/] float32",
" [0] [2, 4096] rank=0 tp=0/2",
" [1] [2, 4096] rank=1 tp=1/2",
]
def test_no_files(self) -> None:
empty: BundleSideInfo = BundleSideInfo(num_files=0, files=[])
bundle: Pair[BundleSideInfo] = Pair(x=empty, y=empty)
lines: list[str] = _format_bundle_section(bundle, verbose=True)
assert lines == [
" baseline [dim](no files)[/]",
" target [dim](no files)[/]",
]
class TestFormatPlanSectionRich:
"""_format_plan_section_rich() snapshot tests."""
@@ -719,6 +899,45 @@ class TestFormatStatsRich:
]
class TestFormatStatsRichVerbose:
"""_format_stats_rich(verbose=True) snapshot tests."""
def test_all_stats_with_percentiles(self) -> None:
baseline: TensorStats = _make_stats()
target: TensorStats = _make_stats()
lines: list[str] = _format_stats_rich(
baseline=baseline, target=target, verbose=True
)
assert lines == [
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]",
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]",
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]",
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]",
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]",
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]",
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]",
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]",
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]",
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]",
]
def test_no_percentiles(self) -> None:
baseline: TensorStats = _make_stats(percentiles={})
target: TensorStats = _make_stats(percentiles={})
lines: list[str] = _format_stats_rich(
baseline=baseline, target=target, verbose=True
)
assert lines == [
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]",
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]",
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]",
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]",
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]",
]
class TestFormatAbsDiffPercentilesRich:
"""_format_abs_diff_percentiles_rich() snapshot tests."""
@@ -294,8 +294,8 @@ class TestEntrypointGroupingRaw:
run(parse_args(argv))
output = capsys.readouterr().out
assert "Config:" in output
assert "Summary:" in output
assert "Comparator Config" in output
assert "SUMMARY" in output
def test_text_output_with_failure(self, tmp_path, capsys):
"""Text output with a failed comparison renders failure info."""
@@ -317,7 +317,7 @@ class TestEntrypointGroupingRaw:
run(parse_args(argv))
output = capsys.readouterr().out
assert "Summary:" in output
assert "SUMMARY" in output
assert "failed" in output.lower()
def test_duplicate_dump_pairing(self, tmp_path, capsys):
@@ -4267,102 +4267,6 @@ class TestReportOutput:
assert isinstance(parsed, ConfigRecord)
class TestEntrypointAutoDescend:
"""Test auto-descend: --baseline-path / --target-path pointing to a parent
directory that contains a single subdirectory with .pt files."""
def test_auto_descend_single_engine(self, tmp_path: Path, capsys) -> None:
"""Parent dir wrapping a single engine subdir is auto-descended and comparison succeeds."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
baseline_wrapper: Path = tmp_path / "baseline_wrap"
target_wrapper: Path = tmp_path / "target_wrap"
baseline_wrapper.mkdir()
target_wrapper.mkdir()
baseline_exp.rename(baseline_wrapper / "engine_0")
target_exp.rename(target_wrapper / "engine_0")
argv = _make_argv(baseline_wrapper, target_wrapper, preset="raw")
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 0
_assert_single_comparison_passed(records)
def test_no_descend_when_pt_at_root(self, tmp_path: Path, capsys) -> None:
"""Direct .pt files — no descend needed, comparison still works."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
argv = _make_argv(baseline_exp, target_exp, preset="raw")
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 0
_assert_single_comparison_passed(records)
def test_auto_descend_emits_log_record(self, tmp_path: Path, capsys) -> None:
"""Auto-descend emits a LogRecord with the info message."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
wrapper: Path = tmp_path / "target_wrap"
wrapper.mkdir()
target_exp.rename(wrapper / "engine_0")
argv = _make_argv(baseline_exp, wrapper, preset="raw")
records, _ = _run_and_parse(argv, capsys)
log_records: list[LogRecord] = [r for r in records if isinstance(r, LogRecord)]
auto_descend_msgs: list[str] = [
info.message
for lr in log_records
for info in lr.infos
if "auto-descend" in info.message
]
assert any("target_path" in m for m in auto_descend_msgs)
def test_auto_descend_single_nonempty_among_empty(
self, tmp_path: Path, capsys
) -> None:
"""Two subdirs but only one has .pt — auto-descend picks the non-empty one."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
wrapper: Path = tmp_path / "target_wrap"
wrapper.mkdir()
target_exp.rename(wrapper / "engine_0")
(wrapper / "empty_subdir").mkdir()
argv = _make_argv(baseline_exp, wrapper, preset="raw")
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 0
_assert_single_comparison_passed(records)
def test_error_multiple_nonempty_subdirs(self, tmp_path: Path) -> None:
"""Two subdirs both with .pt — raises ValueError with clear message."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
wrapper: Path = tmp_path / "target_wrap"
wrapper.mkdir()
target_exp.rename(wrapper / "engine_0")
engine_1: Path = wrapper / "engine_1"
engine_1.mkdir()
torch.save(torch.tensor([1.0]), engine_1 / "dummy.pt")
argv: list[str] = _make_argv(baseline_exp, wrapper, preset="raw")
with pytest.raises(ValueError, match="multiple subdirectories contain data"):
run(parse_args(argv))
def test_error_no_data_found(self, tmp_path: Path) -> None:
"""No .pt files anywhere — raises ValueError."""
baseline_exp, _ = _create_dumps(tmp_path, ["tensor_a"])
empty_dir: Path = tmp_path / "empty_target"
empty_dir.mkdir()
(empty_dir / "subdir").mkdir()
argv: list[str] = _make_argv(baseline_exp, empty_dir, preset="raw")
with pytest.raises(ValueError, match="no .pt files found"):
run(parse_args(argv))
class TestEntrypointDpAttentionMissingAlias:
"""Regression: dp-attention without ``# dp:=attn_dp`` → shape mismatch failure.
@@ -4460,5 +4364,101 @@ class TestEntrypointDpAttentionMissingAlias:
assert comparison.category == "failed"
class TestEntrypointAutoDescend:
"""Test auto-descend: --baseline-path / --target-path pointing to a parent
directory that contains a single subdirectory with .pt files."""
def test_auto_descend_single_engine(self, tmp_path: Path, capsys) -> None:
"""Parent dir wrapping a single engine subdir is auto-descended and comparison succeeds."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
baseline_wrapper: Path = tmp_path / "baseline_wrap"
target_wrapper: Path = tmp_path / "target_wrap"
baseline_wrapper.mkdir()
target_wrapper.mkdir()
baseline_exp.rename(baseline_wrapper / "engine_0")
target_exp.rename(target_wrapper / "engine_0")
argv = _make_argv(baseline_wrapper, target_wrapper, preset="raw")
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 0
_assert_single_comparison_passed(records)
def test_no_descend_when_pt_at_root(self, tmp_path: Path, capsys) -> None:
"""Direct .pt files — no descend needed, comparison still works."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
argv = _make_argv(baseline_exp, target_exp, preset="raw")
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 0
_assert_single_comparison_passed(records)
def test_auto_descend_emits_log_record(self, tmp_path: Path, capsys) -> None:
"""Auto-descend emits a LogRecord with the info message."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
wrapper: Path = tmp_path / "target_wrap"
wrapper.mkdir()
target_exp.rename(wrapper / "engine_0")
argv = _make_argv(baseline_exp, wrapper, preset="raw")
records, _ = _run_and_parse(argv, capsys)
log_records: list[LogRecord] = [r for r in records if isinstance(r, LogRecord)]
auto_descend_msgs: list[str] = [
info.message
for lr in log_records
for info in lr.infos
if "auto-descend" in info.message
]
assert any("target_path" in m for m in auto_descend_msgs)
def test_auto_descend_single_nonempty_among_empty(
self, tmp_path: Path, capsys
) -> None:
"""Two subdirs but only one has .pt — auto-descend picks the non-empty one."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
wrapper: Path = tmp_path / "target_wrap"
wrapper.mkdir()
target_exp.rename(wrapper / "engine_0")
(wrapper / "empty_subdir").mkdir()
argv = _make_argv(baseline_exp, wrapper, preset="raw")
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 0
_assert_single_comparison_passed(records)
def test_error_multiple_nonempty_subdirs(self, tmp_path: Path) -> None:
"""Two subdirs both with .pt — raises ValueError with clear message."""
baseline_exp, target_exp = _create_dumps(tmp_path, ["tensor_a"])
wrapper: Path = tmp_path / "target_wrap"
wrapper.mkdir()
target_exp.rename(wrapper / "engine_0")
engine_1: Path = wrapper / "engine_1"
engine_1.mkdir()
torch.save(torch.tensor([1.0]), engine_1 / "dummy.pt")
argv: list[str] = _make_argv(baseline_exp, wrapper, preset="raw")
with pytest.raises(ValueError, match="multiple subdirectories contain data"):
run(parse_args(argv))
def test_error_no_data_found(self, tmp_path: Path) -> None:
"""No .pt files anywhere — raises ValueError."""
baseline_exp, _ = _create_dumps(tmp_path, ["tensor_a"])
empty_dir: Path = tmp_path / "empty_target"
empty_dir.mkdir()
(empty_dir / "subdir").mkdir()
argv: list[str] = _make_argv(baseline_exp, empty_dir, preset="raw")
with pytest.raises(ValueError, match="no .pt files found"):
run(parse_args(argv))
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))