diff --git a/miles/utils/test_utils/comparisons/__init__.py b/miles/utils/test_utils/comparisons/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/test_utils/comparisons/comparators.py b/miles/utils/test_utils/comparisons/comparators.py new file mode 100644 index 00000000000..4051c1804ff --- /dev/null +++ b/miles/utils/test_utils/comparisons/comparators.py @@ -0,0 +1,55 @@ +import subprocess +import sys +from pathlib import Path + + +def run_comparator( + *, + baseline_path: Path, + target_path: Path, + diff_thresholds: list[tuple[str, str]], + allow_skipped_pattern: str, + allow_failed_pattern: str, + grouping_skip_keys: list[str] | None = None, + extra_args: list[str] | None, +) -> subprocess.CompletedProcess[str]: + # Skip 'rank' when grouping bundles: under FT (target) and non-FT (baseline) the same + # logical (pp_rank, cp_rank, ep_rank, tp_rank) coordinate maps to a different absolute + # rank ID (e.g. baseline rank=4 vs target cell0 rank=2 for PP=1, CP=0). Without skipping + # 'rank' the comparator gets `baseline_load_failed` for every tensor and fails with rc=1. + # Callers may pass extra keys (e.g. no_failure skips 'dp'/'edp' too). (Grouping is a + # comparator-matching detail, not a pass/fail threshold.) + skip_keys: list[str] = list(grouping_skip_keys) if grouping_skip_keys is not None else ["rank"] + assert "rank" in skip_keys, f"grouping_skip_keys must include 'rank', got {skip_keys}" + + cmd: list[str] = [ + sys.executable, + "-m", + "sglang.srt.debug_utils.comparator", + "--baseline-path", + str(baseline_path), + "--target-path", + str(target_path), + "--output-format", + "json", + "--grouping-skip-keys", + *skip_keys, + "--allow-skipped-pattern", + allow_skipped_pattern, + "--allow-failed-pattern", + allow_failed_pattern, + ] + if extra_args: + cmd.extend(extra_args) + # Keep --diff-threshold strictly last: its nargs="*" greedily consumes every + # following token, so no flag with a bare value may come after it. + cmd.append("--diff-threshold") + for pattern, predicate in diff_thresholds: + cmd.extend([pattern, predicate]) + + result: subprocess.CompletedProcess[str] = subprocess.run( + cmd, + text=True, + ) + + return result diff --git a/miles/utils/test_utils/comparisons/dumps.py b/miles/utils/test_utils/comparisons/dumps.py new file mode 100644 index 00000000000..84810671b5b --- /dev/null +++ b/miles/utils/test_utils/comparisons/dumps.py @@ -0,0 +1,67 @@ +from pathlib import Path + +from miles.utils.test_utils.comparisons.comparators import run_comparator + +# Shared regexes for model-input / metadata tensors that are not weights or grads to +# compare. Exposed as named constants (not as defaults) so each test passes them +# explicitly -- every pass/fail knob is visible at the call site, nothing is implicit. +INPUT_TENSORS_SKIP_PATTERN: str = "input_ids|positions|cu_seqlens_q|cu_seqlens_kv|qkv_format|.*witness.*" +INPUT_TENSORS_ALLOW_FAILED_PATTERN: str = "input_ids|positions|cu_seqlens_q|cu_seqlens_kv|qkv_format" + + +def compare_dumps( + baseline_dir: str, + target_dir: str, + *, + diff_thresholds: list[tuple[str, str]], + allow_skipped_pattern: str, + allow_failed_pattern: str, + phase_subdir: str | None = None, + grouping_skip_keys: list[str] | None = None, + extra_args: list[str] | None = None, +) -> None: + subdir = phase_subdir or "" + baseline_root = Path(baseline_dir) / "dumps" / subdir + target_root = Path(target_dir) / "dumps" / subdir + + assert baseline_root.exists(), f"Baseline dump dir does not exist: {baseline_root}" + assert target_root.exists(), f"Target dump dir does not exist: {target_root}" + + # Dumps are segmented into leaf dirs (e.g. fwd_bwd/rollout_), each a flat set of + # .pt files with its own per-leaf step numbering. The sglang comparator compares one + # flat dir at a time, so compare each matching leaf pair independently. + baseline_leaves = _find_leaf_dump_dirs(baseline_root) + target_leaves = _find_leaf_dump_dirs(target_root) + + assert baseline_leaves, f"No .pt dump files found under {baseline_root}" + assert baseline_leaves == target_leaves, ( + f"Dump leaf-dir mismatch: baseline={baseline_leaves} vs target={target_leaves} " + f"(under {baseline_root} vs {target_root})" + ) + + failed_leaves: list[str] = [] + for leaf in baseline_leaves: + result = run_comparator( + baseline_path=baseline_root / leaf, + target_path=target_root / leaf, + diff_thresholds=diff_thresholds, + allow_skipped_pattern=allow_skipped_pattern, + allow_failed_pattern=allow_failed_pattern, + grouping_skip_keys=grouping_skip_keys, + extra_args=extra_args, + ) + if result.returncode != 0: + failed_leaves.append(leaf) + + assert not failed_leaves, ( + f"Dump comparator failed (rc!=0) for {len(failed_leaves)}/{len(baseline_leaves)} leaf dir(s): " + f"{failed_leaves} (baseline {baseline_root} vs target {target_root}). The comparator applies the " + f"per-tensor predicates ({diff_thresholds}) and the allow/skip patterns itself; see " + f"comparator_report.jsonl under {target_root}/ for the offending tensors." + ) + print(f"Dump comparison passed: {len(baseline_leaves)} leaf dir(s) under {baseline_root} vs {target_root}") + + +def _find_leaf_dump_dirs(root: Path) -> list[str]: + leaves: set[str] = {str(p.parent.relative_to(root)) for p in root.rglob("*.pt")} + return sorted(leaves) diff --git a/miles/utils/test_utils/comparisons/inference_engine_checksums.py b/miles/utils/test_utils/comparisons/inference_engine_checksums.py new file mode 100644 index 00000000000..57e5ab3c644 --- /dev/null +++ b/miles/utils/test_utils/comparisons/inference_engine_checksums.py @@ -0,0 +1,65 @@ +from pathlib import Path + +from miles.utils.event_analyzer.rules import inference_engine_weight_checksum_consistency +from miles.utils.event_analyzer.rules.checksum_compare import ChecksumMismatchIssue, compare_flat_dicts +from miles.utils.event_logger.logger import read_events +from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent + + +def compare_inference_engine_checksums(baseline_dir: str, target_dir: str) -> None: + baseline = _read_inference_engine_checksum_events(Path(baseline_dir)) + target = _read_inference_engine_checksum_events(Path(target_dir)) + assert baseline, f"No InferenceEngineWeightChecksumEvents found in baseline dir: {baseline_dir}" + assert target, f"No InferenceEngineWeightChecksumEvents found in target dir: {target_dir}" + + # Each side's engines must already agree internally (same invariant as the production rule), so + # one representative engine per rollout then proves baseline == target regardless of engine count. + assert not inference_engine_weight_checksum_consistency.check( + baseline + ), "Baseline engines disagree with each other" + assert not inference_engine_weight_checksum_consistency.check(target), "Target engines disagree with each other" + + baseline_by_rollout = _checksums_by_rollout_id(baseline) + target_by_rollout = _checksums_by_rollout_id(target) + assert baseline_by_rollout.keys() == target_by_rollout.keys(), ( + f"Engine checksum rollout_id sets differ: " + f"baseline={sorted(baseline_by_rollout)} " + f"vs target={sorted(target_by_rollout)}" + ) + + mismatches: list[ChecksumMismatchIssue] = [] + for rollout_id in sorted(baseline_by_rollout): + mismatches += list( + compare_flat_dicts( + a=baseline_by_rollout[rollout_id], + b=target_by_rollout[rollout_id], + label_a=f"baseline/rollout_{rollout_id}", + label_b=f"target/rollout_{rollout_id}", + ) + ) + assert not mismatches, "Engine weight checksum baseline-vs-target mismatch:\n" + "\n".join( + f" - {m.label_a} vs {m.label_b} key {m.key}: {m.value_a} != {m.value_b}" for m in mismatches + ) + print(f"Engine weight checksum comparison passed: {len(baseline_by_rollout)} rollout(s) compared") + + +def _checksums_by_rollout_id(events: list[InferenceEngineWeightChecksumEvent]) -> dict[int, dict[str, str]]: + by_rollout: dict[int, dict[str, str]] = {} + for event in events: + if event.rollout_id is None: + continue + assert ( + event.rollout_id not in by_rollout + ), f"Duplicate InferenceEngineWeightChecksumEvent for rollout {event.rollout_id}" + assert event.engine_checksums, f"No engine checksums for rollout {event.rollout_id}" + by_rollout[event.rollout_id] = event.engine_checksums[0] + return by_rollout + + +def _read_inference_engine_checksum_events(dump_dir: Path) -> list[InferenceEngineWeightChecksumEvent]: + """Read all InferenceEngineWeightChecksumEvents from the events directory.""" + events_dir: Path = dump_dir / "events" + if not events_dir.exists(): + return [] + all_events = read_events(events_dir) + return [e for e in all_events if isinstance(e, InferenceEngineWeightChecksumEvent)] diff --git a/tests/fast/utils/test_utils/comparisons/__init__.py b/tests/fast/utils/test_utils/comparisons/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/utils/test_utils/comparisons/test_dumps.py b/tests/fast/utils/test_utils/comparisons/test_dumps.py new file mode 100644 index 00000000000..59716c2c305 --- /dev/null +++ b/tests/fast/utils/test_utils/comparisons/test_dumps.py @@ -0,0 +1,45 @@ +from pathlib import Path + +from miles.utils.test_utils.comparisons.dumps import _find_leaf_dump_dirs + + +class TestFindLeafDumpDirs: + def test_two_pt_files_in_one_leaf_yield_single_entry(self, tmp_path: Path) -> None: + """Multiple .pt files sharing one leaf dir dedup to a single relative entry.""" + leaf = tmp_path / "fwd_bwd" / "rollout_0" + leaf.mkdir(parents=True) + (leaf / "step_0.pt").touch() + (leaf / "step_1.pt").touch() + + assert _find_leaf_dump_dirs(tmp_path) == ["fwd_bwd/rollout_0"] + + def test_two_leaves_returned_sorted(self, tmp_path: Path) -> None: + """Distinct leaf dirs are returned sorted by their relative path string.""" + leaf_b = tmp_path / "leaf_b" + leaf_a = tmp_path / "leaf_a" + leaf_b.mkdir() + leaf_a.mkdir() + (leaf_b / "x.pt").touch() + (leaf_a / "y.pt").touch() + + assert _find_leaf_dump_dirs(tmp_path) == ["leaf_a", "leaf_b"] + + def test_pt_file_directly_in_root_yields_dot(self, tmp_path: Path) -> None: + """A .pt file directly under root has parent equal to root, reported as '.'.""" + (tmp_path / "step_0.pt").touch() + + assert _find_leaf_dump_dirs(tmp_path) == ["."] + + def test_no_pt_files_yields_empty_list(self, tmp_path: Path) -> None: + """A tree with no .pt files produces an empty list.""" + (tmp_path / "sub").mkdir() + + assert _find_leaf_dump_dirs(tmp_path) == [] + + def test_non_pt_files_are_ignored(self, tmp_path: Path) -> None: + """Files not matching *.pt (including *.pth) are ignored by the glob.""" + (tmp_path / "notes.txt").touch() + (tmp_path / "weights.pth").touch() + (tmp_path / "data.json").touch() + + assert _find_leaf_dump_dirs(tmp_path) == [] diff --git a/tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py b/tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py new file mode 100644 index 00000000000..416cfe4078d --- /dev/null +++ b/tests/fast/utils/test_utils/comparisons/test_inference_engine_checksums.py @@ -0,0 +1,149 @@ +"""Tests for test_utils.comparisons.inference_engine_checksums.compare_inference_engine_checksums.""" + +from pathlib import Path +from typing import Any + +import pytest + +from miles.utils.event_logger.logger import EventLogger +from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent +from miles.utils.process_identity import MainProcessIdentity +from miles.utils.test_utils.comparisons.inference_engine_checksums import compare_inference_engine_checksums + + +def _write_inference_engine_events(side_dir: Path, partials: list[dict[str, Any]]) -> None: + events_dir = side_dir / "events" + event_logger = EventLogger(log_dir=events_dir, source=MainProcessIdentity()) + for partial in partials: + event_logger.log(InferenceEngineWeightChecksumEvent, partial, print_log=False) + event_logger.close() + + +def _partial(*, rollout_id: int | None, engine_checksums: list[dict[str, str]]) -> dict[str, Any]: + return dict(rollout_id=rollout_id, engine_checksums=engine_checksums) + + +class TestCompareInferenceEngineChecksums: + def test_identical_passes(self, tmp_path: Path) -> None: + """Internally-consistent sides with equal representative checksums pass.""" + partials = [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}])] + _write_inference_engine_events(tmp_path / "baseline", partials) + _write_inference_engine_events(tmp_path / "target", partials) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_differing_engine_counts_still_pass(self, tmp_path: Path) -> None: + """Engine count may differ between sides; only internal agreement + representative equality matter.""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + _write_inference_engine_events( + tmp_path / "target", + [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}, {"rank0/w": "aaa"}])], + ) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_none_rollout_id_skipped(self, tmp_path: Path) -> None: + """The initial out-of-loop sync (rollout_id=None) is not compared: it differs here yet the + per-rollout checksums match, so the comparison still passes.""" + _write_inference_engine_events( + tmp_path / "baseline", + [ + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_baseline"}]), + _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]), + ], + ) + _write_inference_engine_events( + tmp_path / "target", + [ + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_target"}]), + _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]), + ], + ) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_recurring_none_across_phases_skipped(self, tmp_path: Path) -> None: + """A multi-phase resume yields several None events per side; all are skipped, so a side with + more None events than the other still passes when the per-rollout checksums match.""" + _write_inference_engine_events( + tmp_path / "baseline", + [ + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_a"}]), + _partial(rollout_id=2, engine_checksums=[{"rank0/w": "aaa"}]), + _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_b"}]), + _partial(rollout_id=5, engine_checksums=[{"rank0/w": "bbb"}]), + ], + ) + _write_inference_engine_events( + tmp_path / "target", + [ + _partial(rollout_id=2, engine_checksums=[{"rank0/w": "aaa"}]), + _partial(rollout_id=5, engine_checksums=[{"rank0/w": "bbb"}]), + ], + ) + + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_baseline_engines_disagree_fails(self, tmp_path: Path) -> None: + """If baseline's own engines disagree, the comparison fails (caught by the consistency rule).""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])] + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + + with pytest.raises(AssertionError, match="Baseline engines disagree"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_target_engines_disagree_fails(self, tmp_path: Path) -> None: + """If target's own engines disagree, the comparison fails.""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])] + ) + + with pytest.raises(AssertionError, match="Target engines disagree"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_representative_mismatch_fails(self, tmp_path: Path) -> None: + """Internally-consistent sides whose representatives differ fail and name the tensor.""" + _write_inference_engine_events( + tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "zzz"}])] + ) + + with pytest.raises(AssertionError, match=r"key rank0/w"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_missing_rollout_fails(self, tmp_path: Path) -> None: + """A rollout present only on one side fails closed.""" + _write_inference_engine_events( + tmp_path / "baseline", + [ + _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]), + _partial(rollout_id=2, engine_checksums=[{"rank0/w": "ccc"}]), + ], + ) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + + with pytest.raises(AssertionError, match="rollout_id sets differ"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target")) + + def test_empty_baseline_fails(self, tmp_path: Path) -> None: + """No baseline events fails closed rather than vacuously passing.""" + _write_inference_engine_events(tmp_path / "baseline", []) + _write_inference_engine_events( + tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])] + ) + + with pytest.raises(AssertionError, match="No InferenceEngineWeightChecksumEvents found in baseline"): + compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))