From 5bf3deb4bc3b54b112e063484340efabe124d15e Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 2 Mar 2026 18:46:27 +0800 Subject: [PATCH] Trace execution information in dump comparator (#19682) --- .../srt/debug_utils/comparator/__init__.py | 3 + .../comparator/aligner/entrypoint/executor.py | 104 +++- .../aligner/entrypoint/traced_types.py | 37 ++ .../aligner/token_aligner/smart/aux_loader.py | 8 +- .../comparator/bundle_comparator.py | 42 +- .../srt/debug_utils/comparator/display.py | 10 +- .../srt/debug_utils/comparator/entrypoint.py | 26 +- .../srt/debug_utils/comparator/log_sink.py | 2 +- .../debug_utils/comparator/output_types.py | 158 +++-- .../srt/debug_utils/comparator/report_sink.py | 87 +++ .../aligner/entrypoint/test_executor.py | 55 +- .../debug_utils/comparator/conftest.py | 27 +- .../debug_utils/comparator/test_display.py | 10 +- .../debug_utils/comparator/test_entrypoint.py | 106 +++- .../debug_utils/comparator/test_log_sink.py | 2 +- .../comparator/test_output_types.py | 567 ++++++++++++++++++ .../debug_utils/comparator/testing_helpers.py | 84 +++ 17 files changed, 1180 insertions(+), 148 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/aligner/entrypoint/traced_types.py create mode 100644 python/sglang/srt/debug_utils/comparator/report_sink.py create mode 100644 test/registered/debug_utils/comparator/testing_helpers.py diff --git a/python/sglang/srt/debug_utils/comparator/__init__.py b/python/sglang/srt/debug_utils/comparator/__init__.py index 2d757adce..ab7ed6c6b 100644 --- a/python/sglang/srt/debug_utils/comparator/__init__.py +++ b/python/sglang/srt/debug_utils/comparator/__init__.py @@ -1,3 +1,6 @@ +from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( # noqa: F401 + TracedAlignerPlan, +) from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( # noqa: F401 AlignerPlan, ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py index 0bbf21be8..bf30dde8b 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py @@ -1,13 +1,19 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Optional +from typing import NamedTuple, Optional import torch from sglang.srt.debug_utils.comparator.aligner.axis_aligner import ( execute_axis_aligner_plan, ) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( + TracedAlignerPlan, + TracedSidePlan, + TracedStepPlan, + TracedSubPlan, +) from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( AlignerPerStepPlan, AlignerPerStepSubPlan, @@ -28,15 +34,31 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import ( execute_unsharder_plan, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan -from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult +from sglang.srt.debug_utils.comparator.output_types import ( + ReplicatedCheckResult, + ShapeSnapshot, +) from sglang.srt.debug_utils.comparator.utils import Pair +class StepPlansResult(NamedTuple): + tensors: dict[int, torch.Tensor] + checks: list[ReplicatedCheckResult] + traced_side: TracedSidePlan + + +class SubPlansResult(NamedTuple): + tensor: Optional[torch.Tensor] + checks: list[ReplicatedCheckResult] + snapshots: list[ShapeSnapshot] + + @dataclass(frozen=True) class AlignerResult: tensors: Optional[Pair[torch.Tensor]] failed_side_xy: Optional[str] # "x" or "y"; None if success replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list) + traced_plan: Optional[TracedAlignerPlan] = None def execute_aligner_plan( @@ -48,26 +70,34 @@ def execute_aligner_plan( all_checks: list[ReplicatedCheckResult] = [] # Per-side: unshard + reorder -> dict[step, tensor] - step_tensors_x, checks_x = _execute_step_plans( + result_x: StepPlansResult = _execute_step_plans( tensors=tensors_pair.x, step_plans=plan.per_step_plans.x ) - all_checks.extend(checks_x) + all_checks.extend(result_x.checks) - step_tensors_y, checks_y = _execute_step_plans( + result_y: StepPlansResult = _execute_step_plans( tensors=tensors_pair.y, step_plans=plan.per_step_plans.y ) - all_checks.extend(checks_y) + all_checks.extend(result_y.checks) - if not step_tensors_x or not step_tensors_y: - failed_side_xy: str = "x" if not step_tensors_x else "y" + traced_plan: TracedAlignerPlan = TracedAlignerPlan( + plan=plan, + per_side=Pair(x=result_x.traced_side, y=result_y.traced_side), + ) + + if not result_x.tensors or not result_y.tensors: + failed_side_xy: str = "x" if not result_x.tensors else "y" return AlignerResult( tensors=None, failed_side_xy=failed_side_xy, replicated_checks=all_checks, + traced_plan=traced_plan, ) # Cross-side: token alignment (or direct extraction for single-step) - step_pair: Pair[dict[int, torch.Tensor]] = Pair(x=step_tensors_x, y=step_tensors_y) + step_pair: Pair[dict[int, torch.Tensor]] = Pair( + x=result_x.tensors, y=result_y.tensors + ) combined: Pair[torch.Tensor] if plan.token_aligner_mode == "concat_steps": combined = execute_token_aligner_concat_steps(tensor_of_step_pair=step_pair) @@ -78,10 +108,10 @@ def execute_aligner_plan( tensor_of_step_pair=step_pair, ) else: - assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1 + assert len(result_x.tensors) == 1 and len(result_y.tensors) == 1 combined = Pair( - x=list(step_tensors_x.values())[0], - y=list(step_tensors_y.values())[0], + x=list(result_x.tensors.values())[0], + y=list(result_y.tensors.values())[0], ) # Cross-side: axis alignment (squeeze singletons + rearrange dim order) @@ -95,50 +125,78 @@ def execute_aligner_plan( tensors=combined, failed_side_xy=None, replicated_checks=all_checks, + traced_plan=traced_plan, ) def _execute_step_plans( tensors: list[torch.Tensor], step_plans: list[AlignerPerStepPlan], -) -> tuple[dict[int, torch.Tensor], list[ReplicatedCheckResult]]: +) -> StepPlansResult: result: dict[int, torch.Tensor] = {} all_checks: list[ReplicatedCheckResult] = [] + traced_steps: list[TracedStepPlan] = [] for step_plan in step_plans: step_tensors: list[torch.Tensor] = [ tensors[i] for i in step_plan.input_object_indices ] - tensor, checks = execute_sub_plans( + sub_result: SubPlansResult = execute_sub_plans( tensors=step_tensors, plans=step_plan.sub_plans ) - all_checks.extend(checks) - if tensor is not None: - result[step_plan.step] = tensor + all_checks.extend(sub_result.checks) - return result, all_checks + traced_subs: list[TracedSubPlan] = [ + TracedSubPlan(plan=sub_plan, snapshot=snapshot) + for sub_plan, snapshot in zip(step_plan.sub_plans, sub_result.snapshots) + ] + traced_steps.append( + TracedStepPlan( + step=step_plan.step, + input_object_indices=step_plan.input_object_indices, + sub_plans=traced_subs, + ) + ) + + if sub_result.tensor is not None: + result[step_plan.step] = sub_result.tensor + + return StepPlansResult( + tensors=result, + checks=all_checks, + traced_side=TracedSidePlan(step_plans=traced_steps), + ) def execute_sub_plans( tensors: list[torch.Tensor], plans: list[AlignerPerStepSubPlan], -) -> tuple[Optional[torch.Tensor], list[ReplicatedCheckResult]]: +) -> SubPlansResult: if not tensors: - return None, [] + return SubPlansResult(tensor=None, checks=[], snapshots=[]) if not plans: if len(tensors) != 1: - return None, [] - return tensors[0], [] + return SubPlansResult(tensor=None, checks=[], snapshots=[]) + return SubPlansResult(tensor=tensors[0], checks=[], snapshots=[]) current: list[torch.Tensor] = tensors all_checks: list[ReplicatedCheckResult] = [] + all_snapshots: list[ShapeSnapshot] = [] for plan in plans: + input_shapes: list[list[int]] = [list(t.shape) for t in current] current, checks = execute_sub_plan(tensors=current, plan=plan) + output_shapes: list[list[int]] = [list(t.shape) for t in current] all_checks.extend(checks) + all_snapshots.append( + ShapeSnapshot( + input_shapes=input_shapes, + output_shapes=output_shapes, + ) + ) assert len(current) == 1 - return current[0], all_checks + return SubPlansResult(tensor=current[0], checks=all_checks, snapshots=all_snapshots) def execute_sub_plan( diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/traced_types.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/traced_types.py new file mode 100644 index 000000000..1ecdad4c2 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/traced_types.py @@ -0,0 +1,37 @@ +"""Traced wrapper types that embed execution traces (ShapeSnapshots) into plan nodes. + +These types are created *after* execution, pairing each sub-plan with its +observed shape snapshot so that downstream formatters never need to manually +zip plan + trace by index. +""" + +from __future__ import annotations + +from typing import Optional + +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepSubPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.output_types import ShapeSnapshot +from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase + + +class TracedSubPlan(_StrictBase): + plan: AlignerPerStepSubPlan + snapshot: Optional[ShapeSnapshot] = None + + +class TracedStepPlan(_StrictBase): + step: int + input_object_indices: list[int] + sub_plans: list[TracedSubPlan] + + +class TracedSidePlan(_StrictBase): + step_plans: list[TracedStepPlan] + + +class TracedAlignerPlan(_StrictBase): + plan: AlignerPlan + per_side: Pair[TracedSidePlan] diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py index a63c26978..7ce7814e3 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py @@ -240,9 +240,11 @@ def _load_and_align_aux_tensor( dim_names: list[str] = resolve_dim_names(dims_str) tensors = [apply_dim_names(t, dim_names) for t in tensors] - result, _replicated_checks = execute_sub_plans(tensors=tensors, plans=sub_plans) - assert result is not None - return result.rename(None) # strip named dims before returning to plugin + sub_result = execute_sub_plans(tensors=tensors, plans=sub_plans) + assert sub_result.tensor is not None + return sub_result.tensor.rename( + None + ) # strip named dims before returning to plugin log_sink.add( InfoLog( diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index f82f4c5c0..35a453cae 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -29,6 +29,8 @@ from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_ra from sglang.srt.debug_utils.comparator.log_sink import log_sink from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider from sglang.srt.debug_utils.comparator.output_types import ( + BundleFileInfo, + BundleSideInfo, ErrorLog, NonTensorComparisonRecord, SkipComparisonRecord, @@ -44,6 +46,37 @@ from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta _FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"} +def _collect_bundle_side_info( + items: list[ValueWithMeta], + metas: list[dict[str, Any]], +) -> BundleSideInfo: + from sglang.srt.debug_utils.comparator.display import ( + PARALLEL_INFO_KEYS, + extract_parallel_info, + ) + + files: list[BundleFileInfo] = [] + for item, meta in zip(items, metas): + assert isinstance(item.value, torch.Tensor) + tensor: torch.Tensor = item.value + + parallel_info: dict[str, str] = {} + for key in PARALLEL_INFO_KEYS: + extract_parallel_info(row_data=parallel_info, info=meta.get(key, {})) + + files.append( + BundleFileInfo( + shape=list(tensor.shape), + dtype=str(tensor.dtype), + rank=meta.get("rank"), + parallel_info=parallel_info if parallel_info else None, + ) + ) + + dims: Optional[str] = metas[0].get("dims") if metas else None + return BundleSideInfo(num_files=len(files), files=files, dims=dims) + + def compare_bundle_pair( *, name: str, @@ -186,6 +219,12 @@ def _compare_bundle_pair_tensor_type( thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, ) + # Collect raw bundle info before alignment + raw_bundle_info: Pair[BundleSideInfo] = Pair( + x=_collect_bundle_side_info(items=valid_pair.x, metas=metas_pair.x), + y=_collect_bundle_side_info(items=valid_pair.y, metas=metas_pair.y), + ) + # Apply dim names to tensors, then execute tensors_pair: Pair[list[torch.Tensor]] = Pair( x=_apply_dim_names_from_meta( @@ -226,8 +265,9 @@ def _compare_bundle_pair_tensor_type( ) record = TensorComparisonRecord( **info.model_dump(), - aligner_plan=plan, + traced_plan=aligner_result.traced_plan, replicated_checks=replicated_checks, + raw_bundle_info=raw_bundle_info, ) if viz_output_dir is not None: diff --git a/python/sglang/srt/debug_utils/comparator/display.py b/python/sglang/srt/debug_utils/comparator/display.py index 0ef7e5ab6..a20141436 100644 --- a/python/sglang/srt/debug_utils/comparator/display.py +++ b/python/sglang/srt/debug_utils/comparator/display.py @@ -10,11 +10,11 @@ import polars as pl from sglang.srt.debug_utils.comparator.output_types import ( InputIdsRecord, RankInfoRecord, - report_sink, ) +from sglang.srt.debug_utils.comparator.report_sink import report_sink from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta -_PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"] +PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"] def emit_display_records( @@ -68,8 +68,8 @@ def _collect_rank_info( meta: dict[str, Any] = ValueWithMeta.load(dump_dir / row["filename"]).meta row_data: dict[str, Any] = {"rank": row["rank"]} - for key in _PARALLEL_INFO_KEYS: - _extract_parallel_info(row_data=row_data, info=meta.get(key, {})) + for key in PARALLEL_INFO_KEYS: + extract_parallel_info(row_data=row_data, info=meta.get(key, {})) table_rows.append(row_data) return table_rows or None @@ -119,7 +119,7 @@ def _collect_input_ids_and_positions( return table_rows or None -def _extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None: +def extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None: if not info or info.get("error"): return diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 80c39f19d..dfb311cf8 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -31,12 +31,12 @@ from sglang.srt.debug_utils.comparator.output_types import ( 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.report_sink import report_sink from sglang.srt.debug_utils.comparator.utils import ( Pair, auto_descend_dir, @@ -53,7 +53,11 @@ def main() -> None: def run(args: argparse.Namespace) -> int: - report_sink.configure(output_format=args.output_format, report_path=None) + report_sink.configure( + output_format=args.output_format, + report_path=None, + verbosity=args.verbosity, + ) dir_pair: Pair[Path] = Pair( x=auto_descend_dir(Path(args.baseline_path), label="baseline_path"), @@ -69,6 +73,16 @@ def run(args: argparse.Namespace) -> int: Path(args.override_config) if args.override_config else None ) + report_path: Optional[Path] = _resolve_report_path( + target_path=dir_pair.y, + report_path_arg=args.report_path, + ) + report_sink.configure( + output_format=args.output_format, + report_path=report_path, + verbosity=args.verbosity, + ) + report_path: Optional[Path] = _resolve_report_path( target_path=dir_pair.y, report_path_arg=args.report_path, @@ -294,6 +308,14 @@ def parse_args(argv: list[str]) -> argparse.Namespace: default="text", help="Output format: text (default) or json (JSONL, one JSON object per line)", ) + parser.add_argument( + "--verbosity", + type=str, + choices=["minimal", "normal", "verbose"], + default="normal", + help="Output verbosity: minimal (1 line per tensor), normal (compact lifecycle), " + "verbose (full detail). Default: normal", + ) parser.add_argument( "--preset", type=str, diff --git a/python/sglang/srt/debug_utils/comparator/log_sink.py b/python/sglang/srt/debug_utils/comparator/log_sink.py index fb22c4605..8515fa847 100644 --- a/python/sglang/srt/debug_utils/comparator/log_sink.py +++ b/python/sglang/srt/debug_utils/comparator/log_sink.py @@ -27,8 +27,8 @@ class LogSink: from sglang.srt.debug_utils.comparator.output_types import ( LogRecord, _split_logs, - report_sink, ) + from sglang.srt.debug_utils.comparator.report_sink import report_sink errors, infos = _split_logs([log]) report_sink.add(LogRecord(errors=errors, infos=infos)) diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index 65e700d42..dea5a9acc 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -1,12 +1,12 @@ from __future__ import annotations -import sys from abc import abstractmethod -from pathlib import Path -from typing import IO, TYPE_CHECKING, Annotated, Any, Literal, Optional, Union +from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union import polars as pl from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator +from rich.console import RenderableType +from rich.markup import escape from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( format_comparison, @@ -16,12 +16,15 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( DiffInfo, TensorComparisonInfo, ) -from sglang.srt.debug_utils.comparator.utils import _StrictBase +from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase if TYPE_CHECKING: - from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( - AlignerPlan, + from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( + TracedAlignerPlan, + TracedSubPlan, ) + from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan + from sglang.srt.debug_utils.comparator.report_sink import Verbosity class BaseLog(_StrictBase): @@ -59,6 +62,26 @@ class ReplicatedCheckResult(_StrictBase): diff: Optional[DiffInfo] = None +class BundleFileInfo(_StrictBase): + """Per-file info within a bundle (one rank's raw tensor).""" + + shape: list[int] + dtype: str + rank: Optional[int] = None + parallel_info: Optional[dict[str, str]] = None # e.g. {"tp": "0/4", "ep": "1/2"} + + +class BundleSideInfo(_StrictBase): + num_files: int + files: list[BundleFileInfo] + dims: Optional[str] = None # e.g. "b s h(tp) d" + + +class ShapeSnapshot(_StrictBase): + input_shapes: list[list[int]] + output_shapes: list[list[int]] + + class _OutputRecord(_StrictBase): errors: list[ErrorLog] = Field(default_factory=list) infos: list[InfoLog] = Field(default_factory=list) @@ -66,6 +89,12 @@ class _OutputRecord(_StrictBase): @abstractmethod def _format_body(self) -> str: ... + def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: + return self._format_body() + + def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType: + return self._format_body() + def to_text(self) -> str: body = self._format_body() if self.errors: @@ -87,6 +116,11 @@ class _BaseComparisonRecord(_OutputRecord): return f"[step={self.location.step}] " return "" + def _format_location_prefix_rich(self) -> str: + if self.location.step is not None: + return escape(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})" @@ -149,8 +183,9 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord): model_config = ConfigDict(extra="forbid", defer_build=True) type: Literal["comparison"] = "comparison" - aligner_plan: Optional[AlignerPlan] = None + traced_plan: Optional[TracedAlignerPlan] = None replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list) + raw_bundle_info: Optional[Pair[BundleSideInfo]] = None @property def category(self) -> str: @@ -164,8 +199,8 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord): 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: - body += "\n" + _format_aligner_plan(self.aligner_plan) + if self.traced_plan is not None: + body += "\n" + _format_aligner_plan(self.traced_plan) return body @@ -225,26 +260,48 @@ class LogRecord(_OutputRecord): return "" -def _format_aligner_plan(plan: AlignerPlan) -> str: +def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str: lines: list[str] = ["Aligner Plan:"] - for side_label, side_plans in [ - ("baseline", plan.per_step_plans.x), - ("target", plan.per_step_plans.y), + for side_label, traced_side in [ + ("baseline", traced_plan.per_side.x), + ("target", traced_plan.per_side.y), ]: - if not side_plans: + if not traced_side.step_plans: lines.append(f" {side_label}: (no steps)") continue step_summaries: list[str] = [] - for step_plan in side_plans: - sub_strs: list[str] = [] - for sub in step_plan.sub_plans: - sub_strs.append(f"{sub.type}") + for traced_step in traced_side.step_plans: + sub_strs: list[str] = [ + _format_sub_plan_text(traced_sub) + for traced_sub in traced_step.sub_plans + ] summary: str = ", ".join(sub_strs) if sub_strs else "passthrough" - step_summaries.append(f"step={step_plan.step}: {summary}") + step_summaries.append(f"step={traced_step.step}: {summary}") lines.append(f" {side_label}: [{'; '.join(step_summaries)}]") + lines.extend(_format_cross_side_plan_text(traced_plan.plan)) + return "\n".join(lines) + + +def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str: + sub_desc: str = f"{traced_sub.plan.type}" + + if traced_sub.snapshot is not None: + snap = traced_sub.snapshot + in_count: int = len(snap.input_shapes) + out_count: int = len(snap.output_shapes) + in_shape: str = str(snap.input_shapes[0]) if snap.input_shapes else "?" + out_shape: str = str(snap.output_shapes[0]) if snap.output_shapes else "?" + sub_desc += f" {in_count}x{in_shape} -> {out_count}x{out_shape}" + + return sub_desc + + +def _format_cross_side_plan_text(plan: AlignerPlan) -> list[str]: + lines: list[str] = [] + if plan.token_aligner_plan is not None: num_tokens: int = len(plan.token_aligner_plan.locators.x.steps) lines.append(f" token_aligner: {num_tokens} tokens aligned") @@ -257,7 +314,7 @@ def _format_aligner_plan(plan: AlignerPlan) -> str: parts.append(f"y: {plan.axis_aligner_plan.pattern.y}") lines.append(f" axis_aligner: {', '.join(parts)}") - return "\n".join(lines) + return lines AnyRecord = Annotated[ @@ -281,64 +338,3 @@ def _get_any_record_adapter() -> TypeAdapter: def parse_record_json(json_str: str | bytes) -> AnyRecord: return _get_any_record_adapter().validate_json(json_str) - - -def _print_to_stdout(record: _OutputRecord, *, output_format: str) -> None: - if output_format == "json": - print(record.model_dump_json()) - else: - print(record.to_text()) - - -class ReportSink: - """Unified entry point for all record output.""" - - def __init__(self) -> None: - self._output_format: str = "text" - self._report_file: Optional[IO[str]] = None - self._report_path: Optional[Path] = None - - def configure( - self, - *, - output_format: str = "text", - report_path: Optional[Path] = None, - ) -> None: - self._output_format = output_format - - if report_path is not None: - try: - report_path.parent.mkdir(parents=True, exist_ok=True) - self._report_file = open(report_path, "w", encoding="utf-8") - self._report_path = report_path - except OSError as exc: - print( - f"Warning: cannot open report file {report_path}: {exc}", - file=sys.stderr, - ) - - def add(self, record: _OutputRecord) -> None: - _print_to_stdout(record, output_format=self._output_format) - - if self._report_file is not None: - self._report_file.write(record.model_dump_json()) - self._report_file.write("\n") - self._report_file.flush() - - def close(self) -> None: - if self._report_file is not None: - self._report_file.close() - self._report_file = None - - @property - def report_path(self) -> Optional[Path]: - return self._report_path - - def _reset(self) -> None: - """Reset state for test isolation.""" - self.close() - self._output_format = "text" - self._report_path = None - - -report_sink = ReportSink() diff --git a/python/sglang/srt/debug_utils/comparator/report_sink.py b/python/sglang/srt/debug_utils/comparator/report_sink.py new file mode 100644 index 000000000..61f9e9ac5 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/report_sink.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import sys +from pathlib import Path +from typing import IO, Literal, Optional + +from rich.console import Console + +from sglang.srt.debug_utils.comparator.output_types import _OutputRecord + +Verbosity = Literal["minimal", "normal", "verbose"] + + +class ReportSink: + """Unified entry point for all record output.""" + + def __init__(self) -> None: + self._output_format: str = "text" + self._verbosity: Verbosity = "normal" + self._report_file: Optional[IO[str]] = None + self._report_path: Optional[Path] = None + self._console: Optional[Console] = None + + @property + def verbosity(self) -> Verbosity: + return self._verbosity + + def configure( + self, + *, + output_format: str = "text", + report_path: Optional[Path] = None, + verbosity: Verbosity = "normal", + ) -> None: + self._output_format = output_format + self._verbosity = verbosity + + if report_path is not None: + try: + report_path.parent.mkdir(parents=True, exist_ok=True) + self._report_file = open(report_path, "w", encoding="utf-8") + self._report_path = report_path + except OSError as exc: + print( + f"Warning: cannot open report file {report_path}: {exc}", + file=sys.stderr, + ) + + def add(self, record: _OutputRecord) -> None: + self._print_to_stdout(record) + + if self._report_file is not None: + self._report_file.write(record.model_dump_json()) + self._report_file.write("\n") + self._report_file.flush() + + def close(self) -> None: + if self._report_file is not None: + self._report_file.close() + self._report_file = None + + @property + def report_path(self) -> Optional[Path]: + return self._report_path + + def _reset(self) -> None: + self.close() + self._output_format = "text" + self._verbosity = "normal" + self._report_path = None + self._console = None + + def _get_console(self) -> Console: + if self._console is None: + self._console = Console() + return self._console + + def _print_to_stdout(self, record: _OutputRecord) -> None: + if self._output_format == "json": + print(record.model_dump_json()) + else: + console: Console = self._get_console() + console.print(record.to_rich(verbosity=self._verbosity)) + console.print() # blank line between records + + +report_sink = ReportSink() diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py index ca0c5bed6..a6b6bad71 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py @@ -5,6 +5,8 @@ import torch from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import ( AlignerResult, + StepPlansResult, + SubPlansResult, _execute_step_plans, execute_aligner_plan, execute_sub_plan, @@ -31,25 +33,28 @@ register_cpu_ci(est_time=15, suite="default", nightly=True) class TestExecuteSubPlans: def test_empty_tensors_returns_none(self) -> None: - result, checks = execute_sub_plans(tensors=[], plans=[]) - assert result is None - assert checks == [] + r: SubPlansResult = execute_sub_plans(tensors=[], plans=[]) + assert r.tensor is None + assert r.checks == [] + assert r.snapshots == [] def test_no_plans_single_tensor_passthrough(self) -> None: tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0]) - result, checks = execute_sub_plans(tensors=[tensor], plans=[]) - assert result is not None - assert torch.equal(result, tensor) - assert checks == [] + r: SubPlansResult = execute_sub_plans(tensors=[tensor], plans=[]) + assert r.tensor is not None + assert torch.equal(r.tensor, tensor) + assert r.checks == [] + assert r.snapshots == [] def test_no_plans_multiple_tensors_returns_none(self) -> None: tensors: list[torch.Tensor] = [ torch.tensor([1.0]), torch.tensor([2.0]), ] - result, checks = execute_sub_plans(tensors=tensors, plans=[]) - assert result is None - assert checks == [] + r: SubPlansResult = execute_sub_plans(tensors=tensors, plans=[]) + assert r.tensor is None + assert r.checks == [] + assert r.snapshots == [] def test_with_unsharder_plan(self) -> None: t0: torch.Tensor = torch.tensor([[1.0, 2.0]]).refine_names("b", "h") @@ -61,12 +66,13 @@ class TestExecuteSubPlans: groups=[[0, 1]], ) - result, checks = execute_sub_plans(tensors=[t0, t1], plans=[plan]) + r: SubPlansResult = execute_sub_plans(tensors=[t0, t1], plans=[plan]) - assert result is not None + assert r.tensor is not None expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) - assert torch.equal(result.rename(None), expected) - assert checks == [] + assert torch.equal(r.tensor.rename(None), expected) + assert r.checks == [] + assert len(r.snapshots) == 1 class TestExecuteSubPlan: @@ -91,10 +97,13 @@ class TestExecuteStepPlans: sub_plans=[], ) - result, checks = _execute_step_plans(tensors=tensors, step_plans=[step_plan]) + r: StepPlansResult = _execute_step_plans( + tensors=tensors, step_plans=[step_plan] + ) - assert result == {} - assert checks == [] + assert r.tensors == {} + assert r.checks == [] + assert len(r.traced_side.step_plans) == 1 def test_single_step_passthrough(self) -> None: tensor: torch.Tensor = torch.tensor([1.0, 2.0]) @@ -105,11 +114,15 @@ class TestExecuteStepPlans: sub_plans=[], ) - result, checks = _execute_step_plans(tensors=[tensor], step_plans=[step_plan]) + r: StepPlansResult = _execute_step_plans( + tensors=[tensor], step_plans=[step_plan] + ) - assert 5 in result - assert torch.equal(result[5], tensor) - assert checks == [] + assert 5 in r.tensors + assert torch.equal(r.tensors[5], tensor) + assert r.checks == [] + assert len(r.traced_side.step_plans) == 1 + assert r.traced_side.step_plans[0].step == 5 class TestExecuteAlignerPlan: diff --git a/test/registered/debug_utils/comparator/conftest.py b/test/registered/debug_utils/comparator/conftest.py index 4b032ede5..3b26f0af8 100644 --- a/test/registered/debug_utils/comparator/conftest.py +++ b/test/registered/debug_utils/comparator/conftest.py @@ -1,12 +1,37 @@ +import sys import warnings +from pathlib import Path warnings.filterwarnings( "ignore", message="builtin type Swig.*", category=DeprecationWarning ) +# Add the test root to sys.path so `registered.debug_utils.comparator.testing_helpers` +# can be imported by test modules. +_TEST_ROOT: Path = Path(__file__).resolve().parents[3] +if str(_TEST_ROOT) not in sys.path: + sys.path.insert(0, str(_TEST_ROOT)) + import pytest -from sglang.srt.debug_utils.comparator.output_types import report_sink +from sglang.srt.debug_utils.comparator.report_sink import report_sink + +collect_ignore_glob: list[str] = [] + + +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "filterwarnings", + "ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning", + ) + config.addinivalue_line( + "filterwarnings", + "ignore:builtin type Swig.*:DeprecationWarning", + ) + config.addinivalue_line( + "filterwarnings", + "ignore:Named tensors and all their associated APIs:UserWarning", + ) collect_ignore_glob: list[str] = [] diff --git a/test/registered/debug_utils/comparator/test_display.py b/test/registered/debug_utils/comparator/test_display.py index 3345b29c7..451ac56b8 100644 --- a/test/registered/debug_utils/comparator/test_display.py +++ b/test/registered/debug_utils/comparator/test_display.py @@ -9,8 +9,8 @@ import torch from sglang.srt.debug_utils.comparator.display import ( _collect_input_ids_and_positions, _collect_rank_info, - _extract_parallel_info, _render_polars_as_text, + extract_parallel_info, ) from sglang.srt.debug_utils.comparator.output_types import ( InputIdsRecord, @@ -359,26 +359,26 @@ class TestExtractParallelInfo: "pp_size": 2, } row_data: dict = {} - _extract_parallel_info(row_data=row_data, info=info) + extract_parallel_info(row_data=row_data, info=info) assert row_data["tp"] == "1/4" assert row_data["pp"] == "0/2" def test_skips_error_info(self) -> None: row_data: dict = {} - _extract_parallel_info( + extract_parallel_info( row_data=row_data, info={"error": True, "tp_rank": 0, "tp_size": 1} ) assert row_data == {} def test_skips_empty_info(self) -> None: row_data: dict = {} - _extract_parallel_info(row_data=row_data, info={}) + extract_parallel_info(row_data=row_data, info={}) assert row_data == {} def test_ignores_rank_without_size(self) -> None: row_data: dict = {} - _extract_parallel_info(row_data=row_data, info={"tp_rank": 0}) + extract_parallel_info(row_data=row_data, info={"tp_rank": 0}) assert "tp" not in row_data diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 07850a18c..964727d5e 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -1458,7 +1458,7 @@ class TestEntrypointConcatMode: assert len(comparisons) == 3 def test_concat_aligner_plan_fields(self, tmp_path, capsys): - """TensorComparisonRecord.aligner_plan reports mode='concat' with plan=None.""" + """TensorComparisonRecord.traced_plan reports mode='concat' with plan=None.""" torch.manual_seed(42) records = self._run_concat( @@ -1470,8 +1470,9 @@ class TestEntrypointConcatMode: ) comparisons = _get_comparisons(records) assert len(comparisons) == 1 - plan = comparisons[0].aligner_plan - assert plan is not None + traced_plan = comparisons[0].traced_plan + assert traced_plan is not None + plan = traced_plan.plan assert plan.token_aligner_mode == "concat_steps" assert plan.token_aligner_plan is None @@ -4250,7 +4251,7 @@ class TestReportOutput: def test_streaming_flush(self, tmp_path, capsys): """Report file is flushed after each record (readable before close).""" - from sglang.srt.debug_utils.comparator.output_types import report_sink + from sglang.srt.debug_utils.comparator.report_sink import report_sink report_file: Path = tmp_path / "stream_report.jsonl" report_sink.configure( @@ -4362,5 +4363,102 @@ class TestEntrypointAutoDescend: run(parse_args(argv)) +class TestEntrypointDpAttentionMissingAlias: + """Regression: dp-attention without ``# dp:=attn_dp`` → shape mismatch failure. + + In dp-attention mode (tp_size=2, attn_dp_size=2), layer_input is dumped + after prepare_attn which DP-distributes tokens. One rank gets 0 tokens + (shape [0, H]), the other gets all tokens (shape [T, H]). + + Without ``# dp:=attn_dp`` in dims, the comparator has no dp_rank/dp_size + to filter on, so it picks one rank via TP pick — potentially the empty + one — causing a shape mismatch with the baseline. + """ + + @staticmethod + def _sglang_dp_attn_parallel_info(*, tp_rank: int) -> dict: + return { + "tp_rank": tp_rank, + "tp_size": 2, + "pp_rank": 0, + "pp_size": 1, + "moe_ep_rank": 0, + "moe_ep_size": 1, + "moe_tp_rank": tp_rank, + "moe_tp_size": 2, + "moe_dp_rank": 0, + "moe_dp_size": 1, + "enable_dp_attention": True, + "attn_tp_rank": 0, + "attn_tp_size": 1, + "attn_dp_rank": tp_rank, + "attn_dp_size": 2, + "local_attn_dp_rank": tp_rank, + "local_attn_dp_size": 2, + "attn_cp_rank": 0, + "attn_cp_size": 1, + } + + def test_missing_dp_alias_causes_shape_mismatch( + self, tmp_path: Path, capsys + ) -> None: + """dims='t h' (no dp:=attn_dp) → comparator picks empty rank → shape_mismatch failure.""" + torch.manual_seed(42) + tensor_data: torch.Tensor = torch.randn(5, 8) + target_data: torch.Tensor = tensor_data + torch.randn(5, 8) * 0.001 + + for side_name, data in [("baseline", tensor_data), ("target", target_data)]: + side_dir: Path = tmp_path / side_name + side_dir.mkdir() + + # Baseline: single rank, no DP attention + if side_name == "baseline": + _create_rank_dump( + side_dir, + rank=0, + name="layer_input", + tensor=data, + dims="t h", + parallel_info={"tp_rank": 0, "tp_size": 1}, + framework="sglang", + ) + else: + # Target: dp-attention, tp_rank=0 gets 0 tokens, tp_rank=1 gets all + _create_rank_dump( + side_dir, + rank=0, + name="layer_input", + tensor=torch.empty(0, 8), + dims="t h", + parallel_info=self._sglang_dp_attn_parallel_info(tp_rank=0), + framework="sglang", + ) + _create_rank_dump( + side_dir, + rank=1, + name="layer_input", + tensor=data, + dims="t h", + parallel_info=self._sglang_dp_attn_parallel_info(tp_rank=1), + framework="sglang", + ) + + argv: list[str] = _make_argv( + tmp_path / "baseline" / _FIXED_EXP_NAME, + tmp_path / "target" / _FIXED_EXP_NAME, + diff_threshold=1e-3, + ) + records, exit_code = _run_and_parse(argv, capsys) + + assert exit_code == 1 + + comparisons: list[TensorComparisonRecord] = _get_comparisons(records) + assert len(comparisons) == 1 + comparison: TensorComparisonRecord = comparisons[0] + assert comparison.shape_mismatch is True + assert comparison.diff is None + assert comparison.category == "failed" + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_log_sink.py b/test/registered/debug_utils/comparator/test_log_sink.py index d553d6c42..348ea168a 100644 --- a/test/registered/debug_utils/comparator/test_log_sink.py +++ b/test/registered/debug_utils/comparator/test_log_sink.py @@ -7,8 +7,8 @@ from sglang.srt.debug_utils.comparator.log_sink import LogSink from sglang.srt.debug_utils.comparator.output_types import ( ErrorLog, InfoLog, - report_sink, ) +from sglang.srt.debug_utils.comparator.report_sink import report_sink from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) diff --git a/test/registered/debug_utils/comparator/test_output_types.py b/test/registered/debug_utils/comparator/test_output_types.py index d7f96eeb4..e7b5f85f3 100644 --- a/test/registered/debug_utils/comparator/test_output_types.py +++ b/test/registered/debug_utils/comparator/test_output_types.py @@ -1,18 +1,67 @@ import sys +from io import StringIO import pytest +from registered.debug_utils.comparator.testing_helpers import make_diff as _make_diff +from registered.debug_utils.comparator.testing_helpers import ( + make_tensor_info as _make_tensor_info, +) +from rich.console import Console +from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan +from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( + TracedAlignerPlan, + TracedSidePlan, + TracedStepPlan, + TracedSubPlan, +) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ( + ReordererPlan, + ZigzagToNaturalParams, +) +from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import ( + TokenAlignerPlan, + TokenLocator, +) +from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( + ConcatParams, + UnsharderPlan, +) +from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout from sglang.srt.debug_utils.comparator.output_types import ( + ConfigRecord, ErrorLog, InfoLog, LogRecord, + NonTensorComparisonRecord, + RecordLocation, + SkipComparisonRecord, + SummaryRecord, + TensorComparisonRecord, + _format_aligner_plan, _split_logs, ) +from sglang.srt.debug_utils.comparator.utils import Pair from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) +def _render_rich(renderable: object) -> str: + buf: StringIO = StringIO() + Console(file=buf, force_terminal=False, width=120).print(renderable) + return buf.getvalue().rstrip("\n") + + +# --------------------------------------------------------------------------- +# Existing tests (preserved) +# --------------------------------------------------------------------------- + + def test_split_logs_mixed_list() -> None: """_split_logs correctly partitions a mixed list of ErrorLog and InfoLog.""" errors, infos = _split_logs( @@ -40,5 +89,523 @@ def test_log_record_to_text_format() -> None: assert "ℹ fyi" in text +class TestLogRecord: + def test_format_body_returns_empty(self) -> None: + record: LogRecord = LogRecord() + assert record._format_body() == "" + + def test_format_rich_body_returns_empty(self) -> None: + record: LogRecord = LogRecord() + assert record._format_rich_body() == "" + + def test_to_text_empty_no_logs(self) -> None: + record: LogRecord = LogRecord() + assert record.to_text() == "" + + def test_to_text_with_errors_and_infos(self) -> None: + record: LogRecord = LogRecord( + errors=[ErrorLog(category="a", message="bad thing")], + infos=[InfoLog(category="b", message="fyi")], + ) + text: str = record.to_text() + assert text == "\n ✗ bad thing\n ℹ fyi" + + +# --------------------------------------------------------------------------- +# ConfigRecord +# --------------------------------------------------------------------------- + + +class TestConfigRecord: + def test_format_body(self) -> None: + record: ConfigRecord = ConfigRecord(config={"a": 1, "b": "two"}) + assert record._format_body() == "Config: {'a': 1, 'b': 'two'}" + + def test_to_text_with_errors(self) -> None: + record: ConfigRecord = ConfigRecord( + config={"x": 1}, + errors=[ErrorLog(category="cfg", message="bad config")], + ) + text: str = record.to_text() + assert text.startswith("Config: {'x': 1}") + assert "✗ bad config" in text + + +# --------------------------------------------------------------------------- +# SkipComparisonRecord +# --------------------------------------------------------------------------- + + +class TestSkipComparisonRecord: + def test_format_body_no_step(self) -> None: + record: SkipComparisonRecord = SkipComparisonRecord( + name="layer.weight", + reason="zero-dim tensor", + ) + assert record._format_body() == "Skip: layer.weight (zero-dim tensor)" + + def test_format_body_with_step(self) -> None: + record: SkipComparisonRecord = SkipComparisonRecord( + name="layer.weight", + reason="scalar", + location=RecordLocation(step=3), + ) + assert record._format_body() == "Skip: layer.weight (step=3) (scalar)" + + def test_category_skipped(self) -> None: + record: SkipComparisonRecord = SkipComparisonRecord( + name="x", + reason="r", + ) + assert record.category == "skipped" + + def test_category_failed(self) -> None: + record: SkipComparisonRecord = SkipComparisonRecord( + name="x", + reason="r", + errors=[ErrorLog(category="e", message="boom")], + ) + assert record.category == "failed" + + +# --------------------------------------------------------------------------- +# NonTensorComparisonRecord +# --------------------------------------------------------------------------- + + +class TestNonTensorComparisonRecord: + def test_format_body_equal(self) -> None: + record: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="config.lr", + baseline_value="0.001", + target_value="0.001", + baseline_type="float", + target_type="float", + values_equal=True, + ) + assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]" + + def test_format_body_not_equal(self) -> None: + record: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="config.lr", + baseline_value="0.001", + target_value="0.01", + baseline_type="float", + target_type="float", + values_equal=False, + ) + assert record._format_body() == ( + "NonTensor: config.lr\n" + " baseline = 0.001 (float)\n" + " target = 0.01 (float)" + ) + + def test_with_step(self) -> None: + record: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="bias", + baseline_value="True", + target_value="True", + baseline_type="bool", + target_type="bool", + values_equal=True, + location=RecordLocation(step=5), + ) + assert "(step=5)" in record._format_body() + + def test_category(self) -> None: + passed: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="x", + baseline_value="1", + target_value="1", + baseline_type="int", + target_type="int", + values_equal=True, + ) + failed: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="x", + baseline_value="1", + target_value="2", + baseline_type="int", + target_type="int", + values_equal=False, + ) + assert passed.category == "passed" + assert failed.category == "failed" + + +# --------------------------------------------------------------------------- +# SummaryRecord +# --------------------------------------------------------------------------- + + +class TestSummaryRecord: + def test_format_body(self) -> None: + record: SummaryRecord = SummaryRecord( + total=10, + passed=7, + failed=2, + skipped=1, + ) + assert record._format_body() == ( + "Summary: 7 passed, 2 failed, 1 skipped (total 10)" + ) + + def test_validation_error(self) -> None: + with pytest.raises(ValueError, match="total=5 !="): + SummaryRecord(total=5, passed=1, failed=1, skipped=1) + + +# --------------------------------------------------------------------------- +# TensorComparisonRecord._format_body +# --------------------------------------------------------------------------- + + +class TestTensorComparisonRecordFormatBody: + def test_basic(self) -> None: + record: TensorComparisonRecord = TensorComparisonRecord( + name="hidden", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + ) + body: str = record._format_body() + + assert body == ( + "Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n" + "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005" + ) + + def test_with_replicated_checks(self) -> None: + from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult + + record: TensorComparisonRecord = TensorComparisonRecord( + name="hidden", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + replicated_checks=[ + ReplicatedCheckResult( + axis="tp", + group_index=0, + compared_index=1, + baseline_index=0, + passed=True, + atol=1e-3, + diff=_make_diff( + rel_diff=1e-6, max_abs_diff=1e-5, mean_abs_diff=1e-6 + ), + ), + ], + ) + body: str = record._format_body() + + assert body == ( + "Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n" + "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n" + "Replicated checks:\n" + " ✅ axis=tp group=0 idx=1 vs 0: " + "rel_diff=1.000000e-06 max_abs_diff=1.000000e-05 mean_abs_diff=1.000000e-06" + ) + + def test_with_aligner_plan(self) -> None: + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair(x=[], y=[]), + ) + traced: TracedAlignerPlan = TracedAlignerPlan( + plan=plan, + per_side=Pair( + x=TracedSidePlan(step_plans=[]), + y=TracedSidePlan(step_plans=[]), + ), + ) + record: TensorComparisonRecord = TensorComparisonRecord( + name="hidden", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + traced_plan=traced, + ) + body: str = record._format_body() + + assert body == ( + "Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n" + "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n" + "Aligner Plan:\n" + " baseline: (no steps)\n" + " target: (no steps)" + ) + + def test_with_step(self) -> None: + record: TensorComparisonRecord = TensorComparisonRecord( + name="hidden", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + location=RecordLocation(step=2), + ) + body: str = record._format_body() + + assert body == ( + "[step=2] " + "Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n" + "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005" + ) + + +# --------------------------------------------------------------------------- +# _format_aligner_plan +# --------------------------------------------------------------------------- + + +def _wrap_plan(plan: AlignerPlan) -> TracedAlignerPlan: + """Wrap an AlignerPlan into a TracedAlignerPlan with no snapshots.""" + baseline_traced_steps: list[TracedStepPlan] = [ + TracedStepPlan( + step=sp.step, + input_object_indices=sp.input_object_indices, + sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans], + ) + for sp in plan.per_step_plans.x + ] + target_traced_steps: list[TracedStepPlan] = [ + TracedStepPlan( + step=sp.step, + input_object_indices=sp.input_object_indices, + sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans], + ) + for sp in plan.per_step_plans.y + ] + return TracedAlignerPlan( + plan=plan, + per_side=Pair( + x=TracedSidePlan(step_plans=baseline_traced_steps), + y=TracedSidePlan(step_plans=target_traced_steps), + ), + ) + + +class TestFormatAlignerPlan: + def test_passthrough(self) -> None: + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair(x=[], y=[]), + ) + result: str = _format_aligner_plan(_wrap_plan(plan)) + + assert result == ( + "Aligner Plan:\n" " baseline: (no steps)\n" " target: (no steps)" + ) + + def test_unsharder(self) -> None: + unsharder: UnsharderPlan = UnsharderPlan( + axis=ParallelAxis.TP, + params=ConcatParams(dim_name="h"), + groups=[[0, 1]], + ) + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair( + x=[], + y=[ + AlignerPerStepPlan( + step=0, input_object_indices=[0, 1], sub_plans=[unsharder] + ) + ], + ), + ) + result: str = _format_aligner_plan(_wrap_plan(plan)) + + assert result == ( + "Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: unsharder]" + ) + + def test_reorderer(self) -> None: + reorderer: ReordererPlan = ReordererPlan( + params=ZigzagToNaturalParams(dim_name="s", cp_size=2), + ) + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair( + x=[], + y=[ + AlignerPerStepPlan( + step=0, input_object_indices=[0], sub_plans=[reorderer] + ) + ], + ), + ) + result: str = _format_aligner_plan(_wrap_plan(plan)) + + assert result == ( + "Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: reorderer]" + ) + + def test_multi_step(self) -> None: + unsharder: UnsharderPlan = UnsharderPlan( + axis=ParallelAxis.TP, + params=ConcatParams(dim_name="h"), + groups=[[0, 1]], + ) + reorderer: ReordererPlan = ReordererPlan( + params=ZigzagToNaturalParams(dim_name="s", cp_size=2), + ) + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair( + x=[], + y=[ + AlignerPerStepPlan( + step=0, input_object_indices=[0, 1], sub_plans=[unsharder] + ), + AlignerPerStepPlan( + step=1, input_object_indices=[0], sub_plans=[reorderer] + ), + ], + ), + ) + result: str = _format_aligner_plan(_wrap_plan(plan)) + + assert result == ( + "Aligner Plan:\n" + " baseline: (no steps)\n" + " target: [step=0: unsharder; step=1: reorderer]" + ) + + def test_with_token_aligner(self) -> None: + ta_plan: TokenAlignerPlan = TokenAlignerPlan( + locators=Pair( + x=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]), + y=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]), + ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), + ) + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair(x=[], y=[]), + token_aligner_plan=ta_plan, + ) + result: str = _format_aligner_plan(_wrap_plan(plan)) + + assert result == ( + "Aligner Plan:\n" + " baseline: (no steps)\n" + " target: (no steps)\n" + " token_aligner: 3 tokens aligned" + ) + + def test_with_axis_aligner(self) -> None: + aa_plan: AxisAlignerPlan = AxisAlignerPlan( + pattern=Pair(x="b s d -> s b d", y=None), + ) + plan: AlignerPlan = AlignerPlan( + per_step_plans=Pair(x=[], y=[]), + axis_aligner_plan=aa_plan, + ) + result: str = _format_aligner_plan(_wrap_plan(plan)) + + assert result == ( + "Aligner Plan:\n" + " baseline: (no steps)\n" + " target: (no steps)\n" + " axis_aligner: x: b s d -> s b d" + ) + + +# --------------------------------------------------------------------------- +# _OutputRecord log attachment (to_text) +# --------------------------------------------------------------------------- + + +class TestOutputRecordLogAttachment: + def test_to_text_no_logs(self) -> None: + record: ConfigRecord = ConfigRecord(config={"a": 1}) + text: str = record.to_text() + + assert text == "Config: {'a': 1}" + + def test_to_text_errors_only(self) -> None: + record: ConfigRecord = ConfigRecord( + config={"a": 1}, + errors=[ErrorLog(category="x", message="err1")], + ) + text: str = record.to_text() + + assert text == "Config: {'a': 1}\n ✗ err1" + + def test_to_text_infos_only(self) -> None: + record: ConfigRecord = ConfigRecord( + config={"a": 1}, + infos=[InfoLog(category="x", message="note1")], + ) + text: str = record.to_text() + + assert text == "Config: {'a': 1}\n ℹ note1" + + def test_to_text_mixed(self) -> None: + record: ConfigRecord = ConfigRecord( + config={"a": 1}, + errors=[ErrorLog(category="x", message="err1")], + infos=[InfoLog(category="y", message="note1")], + ) + text: str = record.to_text() + + assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1" + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/testing_helpers.py b/test/registered/debug_utils/comparator/testing_helpers.py new file mode 100644 index 000000000..1106e3a1a --- /dev/null +++ b/test/registered/debug_utils/comparator/testing_helpers.py @@ -0,0 +1,84 @@ +"""Shared test helpers for comparator tests.""" + +from __future__ import annotations + +from typing import Optional + +from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( + DiffInfo, + TensorInfo, + TensorStats, +) + +DEFAULT_PERCENTILES: dict[int, float] = { + 1: -1.8, + 5: -1.5, + 50: 0.0, + 95: 1.5, + 99: 1.8, +} + +DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = { + 1: 0.0001, + 5: 0.0001, + 50: 0.0002, + 95: 0.0004, + 99: 0.0005, +} + + +def make_stats( + mean: float = 0.0, + abs_mean: float = 0.8, + std: float = 1.0, + min: float = -2.0, + max: float = 2.0, + percentiles: Optional[dict[int, float]] = None, +) -> TensorStats: + return TensorStats( + mean=mean, + abs_mean=abs_mean, + std=std, + min=min, + max=max, + percentiles=percentiles if percentiles is not None else DEFAULT_PERCENTILES, + ) + + +def make_diff( + rel_diff: float = 0.0001, + max_abs_diff: float = 0.0005, + mean_abs_diff: float = 0.0002, + abs_diff_percentiles: Optional[dict[int, float]] = None, + diff_threshold: float = 1e-3, + passed: bool = True, +) -> DiffInfo: + return DiffInfo( + rel_diff=rel_diff, + max_abs_diff=max_abs_diff, + mean_abs_diff=mean_abs_diff, + abs_diff_percentiles=( + abs_diff_percentiles + if abs_diff_percentiles is not None + else DEFAULT_ABS_DIFF_PERCENTILES + ), + max_diff_coord=[2, 3], + baseline_at_max=1.0, + target_at_max=1.0005, + diff_threshold=diff_threshold, + passed=passed, + ) + + +def make_tensor_info( + shape: Optional[list[int]] = None, + dtype: str = "torch.float32", + stats: Optional[TensorStats] = None, + sample: Optional[str] = None, +) -> TensorInfo: + return TensorInfo( + shape=shape if shape is not None else [4, 8], + dtype=dtype, + stats=stats if stats is not None else make_stats(), + sample=sample, + )