diff --git a/python/sglang/srt/debug_utils/comparator/__init__.py b/python/sglang/srt/debug_utils/comparator/__init__.py index 90e89f6ec..2d757adce 100644 --- a/python/sglang/srt/debug_utils/comparator/__init__.py +++ b/python/sglang/srt/debug_utils/comparator/__init__.py @@ -1,6 +1,6 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( # noqa: F401 AlignerPlan, ) -from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord +from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord -ComparisonRecord.model_rebuild() +TensorComparisonRecord.model_rebuild() diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py index 9ccc27e82..8824f62d4 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py @@ -48,15 +48,15 @@ def compute_maybe_token_aligner_result( args: argparse.Namespace, dfs: Pair[pl.DataFrame], ) -> TokenAlignerResult: - if args.grouping != "logical": + token_aligner_mode: Optional[TokenAlignerMode] = getattr( + args, "token_aligner", None + ) + + if token_aligner_mode is None: return TokenAlignerResult( mode=None, plan=None, thd_seq_lens_by_step_pair=_NONE_THD ) - token_aligner_mode: TokenAlignerMode = getattr( - args, "token_aligner", "concat_steps" - ) - if token_aligner_mode == "concat_steps": thd_pair: Pair[Optional[dict[int, list[int]]]] = _load_thd_seq_lens_pair( args=args, dfs=dfs diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 0c4c3ef31..d2f3c4a46 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -28,10 +28,10 @@ from sglang.srt.debug_utils.comparator.dims import ( from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider from sglang.srt.debug_utils.comparator.output_types import ( - ComparisonRecord, GeneralWarning, - NonTensorRecord, - SkipRecord, + NonTensorComparisonRecord, + SkipComparisonRecord, + TensorComparisonRecord, ) from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( compare_tensor_pair, @@ -58,7 +58,7 @@ def compare_bundle_pair( viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, meta_overrider: Optional[MetaOverrider] = None, -) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]: +) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]: with warning_sink.context() as collected_warnings: result = _compare_bundle_pair_inner( name=name, @@ -92,7 +92,7 @@ def _compare_bundle_pair_inner( viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, meta_overrider: Optional[MetaOverrider] = None, -) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]: +) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]: # 1. Load all successfully loaded values all_pair: Pair[list[ValueWithMeta]] = Pair( x=_load_all_values(filenames=filenames_pair.x, base_path=baseline_path), @@ -101,7 +101,7 @@ def _compare_bundle_pair_inner( if not all_pair.x or not all_pair.y: reason = "baseline_load_failed" if not all_pair.x else "target_load_failed" - return SkipRecord(name=name, reason=reason) + return SkipComparisonRecord(name=name, reason=reason) # 1b. Dims override: patch meta["dims"] before DP filter reads it # (--override-dims may add ``# dp:=moe_dp``, so it must run first) @@ -171,10 +171,10 @@ def _compare_bundle_pair_tensor_type( ), viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, -) -> Union[ComparisonRecord, SkipRecord]: +) -> Union[TensorComparisonRecord, SkipComparisonRecord]: if not valid_pair.x or not valid_pair.y: reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed" - return SkipRecord(name=name, reason=reason) + return SkipComparisonRecord(name=name, reason=reason) # Plan (meta only, no tensor) metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map( @@ -207,7 +207,7 @@ def _compare_bundle_pair_tensor_type( assert aligner_result.failed_side_xy is not None side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy] reason: str = f"{side_name}_load_failed" - return SkipRecord(name=name, reason=reason) + return SkipComparisonRecord(name=name, reason=reason) # Resolve seq_dim for per-token computation seq_dim: Optional[int] = ( @@ -225,7 +225,7 @@ def _compare_bundle_pair_tensor_type( diff_threshold=diff_threshold, seq_dim=seq_dim, ) - record = ComparisonRecord( + record = TensorComparisonRecord( **info.model_dump(), aligner_plan=plan, replicated_checks=replicated_checks, @@ -292,7 +292,7 @@ def _compare_bundle_pair_non_tensor_type( *, name: str, value_pair: Pair[list[ValueWithMeta]], -) -> NonTensorRecord: +) -> NonTensorComparisonRecord: baseline_value: Any = value_pair.x[0].value target_value: Any = value_pair.y[0].value @@ -301,7 +301,7 @@ def _compare_bundle_pair_non_tensor_type( except Exception: values_equal = False - return NonTensorRecord( + return NonTensorComparisonRecord( name=name, baseline_value=repr(baseline_value), target_value=repr(target_value), diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 8d290804a..162c72ff2 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -26,22 +26,26 @@ from sglang.srt.debug_utils.comparator.bundle_matcher import ( from sglang.srt.debug_utils.comparator.display import emit_display_records from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider from sglang.srt.debug_utils.comparator.output_types import ( - ComparisonRecord, ConfigRecord, - NonTensorRecord, - SkipRecord, + NonTensorComparisonRecord, + RecordLocation, + SkipComparisonRecord, SummaryRecord, + TensorComparisonRecord, report_sink, ) from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, ) +from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset from sglang.srt.debug_utils.comparator.utils import Pair from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path +_DEFAULT_SKIP_KEYS: set[str] = {"dump_index", "filename"} + def main() -> None: - args = _parse_args() + args = parse_args(sys.argv[1:]) sys.exit(run(args)) @@ -76,9 +80,7 @@ def run(args: argparse.Namespace) -> int: bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles( dfs=dfs, - skip_keys=_compute_skip_keys( - args, has_token_aligner=ta_result.mode is not None - ), + skip_keys=_compute_skip_keys(args), ) viz_output_dir: Optional[Path] = ( @@ -202,13 +204,8 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]: return Pair(x=df_baseline, y=df_target) -def _compute_skip_keys(args, *, has_token_aligner: bool) -> set[str]: - skip_keys: set[str] = {"dump_index", "filename"} - if args.grouping == "logical": - skip_keys |= {"rank", "recompute_status"} - if has_token_aligner: - skip_keys |= {"step"} - return skip_keys +def _compute_skip_keys(args: argparse.Namespace) -> set[str]: + return _DEFAULT_SKIP_KEYS | set(args.grouping_skip_keys or []) def _compare_bundle_pairs( @@ -223,7 +220,9 @@ def _compare_bundle_pairs( viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, meta_overrider: Optional[MetaOverrider] = None, -) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]: +) -> Iterator[ + Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord] +]: for bundle_info_pair in bundle_info_pairs: if not bundle_info_pair.y: continue @@ -232,7 +231,9 @@ def _compare_bundle_pairs( filenames_pair: Pair[list[str]] = bundle_info_pair.map( lambda infos: [info.filename for info in infos] ) - yield compare_bundle_pair( + record: Union[ + TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord + ] = compare_bundle_pair( name=name, filenames_pair=filenames_pair, baseline_path=baseline_path, @@ -246,22 +247,33 @@ def _compare_bundle_pairs( meta_overrider=meta_overrider, ) + target_steps: set[int] = {info.step for info in bundle_info_pair.y} + step: Optional[int] = target_steps.pop() if len(target_steps) == 1 else None + if step is not None: + record = record.model_copy(update={"location": RecordLocation(step=step)}) + + yield record + def _consume_comparison_records( *, - comparison_records: Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]], + comparison_records: Iterator[ + Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord] + ], visualize_per_token: Optional[Path] = None, ) -> tuple[SummaryRecord, list[str]]: counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} - collected_comparisons: list[ComparisonRecord] = [] + collected_comparisons: list[TensorComparisonRecord] = [] skipped_names: list[str] = [] for record in comparison_records: counts[record.category] += 1 report_sink.add(record) - if isinstance(record, SkipRecord) and record.category == "skipped": + if isinstance(record, SkipComparisonRecord) and record.category == "skipped": skipped_names.append(record.name) - if visualize_per_token is not None and isinstance(record, ComparisonRecord): + if visualize_per_token is not None and isinstance( + record, TensorComparisonRecord + ): collected_comparisons.append(record) summary: SummaryRecord = SummaryRecord(total=sum(counts.values()), **counts) @@ -276,7 +288,10 @@ def _consume_comparison_records( return summary, skipped_names -def _parse_args() -> argparse.Namespace: +def parse_args(argv: list[str]) -> argparse.Namespace: + """Parse CLI arguments from an argv list. Applies preset expansion.""" + argv = expand_preset(argv, presets=PRESETS) + parser = argparse.ArgumentParser() parser.add_argument("--baseline-path", type=str) parser.add_argument("--target-path", type=str) @@ -294,18 +309,28 @@ def _parse_args() -> argparse.Namespace: help="Output format: text (default) or json (JSONL, one JSON object per line)", ) parser.add_argument( - "--grouping", + "--preset", type=str, - choices=["logical", "raw"], - default="logical", - help="Grouping mode: logical (cross-rank unshard) or raw (rank-by-rank)", + choices=list(PRESETS.keys()), + default=None, + help="Preset configuration (expanded before parsing). " + f"Available: {list(PRESETS.keys())}", + ) + parser.add_argument( + "--grouping-skip-keys", + nargs="*", + default=None, + help="Metadata keys to skip when grouping bundles (additive on top of " + "always-skipped dump_index and filename). " + "E.g. '--grouping-skip-keys rank step' skips rank and step.", ) parser.add_argument( "--token-aligner", type=str, choices=["smart", "concat_steps"], - default="concat_steps", - help="Token aligner mode: concat_steps (BS=1, no aux needed) or smart (BS>1, sequence matching)", + default=None, + help="Token aligner mode: concat_steps (BS=1, no aux needed) or smart (BS>1, sequence matching). " + "Default None (per-step comparison).", ) parser.add_argument( "--tokenizer", @@ -374,4 +399,4 @@ def _parse_args() -> argparse.Namespace: "Pass empty string '' to disable.", ) - return parser.parse_args() + return parser.parse_args(argv) diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index f75179276..cd0ca602f 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -61,6 +61,24 @@ class _OutputRecord(_StrictBase): return body +class RecordLocation(_StrictBase): + step: Optional[int] = None + + +class _BaseComparisonRecord(_OutputRecord): + location: RecordLocation = Field(default_factory=RecordLocation) + + def _format_location_prefix(self) -> str: + if self.location.step is not None: + return f"[step={self.location.step}] " + return "" + + def _format_location_suffix(self) -> str: + if self.location.step is not None: + return f" (step={self.location.step})" + return "" + + class ConfigRecord(_OutputRecord): type: Literal["config"] = "config" config: dict[str, Any] @@ -74,7 +92,7 @@ class ConfigRecord(_OutputRecord): return f"Config: {self.config}" -class SkipRecord(_OutputRecord): +class SkipComparisonRecord(_BaseComparisonRecord): type: Literal["skip"] = "skip" name: str reason: str @@ -86,7 +104,7 @@ class SkipRecord(_OutputRecord): return "skipped" def _format_body(self) -> str: - return f"Skip: {self.name} ({self.reason})" + return f"Skip: {self.name}{self._format_location_suffix()} ({self.reason})" class _TableRecord(_OutputRecord): @@ -118,7 +136,7 @@ class InputIdsRecord(_TableRecord): return f"{self.label} input_ids & positions" -class ComparisonRecord(TensorComparisonInfo, _OutputRecord): +class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord): model_config = ConfigDict(extra="forbid", defer_build=True) type: Literal["comparison"] = "comparison" @@ -134,7 +152,7 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord): return "passed" if self.diff is not None and self.diff.passed else "failed" def _format_body(self) -> str: - body: str = format_comparison(self) + body: str = self._format_location_prefix() + format_comparison(self) if self.replicated_checks: body += "\n" + format_replicated_checks(self.replicated_checks) if self.aligner_plan is not None: @@ -142,7 +160,7 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord): return body -class NonTensorRecord(_OutputRecord): +class NonTensorComparisonRecord(_BaseComparisonRecord): type: Literal["non_tensor"] = "non_tensor" name: str baseline_value: str @@ -158,10 +176,11 @@ class NonTensorRecord(_OutputRecord): return "passed" if self.values_equal else "failed" def _format_body(self) -> str: + suffix: str = self._format_location_suffix() if self.values_equal: - return f"NonTensor: {self.name} = {self.baseline_value} ({self.baseline_type}) [equal]" + return f"NonTensor: {self.name}{suffix} = {self.baseline_value} ({self.baseline_type}) [equal]" return ( - f"NonTensor: {self.name}\n" + f"NonTensor: {self.name}{suffix}\n" f" baseline = {self.baseline_value} ({self.baseline_type})\n" f" target = {self.target_value} ({self.target_type})" ) @@ -237,9 +256,9 @@ AnyRecord = Annotated[ ConfigRecord, RankInfoRecord, InputIdsRecord, - SkipRecord, - ComparisonRecord, - NonTensorRecord, + SkipComparisonRecord, + TensorComparisonRecord, + NonTensorComparisonRecord, SummaryRecord, WarningRecord, ], diff --git a/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py b/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py index 899b50777..705fd0490 100644 --- a/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py +++ b/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py @@ -9,12 +9,12 @@ from __future__ import annotations from pathlib import Path from typing import Optional -from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord +from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord def generate_per_token_heatmap( *, - records: list[ComparisonRecord], + records: list[TensorComparisonRecord], output_path: Path, ) -> Optional[Path]: """Generate a per-token relative difference heatmap PNG. @@ -31,7 +31,7 @@ def generate_per_token_heatmap( def _collect_per_token_data( *, - records: list[ComparisonRecord], + records: list[TensorComparisonRecord], ) -> list[tuple[str, list[float]]]: rows: list[tuple[str, list[float]]] = [] for record in records: diff --git a/python/sglang/srt/debug_utils/comparator/preset.py b/python/sglang/srt/debug_utils/comparator/preset.py new file mode 100644 index 000000000..dc315de99 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/preset.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +PRESETS: dict[str, list[str]] = { + "raw": [ + "--grouping-skip-keys", + ], + "sglang_dev": [ + "--grouping-skip-keys", + "rank", + ], + "sglang_megatron": [ + "--grouping-skip-keys", + "rank", + "step", + "--token-aligner", + "concat_steps", + ], +} + +DEFAULT_PRESET: str = "sglang_dev" + + +def expand_preset(argv: list[str], presets: dict[str, list[str]]) -> list[str]: + """Expand ``--preset `` into the corresponding argv fragment. + + If ``--preset`` is absent **and** ``--grouping-skip-keys`` is also absent, + the DEFAULT_PRESET is applied automatically. + """ + if (expanded := _expand_flag(argv, "--preset", presets)) is not None: + return expanded + + if "--grouping-skip-keys" not in argv: + return presets[DEFAULT_PRESET] + argv + + return argv + + +def _expand_flag( + argv: list[str], flag: str, mapping: dict[str, list[str]] +) -> list[str] | None: + """Replace ``flag `` in *argv* with the corresponding argv fragment from *mapping*.""" + if flag not in argv: + return None + + idx: int = argv.index(flag) + name: str = argv[idx + 1] + if name not in mapping: + raise ValueError( + f"Unknown value for {flag}: {name}. Available: {list(mapping.keys())}" + ) + + return argv[:idx] + mapping[name] + argv[idx + 2 :] diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py index 7dddc1783..ae8adedc9 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py @@ -4,12 +4,12 @@ import sys import pytest from sglang.srt.debug_utils.comparator.output_types import ( - ComparisonRecord, ConfigRecord, GeneralWarning, ReplicatedCheckResult, - SkipRecord, + SkipComparisonRecord, SummaryRecord, + TensorComparisonRecord, WarningRecord, parse_record_json, ) @@ -83,7 +83,7 @@ class TestStrictBase: class TestRecordTypes: def test_comparison_record_inherits_tensor_fields(self): - record = ComparisonRecord( + record = TensorComparisonRecord( name="hidden_states", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -108,8 +108,8 @@ class TestRecordTypes: "end_step": 100, } ), - SkipRecord(name="attn", reason="no_baseline"), - ComparisonRecord( + SkipComparisonRecord(name="attn", reason="no_baseline"), + TensorComparisonRecord( name="mlp", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -148,8 +148,8 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult: class TestWarnings: def test_comparison_record_failed_when_diff_passed_but_warnings(self): - """ComparisonRecord with diff.passed=True but warnings → category=='failed'.""" - record = ComparisonRecord( + """TensorComparisonRecord with diff.passed=True but warnings → category=='failed'.""" + record = TensorComparisonRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -161,8 +161,8 @@ class TestWarnings: assert record.category == "failed" def test_skip_record_failed_when_warnings(self): - """SkipRecord with warnings → category=='failed' instead of 'skipped'.""" - record = SkipRecord( + """SkipComparisonRecord with warnings → category=='failed' instead of 'skipped'.""" + record = SkipComparisonRecord( name="x", reason="no_baseline", warnings=[GeneralWarning(category="test", message="some warning")], @@ -170,8 +170,8 @@ class TestWarnings: assert record.category == "failed" def test_replicated_checks_all_passed(self): - """ComparisonRecord with all replicated_checks passed → category=='passed'.""" - record = ComparisonRecord( + """TensorComparisonRecord with all replicated_checks passed → category=='passed'.""" + record = TensorComparisonRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -183,8 +183,8 @@ class TestWarnings: assert record.category == "passed" def test_replicated_checks_failed_means_record_failed(self): - """ComparisonRecord with any replicated_check.passed=False → category=='failed'.""" - record = ComparisonRecord( + """TensorComparisonRecord with any replicated_check.passed=False → category=='failed'.""" + record = TensorComparisonRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -196,7 +196,7 @@ class TestWarnings: assert record.category == "failed" def test_replicated_check_json_round_trip(self): - """ReplicatedCheckResult survives JSON round-trip via ComparisonRecord.""" + """ReplicatedCheckResult survives JSON round-trip via TensorComparisonRecord.""" check = _make_replicated_check( axis="cp", group_index=2, @@ -204,7 +204,7 @@ class TestWarnings: baseline_index=0, passed=False, ) - record = ComparisonRecord( + record = TensorComparisonRecord( name="mlp", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -215,7 +215,7 @@ class TestWarnings: ) restored = parse_record_json(record.model_dump_json()) - assert isinstance(restored, ComparisonRecord) + assert isinstance(restored, TensorComparisonRecord) assert len(restored.replicated_checks) == 1 restored_check: ReplicatedCheckResult = restored.replicated_checks[0] diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 52f516ca6..06e27ca11 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -8,16 +8,20 @@ import pytest import torch import sglang.srt.debug_utils.dumper as _dumper_module -from sglang.srt.debug_utils.comparator.entrypoint import _compute_exit_code, run +from sglang.srt.debug_utils.comparator.entrypoint import ( + _compute_exit_code, + parse_args, + run, +) from sglang.srt.debug_utils.comparator.output_types import ( AnyRecord, - ComparisonRecord, ConfigRecord, GeneralWarning, - NonTensorRecord, + NonTensorComparisonRecord, ReplicatedCheckResult, - SkipRecord, + SkipComparisonRecord, SummaryRecord, + TensorComparisonRecord, WarningRecord, _OutputRecord, parse_record_json, @@ -33,14 +37,14 @@ _FIXED_EXP_NAME = "my_exp_name" class TestEntrypointGroupingRaw: - """Test `--grouping raw` scenarios""" + """Test `--grouping-skip-keys` empty (raw) scenarios""" def test_run_basic(self, tmp_path, capsys): - """Two matching tensors produce ConfigRecord, 2 ComparisonRecords, and SummaryRecord.""" + """Two matching tensors produce ConfigRecord, 2 TensorComparisonRecords, and SummaryRecord.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) assert isinstance(records[0], ConfigRecord) assert len(_get_comparisons(records)) == 2 @@ -51,24 +55,24 @@ class TestEntrypointGroupingRaw: assert summary.skipped == 0 def test_filter(self, tmp_path, capsys): - """--filter selects only the matching tensor, producing 1 ComparisonRecord.""" + """--filter selects only the matching tensor, producing 1 TensorComparisonRecord.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) - args = _make_args(baseline_path, target_path, filter="tensor_a", grouping="raw") + argv = _make_argv(baseline_path, target_path, filter="tensor_a", preset="raw") - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) assert len(_get_comparisons(records)) == 1 def test_no_baseline_skip(self, tmp_path, capsys): - """Target tensor missing from baseline emits a SkipRecord with reason baseline_load_failed.""" + """Target tensor missing from baseline emits a SkipComparisonRecord with reason baseline_load_failed.""" baseline_path, target_path = _create_dumps( tmp_path, tensor_names=["tensor_a", "tensor_extra"], baseline_names=["tensor_a"], ) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, _ = _run_and_parse(args, capsys) - skips = [r for r in records if isinstance(r, SkipRecord)] + records, _ = _run_and_parse(argv, capsys) + skips = [r for r in records if isinstance(r, SkipComparisonRecord)] assert len(skips) == 1 assert skips[0].reason == "baseline_load_failed" @@ -79,11 +83,11 @@ class TestEntrypointGroupingRaw: def test_step_range(self, tmp_path, capsys): """--start_step/--end_step restricts comparison to a single step out of three.""" baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=3) - args = _make_args( - baseline_path, target_path, start_step=1, end_step=1, grouping="raw" + argv = _make_argv( + baseline_path, target_path, start_step=1, end_step=1, preset="raw" ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.total == 1 @@ -91,13 +95,13 @@ class TestEntrypointGroupingRaw: def test_all_valid_records(self, tmp_path, capsys): """Every emitted JSON record is a valid _OutputRecord subclass.""" baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=2) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) assert all(isinstance(r, _OutputRecord) for r in records) def test_comparison_failed(self, tmp_path, capsys): - """Completely different tensors produce a failed ComparisonRecord.""" + """Completely different tensors produce a failed TensorComparisonRecord.""" torch.manual_seed(42) baseline_path = _create_rank_dump( tmp_path / "baseline", rank=0, name="tensor_a", tensor=torch.randn(10, 10) @@ -108,11 +112,9 @@ class TestEntrypointGroupingRaw: name="tensor_a", tensor=torch.randn(10, 10) * 100, ) - args = _make_args( - baseline_path, target_path, grouping="raw", diff_threshold=1e-3 - ) + argv = _make_argv(baseline_path, target_path, preset="raw", diff_threshold=1e-3) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].diff is not None @@ -132,9 +134,9 @@ class TestEntrypointGroupingRaw: target_path = _create_rank_dump( tmp_path / "target", rank=0, name="tensor_a", tensor=torch.randn(4, 10) ) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].shape_mismatch is True @@ -158,9 +160,9 @@ class TestEntrypointGroupingRaw: target_path = _create_rank_dump( tmp_path / "target", rank=0, name="tensor_a", tensor=target_tensor ) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 @@ -186,11 +188,9 @@ class TestEntrypointGroupingRaw: target_path = _create_rank_dump( tmp_path / "target", rank=0, name="tensor_a", tensor=target_tensor ) - args = _make_args( - baseline_path, target_path, grouping="raw", diff_threshold=0.01 - ) + argv = _make_argv(baseline_path, target_path, preset="raw", diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].diff_downcast is not None @@ -221,14 +221,14 @@ class TestEntrypointGroupingRaw: _create_rank_dump(target_dir, rank=0, name="different", tensor=different_target) _create_rank_dump(target_dir, rank=0, name="extra", tensor=extra_tensor) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.passed == 1 @@ -239,17 +239,20 @@ class TestEntrypointGroupingRaw: def test_filter_empty_result(self, tmp_path, capsys): """--filter matching nothing produces summary with total=0.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) - args = _make_args( - baseline_path, target_path, filter="nonexistent_pattern", grouping="raw" + argv = _make_argv( + baseline_path, + target_path, + filter="nonexistent_pattern", + preset="raw", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.total == 0 def test_raw_multi_rank(self, tmp_path, capsys): - """Two ranks in raw grouping produce two ComparisonRecords (one per rank).""" + """Two ranks in raw grouping produce two TensorComparisonRecords (one per rank).""" torch.manual_seed(42) tensor = torch.randn(4, 4) @@ -265,14 +268,14 @@ class TestEntrypointGroupingRaw: tensor=tensor + torch.randn(4, 4) * 0.0001, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 2 @@ -284,12 +287,12 @@ class TestEntrypointGroupingRaw: def test_text_output_smoke(self, tmp_path, capsys): """Text output format renders without errors and contains Config/Summary sections.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) - args = _make_args( - baseline_path, target_path, output_format="text", grouping="raw" + argv = _make_argv( + baseline_path, target_path, output_format="text", preset="raw" ) capsys.readouterr() - run(args) + run(parse_args(argv)) output = capsys.readouterr().out assert "Config:" in output @@ -307,12 +310,12 @@ class TestEntrypointGroupingRaw: name="tensor_a", tensor=torch.randn(10, 10) * 100, ) - args = _make_args( - baseline_path, target_path, output_format="text", grouping="raw" + argv = _make_argv( + baseline_path, target_path, output_format="text", preset="raw" ) capsys.readouterr() - run(args) + run(parse_args(argv)) output = capsys.readouterr().out assert "Summary:" in output @@ -343,13 +346,13 @@ class TestEntrypointGroupingRaw: dumper.dump("tensor_a", tensor_v1) dumper.step() - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 2 assert all(c.diff is not None and c.diff.passed for c in comparisons) @@ -361,14 +364,14 @@ class TestEntrypointGroupingRaw: class TestEntrypointGroupingLogical: - """Test `--grouping logical` scenarios""" + """Test `--grouping-skip-keys rank` (logical) scenarios""" def test_no_dims_single_rank(self, tmp_path, capsys): """Single-rank dumps without dims fall back to raw loading.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) - args = _make_args(baseline_path, target_path) + argv = _make_argv(baseline_path, target_path) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) assert len(_get_comparisons(records)) == 2 summary = records[-1] assert isinstance(summary, SummaryRecord) @@ -401,9 +404,9 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp)", ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -438,9 +441,9 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp)", ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) _assert_single_comparison_passed(records) def test_one_side_dims_single_baseline(self, tmp_path, capsys): @@ -465,9 +468,9 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp)", ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) _assert_single_comparison_passed(records) @pytest.mark.parametrize( @@ -478,7 +481,7 @@ class TestEntrypointGroupingLogical: ], ) def test_ambiguous_no_dims_skip(self, tmp_path, capsys, bad_side, expected_reason): - """Multi-rank without dims on one side produces a SkipRecord with the appropriate reason.""" + """Multi-rank without dims on one side produces a SkipComparisonRecord with the appropriate reason.""" torch.manual_seed(42) tensor = torch.randn(4, 8) @@ -492,13 +495,13 @@ class TestEntrypointGroupingLogical: for rank, shard in [(0, tensor[:, :4]), (1, tensor[:, 4:])]: _create_rank_dump(bad_dir, rank=rank, name="hidden", tensor=shard) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, ) - records, _ = _run_and_parse(args, capsys) - skips = [r for r in records if isinstance(r, SkipRecord)] + records, _ = _run_and_parse(argv, capsys) + skips = [r for r in records if isinstance(r, SkipComparisonRecord)] assert len(skips) == 1 assert skips[0].reason == expected_reason @@ -534,9 +537,9 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp)", ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.total == 2 @@ -571,9 +574,14 @@ class TestEntrypointGroupingLogical: num_steps=2, ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv( + baseline_path, + target_path, + diff_threshold=0.01, + preset="sglang_megatron", + ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 # concat along dim 0 (fallback, no token dim) → 2 steps × [4, 8] = [8, 8] @@ -608,13 +616,13 @@ class TestEntrypointGroupingLogical: parallel_info={"cp_rank": cp_rank, "cp_size": 2}, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "attn_out" @@ -645,14 +653,14 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, filter="t_a", diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].name == "t_a" @@ -691,13 +699,13 @@ class TestEntrypointGroupingLogical: tensor=single_tensor + torch.randn(4, 4) * 0.0001, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 2 assert all(c.diff is not None and c.diff.passed for c in comparisons) @@ -732,13 +740,13 @@ class TestEntrypointGroupingLogical: dims_str="b s(cp) h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -771,13 +779,13 @@ class TestEntrypointGroupingLogical: dims_str="b s h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) _assert_single_comparison_passed(records) def test_ep_cp_tp_three_axis_unshard(self, tmp_path, capsys): @@ -806,13 +814,13 @@ class TestEntrypointGroupingLogical: dims_str="b e(ep) s(cp) h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -840,13 +848,13 @@ class TestEntrypointGroupingLogical: dims_str="b s(cp:zigzag) h", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "attn_out" @@ -874,13 +882,13 @@ class TestEntrypointGroupingLogical: dims_str="b s(cp:zigzag) h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -901,13 +909,14 @@ class TestEntrypointGroupingLogical: recompute_tensor=tensor.clone(), ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, + grouping_skip_keys=["rank", "recompute_status"], diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -929,13 +938,14 @@ class TestEntrypointGroupingLogical: recompute_tensor=mismatched_tensor, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, + grouping_skip_keys=["rank", "recompute_status"], diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 @@ -969,9 +979,9 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp:partial)", ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "attn_out" @@ -1000,9 +1010,9 @@ class TestEntrypointGroupingLogical: dims_str="b h(tp:partial)", ) - args = _make_args(baseline_path, target_path, diff_threshold=0.01) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "attn_out" @@ -1036,13 +1046,13 @@ class TestEntrypointGroupingLogical: ) rank += 1 - args = _make_args( + argv = _make_argv( tmp_path / "baseline" / _FIXED_EXP_NAME, tmp_path / "target" / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -1068,17 +1078,88 @@ class TestEntrypointGroupingLogical: dims_str="b s(cp:zigzag,sp) h", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" +class TestEntrypointPerStepMode: + """Test per-step comparison mode (sglang_dev preset behavior).""" + + def test_multi_step_per_step_comparison(self, tmp_path, capsys): + """Multiple steps produce one TensorComparisonRecord per step with step field set.""" + torch.manual_seed(42) + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=3) + argv = _make_argv(baseline_path, target_path, diff_threshold=0.1) + + records, _ = _run_and_parse(argv, capsys) + comparisons = _get_comparisons(records) + assert len(comparisons) == 3 + + steps: list[int] = sorted(c.location.step for c in comparisons) + assert steps == [0, 1, 2] + assert all(c.diff is not None and c.diff.passed for c in comparisons) + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.total == 3 + assert summary.passed == 3 + + def test_per_step_with_tp_unshard(self, tmp_path, capsys): + """Per-step mode with TP=2: each step independently unsharded and compared.""" + torch.manual_seed(42) + full_tensor = torch.randn(4, 8) + + baseline_dir = tmp_path / "baseline" + target_dir = tmp_path / "target" + + baseline_path = _create_tp_sharded_dumps( + baseline_dir, + full_tensor=full_tensor, + name="hidden", + tp_size=2, + shard_dim=1, + dims_str="b h(tp)", + num_steps=2, + ) + target_path = _create_tp_sharded_dumps( + target_dir, + full_tensor=full_tensor + torch.randn(4, 8) * 0.0001, + name="hidden", + tp_size=2, + shard_dim=1, + dims_str="b h(tp)", + num_steps=2, + ) + + argv = _make_argv(baseline_path, target_path, diff_threshold=0.01) + + records, _ = _run_and_parse(argv, capsys) + comparisons = _get_comparisons(records) + assert len(comparisons) == 2 + + steps: list[int] = sorted(c.location.step for c in comparisons) + assert steps == [0, 1] + assert all(c.diff is not None and c.diff.passed for c in comparisons) + assert all(c.baseline.shape == [4, 8] for c in comparisons) + + def test_single_step_has_step_field(self, tmp_path, capsys): + """Single step produces TensorComparisonRecord with location.step=0.""" + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=1) + argv = _make_argv(baseline_path, target_path) + + records, _ = _run_and_parse(argv, capsys) + comparisons = _get_comparisons(records) + assert len(comparisons) == 1 + assert comparisons[0].location.step == 0 + + class TestEntrypointConcatMode: """Test concat token-aligner mode through the full entrypoint pipeline.""" @@ -1135,10 +1216,13 @@ class TestEntrypointConcatMode: name=name, dims=dims, ) - args: Namespace = _make_args( - baseline_path, target_path, diff_threshold=diff_threshold + argv: list[str] = _make_argv( + baseline_path, + target_path, + diff_threshold=diff_threshold, + preset="sglang_megatron", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) return records def test_concat_multi_step_different_data(self, tmp_path, capsys): @@ -1196,13 +1280,14 @@ class TestEntrypointConcatMode: dims_str="b h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, + preset="sglang_megatron", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 # 2 steps × [4, 8] concat along dim 0 (fallback) → [8, 8] @@ -1358,13 +1443,13 @@ class TestEntrypointConcatMode: extra_dumps=[("input_ids", input_ids), ("positions", positions)], ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) # all 3 tensors should be compared (not filtered out) names = {c.name for c in comparisons} @@ -1374,7 +1459,7 @@ class TestEntrypointConcatMode: assert len(comparisons) == 3 def test_concat_aligner_plan_fields(self, tmp_path, capsys): - """ComparisonRecord.aligner_plan reports mode='concat' with plan=None.""" + """TensorComparisonRecord.aligner_plan reports mode='concat' with plan=None.""" torch.manual_seed(42) records = self._run_concat( @@ -1453,13 +1538,14 @@ class TestEntrypointConcatMode: parallel_info={"cp_rank": cp_rank, "cp_size": 2}, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, + preset="sglang_megatron", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 # CP unshard: [4,4,6] × 2 ranks → [4,8,6] per step @@ -1505,16 +1591,16 @@ class TestEntrypointConcatMode: num_steps=num_steps, ) - args: Namespace = _make_args( + argv: list[str] = _make_argv( baseline_path, target_path, - token_aligner="concat_steps", + preset="sglang_megatron", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) - comparisons: list[ComparisonRecord] = _get_comparisons(records) - hidden_comparisons: list[ComparisonRecord] = [ + comparisons: list[TensorComparisonRecord] = _get_comparisons(records) + hidden_comparisons: list[TensorComparisonRecord] = [ c for c in comparisons if c.name == "hidden_states" ] assert len(hidden_comparisons) >= 1 @@ -1547,13 +1633,13 @@ class TestEntrypointAxisAligner: dims="b d h", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" assert comp.baseline.shape == [4, 16, 8] @@ -1584,13 +1670,13 @@ class TestEntrypointAxisAligner: dims_str="b d h(tp)", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" @@ -1617,13 +1703,13 @@ class TestEntrypointAxisAligner: dims="t 1 h", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.name == "hidden" assert comp.baseline.shape == [4, 8] @@ -1656,13 +1742,13 @@ class TestEntrypointReplicatedAxis: dims_str="b s(cp) d", ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comp = _assert_single_comparison_passed(records) assert comp.warnings == [] assert all(c.passed for c in comp.replicated_checks) @@ -1695,13 +1781,13 @@ class TestEntrypointReplicatedAxis: tp_noise=0.5, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.01, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].category == "failed" @@ -1741,13 +1827,13 @@ class TestEntrypointReplicatedAxis: tp_noise=0.5, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, diff_threshold=0.5, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 @@ -1764,7 +1850,7 @@ class TestEntrypointReplicatedAxis: class TestEntrypointAlignment: - """Test `--grouping logical` with token alignment (aux tensors present).""" + """Test smart token alignment with aux tensors.""" def test_sglang_multi_step_alignment(self, tmp_path, capsys): """SGLang multi-step dumps with aux tensors auto-trigger alignment.""" @@ -1807,10 +1893,13 @@ class TestEntrypointAlignment: exp_paths.append(d / _FIXED_EXP_NAME) - args = _make_args( - exp_paths[0], exp_paths[1], grouping="logical", token_aligner="smart" + argv = _make_argv( + exp_paths[0], + exp_paths[1], + grouping_skip_keys=["rank", "step"], + token_aligner="smart", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) # AUX_NAMES are filtered out after plan computation → only hidden_states remains @@ -1917,14 +2006,14 @@ class TestEntrypointAlignment: megatron_dumper.step() # --- Run comparison --- - args = _make_args( + argv = _make_argv( sglang_dir / _FIXED_EXP_NAME, megatron_dir / _FIXED_EXP_NAME, - grouping="logical", + grouping_skip_keys=["rank", "step"], token_aligner="smart", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) warning_records = [r for r in records if isinstance(r, WarningRecord)] layout_warnings = [ @@ -1950,18 +2039,17 @@ class TestEntrypointAlignment: assert summary.skipped == 0 def test_alignment_fallback_when_no_aux(self, tmp_path, capsys): - """Without aux tensors, logical grouping skips alignment and compares per-step.""" + """Without aux tensors, smart alignment falls back to per-step comparison.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=2) - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="logical", token_aligner="smart", diff_threshold=0.1, ) capsys.readouterr() - run(args) + run(parse_args(argv)) captured = capsys.readouterr() records = _parse_jsonl(captured.out) warning_records = [r for r in records if isinstance(r, WarningRecord)] @@ -1986,12 +2074,12 @@ class TestEntrypointNonTensorValues: """Test non-tensor value comparison through the full entrypoint pipeline.""" def test_non_tensor_float_same_value(self, tmp_path: Path, capsys) -> None: - """Two sides dump the same float → NonTensorRecord with values_equal=True, category=passed.""" + """Two sides dump the same float → NonTensorComparisonRecord with values_equal=True, category=passed.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125 ) - args = _make_args(baseline_path, target_path, grouping="raw") - records, _ = _run_and_parse(args, capsys) + argv = _make_argv(baseline_path, target_path, preset="raw") + records, _ = _run_and_parse(argv, capsys) non_tensors = _get_non_tensors(records) assert len(non_tensors) == 1 @@ -2005,12 +2093,12 @@ class TestEntrypointNonTensorValues: assert summary.failed == 0 def test_non_tensor_float_different_value(self, tmp_path: Path, capsys) -> None: - """Two sides dump different floats → NonTensorRecord with values_equal=False, category=failed.""" + """Two sides dump different floats → NonTensorComparisonRecord with values_equal=False, category=failed.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.25 ) - args = _make_args(baseline_path, target_path, grouping="raw") - records, _ = _run_and_parse(args, capsys) + argv = _make_argv(baseline_path, target_path, preset="raw") + records, _ = _run_and_parse(argv, capsys) non_tensors = _get_non_tensors(records) assert len(non_tensors) == 1 @@ -2029,8 +2117,8 @@ class TestEntrypointNonTensorValues: baseline_value="flash_attn", target_value="flash_attn", ) - args = _make_args(baseline_path, target_path, grouping="raw") - records, _ = _run_and_parse(args, capsys) + argv = _make_argv(baseline_path, target_path, preset="raw") + records, _ = _run_and_parse(argv, capsys) non_tensors = _get_non_tensors(records) assert len(non_tensors) == 1 @@ -2055,12 +2143,12 @@ class TestEntrypointNonTensorValues: extra_tensor_dumps=[("hidden", tensor)], ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) non_tensors = _get_non_tensors(records) @@ -2080,8 +2168,8 @@ class TestEntrypointNonTensorValues: baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="debug_info", baseline_value=value, target_value=value ) - args = _make_args(baseline_path, target_path, grouping="raw") - records, _ = _run_and_parse(args, capsys) + argv = _make_argv(baseline_path, target_path, preset="raw") + records, _ = _run_and_parse(argv, capsys) non_tensors = _get_non_tensors(records) assert len(non_tensors) == 1 @@ -2090,12 +2178,12 @@ class TestEntrypointNonTensorValues: assert non_tensors[0].target_type == "dict" def test_non_tensor_none_value(self, tmp_path: Path, capsys) -> None: - """Dumping None is displayed as NonTensorRecord, not skipped as load failure.""" + """Dumping None is displayed as NonTensorComparisonRecord, not skipped as load failure.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="optional_param", baseline_value=None, target_value=None ) - args = _make_args(baseline_path, target_path, grouping="raw") - records, _ = _run_and_parse(args, capsys) + argv = _make_argv(baseline_path, target_path, preset="raw") + records, _ = _run_and_parse(argv, capsys) non_tensors = _get_non_tensors(records) assert len(non_tensors) == 1 @@ -2106,19 +2194,19 @@ class TestEntrypointNonTensorValues: assert non_tensors[0].category == "passed" def test_non_tensor_json_roundtrip(self, tmp_path: Path, capsys) -> None: - """NonTensorRecord JSON output can be parsed back correctly.""" + """NonTensorComparisonRecord JSON output can be parsed back correctly.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125 ) - args = _make_args(baseline_path, target_path, grouping="raw") - records, _ = _run_and_parse(args, capsys) + argv = _make_argv(baseline_path, target_path, preset="raw") + records, _ = _run_and_parse(argv, capsys) non_tensors = _get_non_tensors(records) assert len(non_tensors) == 1 json_str: str = non_tensors[0].model_dump_json() roundtripped = parse_record_json(json_str) - assert isinstance(roundtripped, NonTensorRecord) + assert isinstance(roundtripped, NonTensorComparisonRecord) assert roundtripped.name == "sm_scale" assert roundtripped.values_equal is True @@ -2137,16 +2225,16 @@ class TestEntrypointVisualize: """--visualize-bundle-details with --filter produces PNG files.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) viz_dir = tmp_path / "viz_out" - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", filter="tensor_a", viz_bundle_details=True, viz_output_dir=str(viz_dir), ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) assert len(_get_comparisons(records)) == 1 png_files = list(viz_dir.glob("*.png")) @@ -2157,30 +2245,32 @@ class TestEntrypointVisualize: """Without --visualize-bundle-details, no PNGs are created.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) viz_dir = tmp_path / "viz_out" - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", viz_bundle_details=False, viz_output_dir=str(viz_dir), ) - _run_and_parse(args, capsys) + _run_and_parse(argv, capsys) assert not viz_dir.exists() or len(list(viz_dir.glob("*.png"))) == 0 # --------------------------- Assertion helpers ------------------- -def _get_comparisons(records: list[AnyRecord]) -> list[ComparisonRecord]: - return [r for r in records if isinstance(r, ComparisonRecord)] +def _get_comparisons(records: list[AnyRecord]) -> list[TensorComparisonRecord]: + return [r for r in records if isinstance(r, TensorComparisonRecord)] -def _get_non_tensors(records: list[AnyRecord]) -> list[NonTensorRecord]: - return [r for r in records if isinstance(r, NonTensorRecord)] +def _get_non_tensors(records: list[AnyRecord]) -> list[NonTensorComparisonRecord]: + return [r for r in records if isinstance(r, NonTensorComparisonRecord)] -def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonRecord: +def _assert_single_comparison_passed( + records: list[AnyRecord], +) -> TensorComparisonRecord: comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].diff is not None @@ -2283,34 +2373,77 @@ def _create_non_tensor_dumps( return baseline_path, target_path -def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace: - defaults = dict( - baseline_path=str(baseline_path), - target_path=str(target_path), - start_step=0, - end_step=1000000, - diff_threshold=1e-3, - filter=None, - output_format="json", - grouping="logical", - token_aligner="concat_steps", - viz_bundle_details=False, - viz_output_dir="/tmp/comparator_viz/", - visualize_per_token=None, - override_dims=[], - override_baseline_dims=[], - override_target_dims=[], - override_config=None, - allow_skip_pattern=".*", - report_path="", - ) - defaults.update(overrides) - return Namespace(**defaults) +def _make_argv( + baseline_path: Path, + target_path: Path, + *, + preset: str | None = None, + grouping_skip_keys: list[str] | None = None, + token_aligner: str | None = None, + diff_threshold: float = 1e-3, + output_format: str = "json", + start_step: int | None = None, + end_step: int | None = None, + filter: str | None = None, + override_dims: list[str] | None = None, + override_baseline_dims: list[str] | None = None, + override_target_dims: list[str] | None = None, + override_config: str | None = None, + allow_skip_pattern: str | None = None, + report_path: str | None = "", + viz_bundle_details: bool = False, + viz_output_dir: str | None = None, + visualize_per_token: str | None = None, +) -> list[str]: + argv: list[str] = [ + "--baseline-path", + str(baseline_path), + "--target-path", + str(target_path), + "--diff-threshold", + str(diff_threshold), + "--output-format", + output_format, + ] + + if preset is not None: + argv += ["--preset", preset] + if grouping_skip_keys is not None: + argv += ["--grouping-skip-keys"] + grouping_skip_keys + if token_aligner is not None: + argv += ["--token-aligner", token_aligner] + if start_step is not None: + argv += ["--start-step", str(start_step)] + if end_step is not None: + argv += ["--end-step", str(end_step)] + if filter is not None: + argv += ["--filter", filter] + for dim in override_dims or []: + argv += ["--override-dims", dim] + for dim in override_baseline_dims or []: + argv += ["--override-baseline-dims", dim] + for dim in override_target_dims or []: + argv += ["--override-target-dims", dim] + if override_config is not None: + argv += ["--override-config", override_config] + if allow_skip_pattern is not None: + argv += ["--allow-skip-pattern", allow_skip_pattern] + if report_path is not None: + argv += ["--report-path", report_path] + if viz_bundle_details: + argv += ["--viz-bundle-details"] + if viz_output_dir is not None: + argv += ["--viz-output-dir", viz_output_dir] + if visualize_per_token is not None: + argv += ["--visualize-per-token", visualize_per_token] + + return argv def _run_and_parse( - args: Namespace, capsys: pytest.CaptureFixture + argv: list[str], capsys: pytest.CaptureFixture ) -> tuple[list[AnyRecord], int]: + args: Namespace = parse_args(argv) capsys.readouterr() exit_code: int = run(args) return _parse_jsonl(capsys.readouterr().out), exit_code @@ -2890,13 +3023,13 @@ class TestEntrypointPerTokenVisualization: target_path: Path = target_dir / _FIXED_EXP_NAME output_png: Path = tmp_path / "per_token.png" - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", visualize_per_token=str(output_png), ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 2 @@ -2911,9 +3044,9 @@ class TestEntrypointPerTokenVisualization: def test_no_visualize_no_per_token(self, tmp_path: Path, capsys) -> None: """Without --visualize-per-token, per_token_rel_diff is None.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) comparisons = _get_comparisons(records) assert len(comparisons) == 1 @@ -3006,17 +3139,17 @@ class TestEntrypointThdCpZigzag: ) # --- Run comparison --- - args: Namespace = _make_args( + argv: list[str] = _make_argv( sglang_dir / _FIXED_EXP_NAME, megatron_dir / _FIXED_EXP_NAME, - grouping="logical", + grouping_skip_keys=["rank", "step"], token_aligner="smart", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) - comparisons: list[ComparisonRecord] = _get_comparisons(records) - hidden_comparisons: list[ComparisonRecord] = [ + comparisons: list[TensorComparisonRecord] = _get_comparisons(records) + hidden_comparisons: list[TensorComparisonRecord] = [ c for c in comparisons if c.name == "hidden_states" ] assert len(hidden_comparisons) >= 1 @@ -3057,18 +3190,18 @@ class TestEntrypointThdCpZigzag: total_per_rank=total_per_rank, ) - args: Namespace = _make_args( + argv: list[str] = _make_argv( baseline_path, target_path, - grouping="logical", + grouping_skip_keys=["rank", "step"], token_aligner="smart", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) # hidden_states should pass comparison (after unshard + reorder) - comparisons: list[ComparisonRecord] = _get_comparisons(records) - hidden_comparisons: list[ComparisonRecord] = [ + comparisons: list[TensorComparisonRecord] = _get_comparisons(records) + hidden_comparisons: list[TensorComparisonRecord] = [ c for c in comparisons if c.name == "hidden_states" ] assert len(hidden_comparisons) >= 1 @@ -3128,15 +3261,14 @@ class TestEntrypointDpFilter: framework="sglang", ) - args: Namespace = _make_args( + argv: list[str] = _make_argv( tmp_path / "baseline" / _FIXED_EXP_NAME, tmp_path / "target" / _FIXED_EXP_NAME, - grouping="logical", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) - comparison: ComparisonRecord = _assert_single_comparison_passed(records) + comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp2_megatron_both_sides(self, tmp_path: Path, capsys) -> None: @@ -3184,15 +3316,14 @@ class TestEntrypointDpFilter: framework="megatron", ) - args: Namespace = _make_args( + argv: list[str] = _make_argv( tmp_path / "baseline" / _FIXED_EXP_NAME, tmp_path / "target" / _FIXED_EXP_NAME, - grouping="logical", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) - comparison: ComparisonRecord = _assert_single_comparison_passed(records) + comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp2_tp2_sglang(self, tmp_path: Path, capsys) -> None: @@ -3233,15 +3364,14 @@ class TestEntrypointDpFilter: ) rank += 1 - args: Namespace = _make_args( + argv: list[str] = _make_argv( tmp_path / "baseline" / _FIXED_EXP_NAME, tmp_path / "target" / _FIXED_EXP_NAME, - grouping="logical", diff_threshold=1e-3, ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) - comparison: ComparisonRecord = _assert_single_comparison_passed(records) + comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp2_both_nonempty_raises(self, tmp_path: Path, capsys) -> None: @@ -3273,17 +3403,158 @@ class TestEntrypointDpFilter: framework="sglang", ) - args: Namespace = _make_args( + argv: list[str] = _make_argv( tmp_path / "baseline" / _FIXED_EXP_NAME, tmp_path / "target" / _FIXED_EXP_NAME, - grouping="logical", diff_threshold=1e-3, ) with pytest.raises( AssertionError, match="Expected exactly 1 non-empty dp_rank" ): - _run_and_parse(args, capsys) + _run_and_parse(argv, capsys) + + +class TestEntrypointDpGroupAlias: + """E2E tests for the ``# dp:=`` dp group alias feature. + + In dp_attn mode, dp_size > 1 but MLP tensors after dp_gather have data + on all ranks. With ``# dp:=moe_dp`` in dims, the dp filter uses + ``moe_dp_rank/moe_dp_size`` instead of ``dp_rank/dp_size``. + """ + + def test_dp_alias_absent_group_noop(self, tmp_path: Path, capsys) -> None: + """Single rank with ``# dp:=moe_dp`` in dims → parse_dims strips ``#``, comparison OK.""" + torch.manual_seed(42) + tensor_data: torch.Tensor = torch.randn(10, 8) + target_data: torch.Tensor = tensor_data + torch.randn(10, 8) * 0.001 + + for side_dir_name, data in [("baseline", tensor_data), ("target", target_data)]: + side_dir: Path = tmp_path / side_dir_name + side_dir.mkdir() + + _create_rank_dump( + side_dir, + rank=0, + name="hidden", + tensor=data, + dims="t h # dp:=moe_dp", + parallel_info={ + "tp_rank": 0, + "tp_size": 1, + "dp_rank": 0, + "dp_size": 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, _ = _run_and_parse(argv, capsys) + + comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + assert comparison.name == "hidden" + + def test_dp_alias_via_override_dims(self, tmp_path: Path, capsys) -> None: + """--override-dims adds ``# dp:=moe_dp`` → dp filter uses alias, filters correctly.""" + torch.manual_seed(42) + tensor_data: torch.Tensor = torch.randn(10, 8) + target_data: torch.Tensor = tensor_data + torch.randn(10, 8) * 0.001 + + for side_dir_name, data in [("baseline", tensor_data), ("target", target_data)]: + side_dir: Path = tmp_path / side_dir_name + side_dir.mkdir() + + # moe_dp_rank=0: non-empty + _create_rank_dump( + side_dir, + rank=0, + name="hidden", + tensor=data, + dims="t h", + parallel_info={ + "tp_rank": 0, + "tp_size": 1, + "dp_rank": 0, + "dp_size": 1, + "moe_dp_rank": 0, + "moe_dp_size": 2, + }, + framework="sglang", + ) + + # moe_dp_rank=1: empty + _create_rank_dump( + side_dir, + rank=1, + name="hidden", + tensor=torch.empty(0, 8), + dims="t h", + parallel_info={ + "tp_rank": 0, + "tp_size": 1, + "dp_rank": 0, + "dp_size": 1, + "moe_dp_rank": 1, + "moe_dp_size": 2, + }, + framework="sglang", + ) + + argv: list[str] = _make_argv( + tmp_path / "baseline" / _FIXED_EXP_NAME, + tmp_path / "target" / _FIXED_EXP_NAME, + diff_threshold=1e-3, + override_dims=["hidden:t h # dp:=moe_dp"], + ) + records, _ = _run_and_parse(argv, capsys) + + comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + assert comparison.name == "hidden" + + def test_dp_alias_with_real_alias_group_filters( + self, tmp_path: Path, capsys + ) -> None: + """Alias group present with moe_dp_size=2, one empty rank → filters correctly.""" + torch.manual_seed(42) + tensor_data: torch.Tensor = torch.randn(10, 8) + target_data: torch.Tensor = tensor_data + torch.randn(10, 8) * 0.001 + + for side_dir_name, data in [("baseline", tensor_data), ("target", target_data)]: + side_dir: Path = tmp_path / side_dir_name + side_dir.mkdir() + + for moe_dp_rank in range(2): + tensor: torch.Tensor = data if moe_dp_rank == 0 else torch.empty(0, 8) + _create_rank_dump( + side_dir, + rank=moe_dp_rank, + name="hidden", + tensor=tensor, + dims="t h # dp:=moe_dp", + parallel_info={ + "tp_rank": 0, + "tp_size": 1, + "dp_rank": 0, + "dp_size": 1, + "moe_dp_rank": moe_dp_rank, + "moe_dp_size": 2, + }, + framework="sglang", + ) + + argv: list[str] = _make_argv( + tmp_path / "baseline" / _FIXED_EXP_NAME, + tmp_path / "target" / _FIXED_EXP_NAME, + diff_threshold=1e-3, + ) + records, _ = _run_and_parse(argv, capsys) + + comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + assert comparison.name == "hidden" class TestEntrypointDpGroupAlias: @@ -3466,7 +3737,7 @@ class TestEntrypointMetaOverride: records: list[AnyRecord], *, expected_count: int = 1 ) -> None: """Assert that exactly expected_count comparisons exist and all passed.""" - comparisons: list[ComparisonRecord] = _get_comparisons(records) + comparisons: list[TensorComparisonRecord] = _get_comparisons(records) assert len(comparisons) == expected_count assert all(c.diff is not None and c.diff.passed for c in comparisons) @@ -3504,13 +3775,12 @@ class TestEntrypointMetaOverride: parallel_info={"tp_rank": tp_rank, "tp_size": 2}, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="logical", override_dims=["hidden:t h(tp)"], ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) @pytest.mark.parametrize( "baseline_dims, target_dims, override_kwarg", @@ -3536,8 +3806,8 @@ class TestEntrypointMetaOverride: target_dims=target_dims, ) - args = _make_args(baseline_path, target_path, grouping="raw", **override_kwarg) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + argv = _make_argv(baseline_path, target_path, preset="raw", **override_kwarg) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_override_config_yaml(self, tmp_path: Path, capsys) -> None: """--override-config YAML overrides dims.""" @@ -3550,13 +3820,13 @@ class TestEntrypointMetaOverride: dims: "t h" """)) - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", override_config=str(yaml_path), ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_no_match_uses_original_dims(self, tmp_path: Path, capsys) -> None: """When override regex doesn't match, original dims from dump are used.""" @@ -3566,13 +3836,13 @@ class TestEntrypointMetaOverride: target_dims="t h", ) - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", override_dims=["no_match_pattern:b s d"], ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_selective_match_multi_tensor(self, tmp_path: Path, capsys) -> None: """Override matches only 'logits'; 'hidden' uses original dims.""" @@ -3597,13 +3867,13 @@ class TestEntrypointMetaOverride: ) _create_rank_dump(target_dir, rank=0, name=name, tensor=t_tensor, dims=dims) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", override_dims=["logits:t v"], ) - self._assert_all_passed(_run_and_parse(args, capsys)[0], expected_count=2) + self._assert_all_passed(_run_and_parse(argv, capsys)[0], expected_count=2) def test_multiple_cli_override_dims(self, tmp_path: Path, capsys) -> None: """Multiple --override-dims for different tensors.""" @@ -3630,13 +3900,13 @@ class TestEntrypointMetaOverride: target_dir, rank=0, name=name, tensor=t_tensor, dims="x y" ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", override_dims=["hidden:t h", "logits:t v"], ) - self._assert_all_passed(_run_and_parse(args, capsys)[0], expected_count=2) + self._assert_all_passed(_run_and_parse(argv, capsys)[0], expected_count=2) def test_per_side_dims_different_parallelism(self, tmp_path: Path, capsys) -> None: """baseline TP-sharded, target EP-sharded — per-side override fixes both.""" @@ -3671,14 +3941,13 @@ class TestEntrypointMetaOverride: parallel_info={"ep_rank": ep_rank, "ep_size": 2}, ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="logical", override_baseline_dims=["hidden:t h(tp)"], override_target_dims=["hidden:t h(ep)"], ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_yaml_first_match_wins_e2e(self, tmp_path: Path, capsys) -> None: """YAML with two matching rules: first rule wins in real pipeline.""" @@ -3693,13 +3962,13 @@ class TestEntrypointMetaOverride: dims: "a b" """)) - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", override_config=str(yaml_path), ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_cli_overrides_yaml_e2e(self, tmp_path: Path, capsys) -> None: """CLI --override-dims wins over YAML rule for the same tensor.""" @@ -3712,14 +3981,14 @@ class TestEntrypointMetaOverride: dims: "a b" """)) - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", override_dims=["hidden:t h"], override_config=str(yaml_path), ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_override_injects_dims_when_absent(self, tmp_path: Path, capsys) -> None: """Override injects dims into meta even when dump had no dims annotation.""" @@ -3729,13 +3998,13 @@ class TestEntrypointMetaOverride: target_dims=None, ) - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", override_dims=["hidden:t h"], ) - self._assert_all_passed(_run_and_parse(args, capsys)[0]) + self._assert_all_passed(_run_and_parse(argv, capsys)[0]) def test_non_tensor_unaffected_by_override(self, tmp_path: Path, capsys) -> None: """Non-tensor values pass through without error even with active override.""" @@ -3756,22 +4025,22 @@ class TestEntrypointMetaOverride: extra_tensor_dumps=[("hidden", tensor)], ) - args = _make_args( + argv = _make_argv( baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME, - grouping="raw", + preset="raw", override_dims=["hidden:x y"], ) - records, _ = _run_and_parse(args, capsys) + records, _ = _run_and_parse(argv, capsys) - non_tensors: list[NonTensorRecord] = [ - r for r in records if isinstance(r, NonTensorRecord) + non_tensors: list[NonTensorComparisonRecord] = [ + r for r in records if isinstance(r, NonTensorComparisonRecord) ] assert len(non_tensors) == 1 assert non_tensors[0].name == "sm_scale" assert non_tensors[0].values_equal - comparisons: list[ComparisonRecord] = _get_comparisons(records) + comparisons: list[TensorComparisonRecord] = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].name == "hidden" @@ -3866,9 +4135,9 @@ class TestExitCode: def test_e2e_all_passed_exit_zero(self, tmp_path, capsys): """Integration: all comparisons pass → run() returns 0.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) - args = _make_args(baseline_path, target_path, grouping="raw") + argv = _make_argv(baseline_path, target_path, preset="raw") - records, exit_code = _run_and_parse(args, capsys) + records, exit_code = _run_and_parse(argv, capsys) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.passed == 2 @@ -3887,11 +4156,9 @@ class TestExitCode: name="tensor_a", tensor=torch.randn(10, 10) * 100, ) - args = _make_args( - baseline_path, target_path, grouping="raw", diff_threshold=1e-3 - ) + argv = _make_argv(baseline_path, target_path, preset="raw", diff_threshold=1e-3) - records, exit_code = _run_and_parse(args, capsys) + records, exit_code = _run_and_parse(argv, capsys) summary = records[-1] assert isinstance(summary, SummaryRecord) assert summary.failed == 1 @@ -3906,7 +4173,7 @@ class TestExitCodeSubprocess: baseline_path: Path, target_path: Path, *, - grouping: str = "raw", + preset: str = "raw", allow_skip_pattern: str = ".*", ) -> subprocess.CompletedProcess[str]: cmd: list[str] = [ @@ -3917,8 +4184,8 @@ class TestExitCodeSubprocess: str(baseline_path), "--target-path", str(target_path), - "--grouping", - grouping, + "--preset", + preset, "--output-format", "json", "--allow-skip-pattern", @@ -3975,9 +4242,9 @@ class TestReportOutput: def test_default_report_path(self, tmp_path, capsys): """Default writes to /comparator_report.jsonl with ConfigRecord + SummaryRecord.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) - args = _make_args(baseline_path, target_path, grouping="raw", report_path=None) + argv = _make_argv(baseline_path, target_path, preset="raw", report_path=None) - exit_code: int = run(args) + exit_code: int = run(parse_args(argv)) report_file: Path = target_path / "comparator_report.jsonl" assert report_file.exists() @@ -3991,14 +4258,14 @@ class TestReportOutput: """--report-path writes to the specified location.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) custom_path: Path = tmp_path / "custom" / "report.jsonl" - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", report_path=str(custom_path), ) - run(args) + run(parse_args(argv)) assert custom_path.exists() report_records: list[AnyRecord] = _parse_jsonl(custom_path.read_text()) @@ -4008,9 +4275,9 @@ class TestReportOutput: def test_disabled_report(self, tmp_path, capsys): """--report-path '' disables file generation.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) - args = _make_args(baseline_path, target_path, grouping="raw", report_path="") + argv = _make_argv(baseline_path, target_path, preset="raw", report_path="") - run(args) + run(parse_args(argv)) report_file: Path = target_path / "comparator_report.jsonl" assert not report_file.exists() @@ -4019,16 +4286,16 @@ class TestReportOutput: """In json mode, report content matches stdout output.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) report_file: Path = tmp_path / "report.jsonl" - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", output_format="json", report_path=str(report_file), ) capsys.readouterr() - run(args) + run(parse_args(argv)) stdout_lines: list[str] = capsys.readouterr().out.strip().splitlines() report_lines: list[str] = report_file.read_text().strip().splitlines() @@ -4038,15 +4305,15 @@ class TestReportOutput: """Text stdout mode still writes JSONL report.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) report_file: Path = tmp_path / "report.jsonl" - args = _make_args( + argv = _make_argv( baseline_path, target_path, - grouping="raw", + preset="raw", output_format="text", report_path=str(report_file), ) - run(args) + run(parse_args(argv)) assert report_file.exists() report_records: list[AnyRecord] = _parse_jsonl(report_file.read_text()) diff --git a/test/registered/debug_utils/comparator/test_manually_verify.py b/test/registered/debug_utils/comparator/test_manually_verify.py index 51a9ef0f8..fd23617bc 100644 --- a/test/registered/debug_utils/comparator/test_manually_verify.py +++ b/test/registered/debug_utils/comparator/test_manually_verify.py @@ -207,7 +207,9 @@ class TestPerTokenHeatmapManualVerify: the left (small diff), bright/hot on the right (large diff). Multiple rows for different tensor names. Colorbar shows log10 scale. """ - from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord + from sglang.srt.debug_utils.comparator.output_types import ( + TensorComparisonRecord, + ) from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, ) @@ -220,7 +222,7 @@ class TestPerTokenHeatmapManualVerify: hidden_dim: int = 128 num_tensors: int = 5 - records: list[ComparisonRecord] = [] + records: list[TensorComparisonRecord] = [] for i in range(num_tensors): baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) noise_scale: torch.Tensor = torch.linspace( @@ -235,7 +237,7 @@ class TestPerTokenHeatmapManualVerify: diff_threshold=1e-3, seq_dim=0, ) - records.append(ComparisonRecord(**info.model_dump())) + records.append(TensorComparisonRecord(**info.model_dump())) output_path: Path = tmp_path / "per_token_increasing_diff.png" result = generate_per_token_heatmap(records=records, output_path=output_path) @@ -250,7 +252,9 @@ class TestPerTokenHeatmapManualVerify: Expected: Heatmap shows one bright vertical stripe at the spike position, rest is dark/cold. """ - from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord + from sglang.srt.debug_utils.comparator.output_types import ( + TensorComparisonRecord, + ) from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, ) @@ -264,7 +268,7 @@ class TestPerTokenHeatmapManualVerify: spike_pos: int = 32 num_tensors: int = 4 - records: list[ComparisonRecord] = [] + records: list[TensorComparisonRecord] = [] for i in range(num_tensors): baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) target: torch.Tensor = baseline.clone() @@ -277,7 +281,7 @@ class TestPerTokenHeatmapManualVerify: diff_threshold=1e-3, seq_dim=0, ) - records.append(ComparisonRecord(**info.model_dump())) + records.append(TensorComparisonRecord(**info.model_dump())) output_path: Path = tmp_path / "per_token_single_spike.png" result = generate_per_token_heatmap(records=records, output_path=output_path) diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index 0648b1392..1996fb5a1 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -22,11 +22,11 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( ) from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout from sglang.srt.debug_utils.comparator.output_types import ( - ComparisonRecord, GeneralWarning, - NonTensorRecord, - SkipRecord, + NonTensorComparisonRecord, + SkipComparisonRecord, SummaryRecord, + TensorComparisonRecord, parse_record_json, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( @@ -208,9 +208,9 @@ def _make_comparison_record( *, diff: DiffInfo | None, warnings: list | None = None, -) -> ComparisonRecord: +) -> TensorComparisonRecord: ti: TensorInfo = _make_tensor_info() - return ComparisonRecord( + return TensorComparisonRecord( name="t", baseline=ti, target=ti, @@ -223,7 +223,7 @@ def _make_comparison_record( class TestOutputRecordCategories: def test_skip_record_with_warnings_is_failed(self) -> None: - record = SkipRecord( + record = SkipComparisonRecord( name="t", reason="test", warnings=[GeneralWarning(category="c", message="m")], @@ -231,28 +231,28 @@ class TestOutputRecordCategories: assert record.category == "failed" def test_skip_record_no_warnings_is_skipped(self) -> None: - record = SkipRecord(name="t", reason="test") + record = SkipComparisonRecord(name="t", reason="test") assert record.category == "skipped" def test_comparison_record_diff_none_is_failed(self) -> None: - record: ComparisonRecord = _make_comparison_record(diff=None) + record: TensorComparisonRecord = _make_comparison_record(diff=None) assert record.category == "failed" def test_comparison_record_passed_with_warnings_is_failed(self) -> None: - record: ComparisonRecord = _make_comparison_record( + record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), warnings=[GeneralWarning(category="c", message="m")], ) assert record.category == "failed" def test_comparison_record_passed_no_warnings_is_passed(self) -> None: - record: ComparisonRecord = _make_comparison_record( + record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) assert record.category == "passed" def test_non_tensor_record_equal_is_passed(self) -> None: - record = NonTensorRecord( + record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", @@ -263,7 +263,7 @@ class TestOutputRecordCategories: assert record.category == "passed" def test_non_tensor_record_different_is_failed(self) -> None: - record = NonTensorRecord( + record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", @@ -274,7 +274,7 @@ class TestOutputRecordCategories: assert record.category == "failed" def test_non_tensor_record_with_warnings_is_failed(self) -> None: - record = NonTensorRecord( + record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", @@ -286,7 +286,7 @@ class TestOutputRecordCategories: assert record.category == "failed" def test_non_tensor_record_json_roundtrip(self) -> None: - record = NonTensorRecord( + record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", @@ -296,14 +296,14 @@ class TestOutputRecordCategories: ) json_str: str = record.model_dump_json() roundtripped = parse_record_json(json_str) - assert isinstance(roundtripped, NonTensorRecord) + assert isinstance(roundtripped, NonTensorComparisonRecord) assert roundtripped.name == "sm_scale" assert roundtripped.values_equal is False assert roundtripped.baseline_value == "0.125" assert roundtripped.target_value == "0.25" def test_non_tensor_record_text_format_equal(self) -> None: - record = NonTensorRecord( + record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", @@ -316,7 +316,7 @@ class TestOutputRecordCategories: assert "[equal]" in text def test_non_tensor_record_text_format_different(self) -> None: - record = NonTensorRecord( + record = NonTensorComparisonRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", @@ -351,10 +351,10 @@ def _make_aligner_plan() -> AlignerPlan: ) -class TestAlignerPlanInComparisonRecord: +class TestAlignerPlanInTensorComparisonRecord: def test_comparison_record_with_aligner_plan(self) -> None: plan: AlignerPlan = _make_aligner_plan() - record: ComparisonRecord = _make_comparison_record( + record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) record_with_plan = record.model_copy(update={"aligner_plan": plan}) @@ -363,7 +363,7 @@ class TestAlignerPlanInComparisonRecord: def test_aligner_plan_json_roundtrip(self) -> None: plan: AlignerPlan = _make_aligner_plan() - record: ComparisonRecord = _make_comparison_record( + record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) record_with_plan = record.model_copy(update={"aligner_plan": plan}) @@ -376,7 +376,7 @@ class TestAlignerPlanInComparisonRecord: == "unsharder" ) - roundtripped: ComparisonRecord = parse_record_json(json_str) + roundtripped: TensorComparisonRecord = parse_record_json(json_str) assert roundtripped.aligner_plan is not None assert ( roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type @@ -384,16 +384,16 @@ class TestAlignerPlanInComparisonRecord: ) def test_comparison_record_without_aligner_plan(self) -> None: - record: ComparisonRecord = _make_comparison_record( + record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) json_str: str = record.model_dump_json() - roundtripped: ComparisonRecord = parse_record_json(json_str) + roundtripped: TensorComparisonRecord = parse_record_json(json_str) assert roundtripped.aligner_plan is None def test_aligner_plan_text_format(self) -> None: plan: AlignerPlan = _make_aligner_plan() - record: ComparisonRecord = _make_comparison_record( + record: TensorComparisonRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) record_with_plan = record.model_copy(update={"aligner_plan": plan}) diff --git a/test/registered/debug_utils/comparator/test_per_token_visualizer.py b/test/registered/debug_utils/comparator/test_per_token_visualizer.py index 3f0e0e0ad..b99f95aba 100644 --- a/test/registered/debug_utils/comparator/test_per_token_visualizer.py +++ b/test/registered/debug_utils/comparator/test_per_token_visualizer.py @@ -9,7 +9,7 @@ from pathlib import Path import pytest import torch -from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord +from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( compare_tensor_pair, ) @@ -31,8 +31,8 @@ def _make_comparison_record( baseline: torch.Tensor, target: torch.Tensor, seq_dim: int = 0, -) -> ComparisonRecord: - """Build a ComparisonRecord with per-token data from raw tensors.""" +) -> TensorComparisonRecord: + """Build a TensorComparisonRecord with per-token data from raw tensors.""" info = compare_tensor_pair( x_baseline=baseline, x_target=target, @@ -40,7 +40,7 @@ def _make_comparison_record( diff_threshold=1e-3, seq_dim=seq_dim, ) - return ComparisonRecord(**info.model_dump()) + return TensorComparisonRecord(**info.model_dump()) class TestPerTokenVisualizer: @@ -68,7 +68,7 @@ class TestPerTokenVisualizer: name="no_per_token", diff_threshold=1e-3, ) - record = ComparisonRecord(**info.model_dump()) + record = TensorComparisonRecord(**info.model_dump()) output_path: Path = tmp_path / "no_data.png" result = generate_per_token_heatmap(records=[record], output_path=output_path) @@ -82,7 +82,7 @@ class TestPerTokenVisualizer: ) torch.manual_seed(42) - records: list[ComparisonRecord] = [ + records: list[TensorComparisonRecord] = [ _make_comparison_record( name=f"tensor_{i}", baseline=torch.randn(16, 32), @@ -108,7 +108,7 @@ class TestPerTokenVisualizer: ) torch.manual_seed(42) - records: list[ComparisonRecord] = [ + records: list[TensorComparisonRecord] = [ _make_comparison_record( name="short", baseline=torch.randn(4, 8), diff --git a/test/registered/debug_utils/comparator/test_preset.py b/test/registered/debug_utils/comparator/test_preset.py new file mode 100644 index 000000000..d7f746790 --- /dev/null +++ b/test/registered/debug_utils/comparator/test_preset.py @@ -0,0 +1,44 @@ +import pytest + +from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="default", nightly=True) + + +class TestExpandPreset: + """Test preset expansion logic.""" + + def test_explicit_preset(self): + """--preset sglang_megatron expands into its argv.""" + argv = [ + "--baseline-path", + "/a", + "--preset", + "sglang_megatron", + "--diff-threshold", + "0.01", + ] + result = expand_preset(argv, presets=PRESETS) + assert "--preset" not in result + assert "--grouping-skip-keys" in result + assert "concat_steps" in result + assert "--baseline-path" in result + assert "--diff-threshold" in result + + def test_default_preset_applied(self): + """No --preset and no --grouping-skip-keys triggers default preset.""" + argv = ["--baseline-path", "/a"] + result = expand_preset(argv, presets=PRESETS) + assert "--grouping-skip-keys" in result + + def test_explicit_skip_keys_prevents_default(self): + """Explicit --grouping-skip-keys prevents default preset injection.""" + argv = ["--grouping-skip-keys", "rank", "--baseline-path", "/a"] + result = expand_preset(argv, presets=PRESETS) + assert result == argv + + def test_unknown_preset_raises(self): + """Unknown preset name raises ValueError.""" + with pytest.raises(ValueError, match="Unknown value for --preset"): + expand_preset(["--preset", "nonexistent"], presets=PRESETS)