From e8dd14519dcd0ced95f4b42ee5c7265ad157abcc Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 26 Feb 2026 10:03:22 +0800 Subject: [PATCH] Add aligner entrypoint and bundle handler in dump comparator (#19375) --- .../comparator/aligner/entrypoint/__init__.py | 0 .../comparator/aligner/entrypoint/executor.py | 105 +++++++++ .../comparator/aligner/entrypoint/planner.py | 75 +++++++ .../comparator/aligner/entrypoint/types.py | 22 ++ .../comparator/bundle_comparator.py | 105 +++++++++ .../debug_utils/comparator/bundle_matcher.py | 46 ++++ .../srt/debug_utils/comparator/entrypoint.py | 104 ++++++--- .../srt/debug_utils/comparator/pipeline.py | 169 --------------- .../comparator/aligner/entrypoint/conftest.py | 0 .../aligner/entrypoint/test_executor.py | 199 ++++++++++++++++++ .../aligner/entrypoint/test_planner.py | 146 +++++++++++++ .../comparator/test_bundle_matcher.py | 150 +++++++++++++ 12 files changed, 926 insertions(+), 195 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/aligner/entrypoint/__init__.py create mode 100644 python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py create mode 100644 python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py create mode 100644 python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py create mode 100644 python/sglang/srt/debug_utils/comparator/bundle_comparator.py create mode 100644 python/sglang/srt/debug_utils/comparator/bundle_matcher.py delete mode 100644 python/sglang/srt/debug_utils/comparator/pipeline.py create mode 100644 test/registered/debug_utils/comparator/aligner/entrypoint/conftest.py create mode 100644 test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py create mode 100644 test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py create mode 100644 test/registered/debug_utils/comparator/test_bundle_matcher.py diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/__init__.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py new file mode 100644 index 000000000..57f1f0717 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepPlan, + AlignerPerStepSubPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import ( + execute_reorderer_plan, +) +from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan +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.utils import Pair + + +@dataclass(frozen=True) +class AlignerResult: + tensors: Optional[Pair[torch.Tensor]] + failed_side_xy: Optional[str] # "x" or "y"; None if success + + +def execute_aligner_plan( + *, + tensors_pair: Pair[list[torch.Tensor]], + plan: AlignerPlan, +) -> AlignerResult: + """Execute unified unshard/reorder per side, then combine.""" + + step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans( + tensors=tensors_pair.x, step_plans=plan.per_step_plans.x + ) + step_tensors_y: dict[int, torch.Tensor] = _execute_step_plans( + tensors=tensors_pair.y, step_plans=plan.per_step_plans.y + ) + + if not step_tensors_x or not step_tensors_y: + failed_side_xy: str = "x" if not step_tensors_x else "y" + return AlignerResult(tensors=None, failed_side_xy=failed_side_xy) + + assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1 + combined = Pair( + x=list(step_tensors_x.values())[0], + y=list(step_tensors_y.values())[0], + ) + + return AlignerResult(tensors=combined, failed_side_xy=None) + + +def _execute_step_plans( + tensors: list[torch.Tensor], + step_plans: list[AlignerPerStepPlan], +) -> dict[int, torch.Tensor]: + result: dict[int, torch.Tensor] = {} + + for step_plan in step_plans: + step_tensors: list[torch.Tensor] = [ + tensors[i] for i in step_plan.input_object_indices + ] + tensor: Optional[torch.Tensor] = execute_sub_plans( + tensors=step_tensors, plans=step_plan.sub_plans + ) + if tensor is not None: + result[step_plan.step] = tensor + + return result + + +def execute_sub_plans( + tensors: list[torch.Tensor], + plans: list[AlignerPerStepSubPlan], +) -> Optional[torch.Tensor]: + if not tensors: + return None + + if not plans: + if len(tensors) != 1: + return None + return tensors[0] + + current = tensors + for plan in plans: + current = execute_sub_plan(tensors=current, plan=plan) + + assert len(current) == 1 + return current[0] + + +def execute_sub_plan( + tensors: list[torch.Tensor], + plan: AlignerPerStepSubPlan, +) -> list[torch.Tensor]: + if isinstance(plan, UnsharderPlan): + return execute_unsharder_plan(plan, tensors) + elif isinstance(plan, ReordererPlan): + return execute_reorderer_plan(plan, tensors) + else: + raise NotImplementedError(f"Unknown {plan=}") diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py new file mode 100644 index 000000000..eb94cb355 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import Any + +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepPlan, + AlignerPerStepSubPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import ( + compute_reorderer_plans, +) +from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import ( + normalize_parallel_info, +) +from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import ( + compute_unsharder_plan, +) +from sglang.srt.debug_utils.comparator.dims import parse_dims +from sglang.srt.debug_utils.comparator.utils import Pair + + +def compute_aligner_plan( + *, + metas_pair: Pair[list[dict[str, Any]]], +) -> AlignerPlan: + return AlignerPlan( + per_step_plans=metas_pair.map( + lambda metas: _compute_per_step_plans(metas=metas) + ), + ) + + +def _compute_per_step_plans(metas: list[dict[str, Any]]) -> list[AlignerPerStepPlan]: + step_to_input_indices: dict[int, list[int]] = {} + for i, meta in enumerate(metas): + step: int = int(meta["step"]) + step_to_input_indices.setdefault(step, []).append(i) + + result: list[AlignerPerStepPlan] = [] + for step in sorted(step_to_input_indices): + input_indices: list[int] = step_to_input_indices[step] + step_metas: list[dict[str, Any]] = [metas[idx] for idx in input_indices] + plans: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( + metas=step_metas + ) + result.append( + AlignerPerStepPlan( + step=step, input_object_indices=input_indices, sub_plans=plans + ) + ) + + return result + + +def compute_per_step_sub_plans( + metas: list[dict[str, Any]], +) -> list[AlignerPerStepSubPlan]: + if not metas or len(metas) == 1: + return [] + + dims_str = metas[0].get("dims") + if dims_str is None: + return [] + + dim_specs = parse_dims(dims_str) + parallel_infos = [normalize_parallel_info(meta) for meta in metas] + + unsharder_plans = compute_unsharder_plan( + dim_specs=dim_specs, parallel_infos=parallel_infos + ) + reorderer_plans = compute_reorderer_plans( + dim_specs=dim_specs, parallel_infos=parallel_infos + ) + return [*unsharder_plans, *reorderer_plans] diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py new file mode 100644 index 000000000..2f14f093b --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Union + +from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan +from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan +from sglang.srt.debug_utils.comparator.utils import Pair + +AlignerPerStepSubPlan = Union[UnsharderPlan, ReordererPlan] + + +@dataclass(frozen=True) +class AlignerPerStepPlan: + step: int + input_object_indices: list[int] + sub_plans: list[AlignerPerStepSubPlan] + + +@dataclass(frozen=True) +class AlignerPlan: + per_step_plans: Pair[list[AlignerPerStepPlan]] diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py new file mode 100644 index 000000000..5475d8c69 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -0,0 +1,105 @@ +"""Compare two tensor bundles.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Union + +import torch + +from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import ( + AlignerResult, + execute_aligner_plan, +) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import ( + compute_aligner_plan, +) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan +from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonRecord, + SkipRecord, +) +from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( + compare_tensor_pair, +) +from sglang.srt.debug_utils.comparator.utils import Pair +from sglang.srt.debug_utils.comparator.warning_sink import warning_sink +from sglang.srt.debug_utils.dump_loader import ValueWithMeta + +_FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"} + + +def compare_bundle_pair( + *, + name: str, + filenames_pair: Pair[list[str]], + baseline_path: Path, + target_path: Path, + diff_threshold: float, +) -> Union[ComparisonRecord, SkipRecord]: + with warning_sink.context() as collected_warnings: + result = _compare_bundle_pair_raw( + name=name, + filenames_pair=filenames_pair, + baseline_path=baseline_path, + target_path=target_path, + diff_threshold=diff_threshold, + ) + + return result.model_copy(update={"warnings": collected_warnings}) + + +def _compare_bundle_pair_raw( + *, + name: str, + filenames_pair: Pair[list[str]], + baseline_path: Path, + target_path: Path, + diff_threshold: float, +) -> Union[ComparisonRecord, SkipRecord]: + # 1. Load (tensor + meta, ungrouped) + valid_pair: Pair[list[ValueWithMeta]] = Pair( + x=_load_valid_tensors(filenames=filenames_pair.x, base_path=baseline_path), + y=_load_valid_tensors(filenames=filenames_pair.y, base_path=target_path), + ) + + 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) + + # 2. Plan (meta only, no tensor) + metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map( + lambda items: [it.meta for it in items] + ) + plan: AlignerPlan = compute_aligner_plan(metas_pair=metas_pair) + + # 3. Execute (tensor + plan only, no meta) + tensors_pair: Pair[list[torch.Tensor]] = valid_pair.map( + lambda items: [it.value for it in items] + ) + aligner_result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + if aligner_result.tensors is None: + assert aligner_result.failed_side_xy is not None + side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy] + reason = f"{side_name}_load_failed" + return SkipRecord(name=name, reason=reason) + + # 4. Compare + info = compare_tensor_pair( + x_baseline=aligner_result.tensors.x, + x_target=aligner_result.tensors.y, + name=name, + diff_threshold=diff_threshold, + ) + return ComparisonRecord(**info.model_dump()) + + +def _load_valid_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]: + return [ + x + for f in filenames + if isinstance((x := ValueWithMeta.load(base_path / f)).value, torch.Tensor) + ] diff --git a/python/sglang/srt/debug_utils/comparator/bundle_matcher.py b/python/sglang/srt/debug_utils/comparator/bundle_matcher.py new file mode 100644 index 000000000..dacf73462 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/bundle_matcher.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from typing import Any + +import polars as pl + +from sglang.srt.debug_utils.comparator.utils import Pair +from sglang.srt.debug_utils.dump_loader import filter_rows + + +@dataclass(frozen=True) +class TensorFileInfo: + filename: str + name: str + step: int + + +TensorBundleInfo = list[TensorFileInfo] + + +def match_bundles( + *, + dfs: Pair[pl.DataFrame], + skip_keys: set[str], +) -> list[Pair[TensorBundleInfo]]: + match_key_cols: list[str] = [c for c in dfs.y.columns if c not in skip_keys] + unique_keys: pl.DataFrame = dfs.y.select(match_key_cols).unique(maintain_order=True) + + results: list[Pair[TensorBundleInfo]] = [] + for key_values in unique_keys.iter_rows(named=True): + result = dfs.map( + lambda df: _rows_to_tensor_infos(filter_rows(df, conditions=key_values)) + ) + results.append(result) + + return results + + +def _rows_to_tensor_infos(rows: list[dict[str, Any]]) -> list[TensorFileInfo]: + tensor_info_fields: set[str] = {f.name for f in dataclasses.fields(TensorFileInfo)} + return [ + TensorFileInfo(**{k: v for k, v in row.items() if k in tensor_info_fields}) + for row in rows + ] diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 2ab6b3b92..367f345bf 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -1,17 +1,26 @@ +from __future__ import annotations + import argparse from pathlib import Path +from typing import Iterator, Union import polars as pl +from sglang.srt.debug_utils.comparator.bundle_comparator import compare_bundle_pair +from sglang.srt.debug_utils.comparator.bundle_matcher import ( + TensorBundleInfo, + match_bundles, +) from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonRecord, ConfigRecord, + SkipRecord, SummaryRecord, print_record, ) -from sglang.srt.debug_utils.comparator.pipeline import process_tensor_group -from sglang.srt.debug_utils.dump_loader import filter_rows, read_meta - -_NON_KEY_COLS = {"dump_index", "filename"} +from sglang.srt.debug_utils.comparator.utils import Pair +from sglang.srt.debug_utils.comparator.warning_sink import warning_sink +from sglang.srt.debug_utils.dump_loader import read_meta def main() -> None: @@ -20,6 +29,32 @@ def main() -> None: def run(args: argparse.Namespace) -> None: + print_record( + ConfigRecord.from_args(args), + output_format=args.output_format, + ) + + warning_sink.set_output_format(args.output_format) + + dfs: Pair[pl.DataFrame] = _read_df(args) + + bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=dfs, + skip_keys=_compute_skip_keys(args), + ) + + comparison_records = _compare_bundle_pairs( + bundle_info_pairs=bundle_info_pairs, + baseline_path=Path(args.baseline_path), + target_path=Path(args.target_path), + diff_threshold=args.diff_threshold, + ) + _consume_comparison_records( + comparison_records=comparison_records, output_format=args.output_format + ) + + +def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]: df_baseline = read_meta(args.baseline_path) df_target = read_meta(args.target_path) @@ -30,37 +65,54 @@ def run(args: argparse.Namespace) -> None: df_target = df_target.filter(pl.col("filename").str.contains(args.filter)) assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"]) - print_record( - ConfigRecord.from_args(args), - output_format=args.output_format, - ) + return Pair(x=df_baseline, y=df_target) - counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} - grouping: str = args.grouping - non_key_cols = _NON_KEY_COLS | ({"rank"} if grouping == "logical" else set()) - key_cols = [c for c in df_target.columns if c not in non_key_cols] - tensor_group_keys = df_target.unique(subset=key_cols) +def _compute_skip_keys(args: argparse.Namespace) -> set[str]: + skip_keys: set[str] = {"dump_index", "filename"} + if args.grouping == "logical": + skip_keys |= {"rank"} + return skip_keys - for tensor_group_key in tensor_group_keys.iter_rows(named=True): - conditions = {k: tensor_group_key[k] for k in key_cols} - baseline_rows = filter_rows(df_baseline, conditions=conditions) - target_rows = filter_rows(df_target, conditions=conditions) - record = process_tensor_group( - name=tensor_group_key["name"], - baseline_filenames=[r["filename"] for r in baseline_rows], - target_filenames=[r["filename"] for r in target_rows], - baseline_path=Path(args.baseline_path), - target_path=Path(args.target_path), - diff_threshold=args.diff_threshold, +def _compare_bundle_pairs( + *, + bundle_info_pairs: list[Pair[TensorBundleInfo]], + baseline_path: Path, + target_path: Path, + diff_threshold: float, +) -> Iterator[Union[ComparisonRecord, SkipRecord]]: + for bundle_info_pair in bundle_info_pairs: + if not bundle_info_pair.y: + continue + + name: str = bundle_info_pair.y[0].name + filenames_pair: Pair[list[str]] = bundle_info_pair.map( + lambda infos: [info.filename for info in infos] ) + yield compare_bundle_pair( + name=name, + filenames_pair=filenames_pair, + baseline_path=baseline_path, + target_path=target_path, + diff_threshold=diff_threshold, + ) + + +def _consume_comparison_records( + *, + comparison_records: Iterator[Union[ComparisonRecord, SkipRecord]], + output_format: str, +) -> None: + counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} + + for record in comparison_records: counts[record.category] += 1 - print_record(record, output_format=args.output_format) + print_record(record, output_format=output_format) print_record( SummaryRecord(total=sum(counts.values()), **counts), - output_format=args.output_format, + output_format=output_format, ) diff --git a/python/sglang/srt/debug_utils/comparator/pipeline.py b/python/sglang/srt/debug_utils/comparator/pipeline.py deleted file mode 100644 index 6612b7f20..000000000 --- a/python/sglang/srt/debug_utils/comparator/pipeline.py +++ /dev/null @@ -1,169 +0,0 @@ -from pathlib import Path -from typing import Any, Optional, Union - -import torch - -from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import ( - execute_reorderer_plan, -) -from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import ( - compute_reorderer_plans, -) -from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan -from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import ( - execute_unsharder_plan, -) -from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import ( - normalize_parallel_info, -) -from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import ( - compute_unsharder_plan, -) -from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan -from sglang.srt.debug_utils.comparator.dims import parse_dims -from sglang.srt.debug_utils.comparator.output_types import ( - AnyWarning, - ComparisonRecord, - SkipRecord, -) -from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( - compare_tensor_pair, -) -from sglang.srt.debug_utils.comparator.warning_sink import warning_sink -from sglang.srt.debug_utils.dump_loader import ValueWithMeta - -Plan = Union[UnsharderPlan, ReordererPlan] - - -def process_tensor_group( - *, - name: str, - baseline_filenames: list[str], - target_filenames: list[str], - baseline_path: Path, - target_path: Path, - diff_threshold: float, -) -> ComparisonRecord | SkipRecord: - with warning_sink.context() as collected_warnings: - return _process_tensor_group_raw( - name=name, - baseline_filenames=baseline_filenames, - target_filenames=target_filenames, - baseline_path=baseline_path, - target_path=target_path, - diff_threshold=diff_threshold, - collected_warnings=collected_warnings, - ) - - -def _process_tensor_group_raw( - *, - name: str, - baseline_filenames: list[str], - target_filenames: list[str], - baseline_path: Path, - target_path: Path, - diff_threshold: float, - collected_warnings: list[AnyWarning], -) -> ComparisonRecord | SkipRecord: - b_tensors = _load_tensors(baseline_filenames, baseline_path) - t_tensors = _load_tensors(target_filenames, target_path) - - b_plans, t_plans = _compute_plans( - baseline_metas=[item.meta for item in b_tensors], - target_metas=[item.meta for item in t_tensors], - ) - - b_extracted = _extract_tensors(b_tensors) - t_extracted = _extract_tensors(t_tensors) - del b_tensors, t_tensors - - b_tensor = _execute_plans(b_extracted, b_plans) - t_tensor = _execute_plans(t_extracted, t_plans) - - if b_tensor is None or t_tensor is None: - reason = "baseline_load_failed" if b_tensor is None else "target_load_failed" - return SkipRecord(name=name, reason=reason, warnings=collected_warnings) - - info = compare_tensor_pair( - x_baseline=b_tensor, - x_target=t_tensor, - name=name, - diff_threshold=diff_threshold, - ) - - return ComparisonRecord(**info.model_dump(), warnings=collected_warnings) - - -def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]: - return [ValueWithMeta.load(base_path / f) for f in filenames] - - -def _compute_plans( - *, - baseline_metas: list[dict[str, Any]], - target_metas: list[dict[str, Any]], -) -> tuple[list[Plan], list[Plan]]: - """This function deliberately takes metadata, since plan computation must never inspect actual tensor data.""" - return ( - _compute_plans_for_group(baseline_metas), - _compute_plans_for_group(target_metas), - ) - - -def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]: - if not metas or len(metas) == 1: - return [] - - dims_str = metas[0].get("dims") - if dims_str is None: - return [] - - dim_specs = parse_dims(dims_str) - parallel_infos = [normalize_parallel_info(meta) for meta in metas] - - unsharder_plans = compute_unsharder_plan( - dim_specs=dim_specs, parallel_infos=parallel_infos - ) - reorderer_plans = compute_reorderer_plans( - dim_specs=dim_specs, parallel_infos=parallel_infos - ) - return [*unsharder_plans, *reorderer_plans] - - -def _extract_tensors( - loaded: list[ValueWithMeta], -) -> Optional[list[torch.Tensor]]: - return [value for item in loaded if isinstance(value := item.value, torch.Tensor)] - - -def _execute_plans( - tensors: list[torch.Tensor], - plans: list[Plan], -) -> Optional[torch.Tensor]: - if not tensors: - return None - - if not plans: - if len(tensors) != 1: - return None - return tensors[0] - - current = tensors - for plan in plans: - current = _execute_plan(current, plan) - - assert len(current) == 1 - return current[0] - - -def _execute_plan( - tensors: list[torch.Tensor], - plan: Plan, -) -> list[torch.Tensor]: - if isinstance(plan, UnsharderPlan): - return execute_unsharder_plan(plan, tensors) - elif isinstance(plan, ReordererPlan): - return execute_reorderer_plan(plan, tensors) - else: - raise NotImplementedError(f"Unknown {plan=}") diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/conftest.py b/test/registered/debug_utils/comparator/aligner/entrypoint/conftest.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py new file mode 100644 index 000000000..9b2688b86 --- /dev/null +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py @@ -0,0 +1,199 @@ +import sys +from typing import Optional + +import pytest +import torch + +from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import ( + AlignerResult, + _execute_step_plans, + execute_aligner_plan, + execute_sub_plan, + execute_sub_plans, +) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( + ConcatParams, + UnsharderPlan, +) +from sglang.srt.debug_utils.comparator.dims import ParallelAxis +from sglang.srt.debug_utils.comparator.utils import Pair +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="default", nightly=True) + + +class TestExecuteSubPlans: + def test_empty_tensors_returns_none(self) -> None: + result: Optional[torch.Tensor] = execute_sub_plans(tensors=[], plans=[]) + assert result is None + + def test_no_plans_single_tensor_passthrough(self) -> None: + tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0]) + result: Optional[torch.Tensor] = execute_sub_plans(tensors=[tensor], plans=[]) + assert result is not None + assert torch.equal(result, tensor) + + def test_no_plans_multiple_tensors_returns_none(self) -> None: + tensors: list[torch.Tensor] = [ + torch.tensor([1.0]), + torch.tensor([2.0]), + ] + result: Optional[torch.Tensor] = execute_sub_plans(tensors=tensors, plans=[]) + assert result is None + + def test_with_unsharder_plan(self) -> None: + t0: torch.Tensor = torch.tensor([[1.0, 2.0]]) + t1: torch.Tensor = torch.tensor([[3.0, 4.0]]) + + plan = UnsharderPlan( + axis=ParallelAxis.TP, + params=ConcatParams(dim=1), + groups=[[0, 1]], + ) + + result: Optional[torch.Tensor] = execute_sub_plans( + tensors=[t0, t1], plans=[plan] + ) + + assert result is not None + expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + assert torch.equal(result, expected) + + +class TestExecuteSubPlan: + def test_unknown_plan_type_raises(self) -> None: + class _FakePlan: + pass + + with pytest.raises(NotImplementedError, match="Unknown"): + execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan()) # type: ignore[arg-type] + + +class TestExecuteStepPlans: + def test_step_with_none_result_omitted(self) -> None: + tensors: list[torch.Tensor] = [ + torch.tensor([1.0]), + torch.tensor([2.0]), + ] + + step_plan = AlignerPerStepPlan( + step=0, + input_object_indices=[0, 1], + sub_plans=[], + ) + + result: dict[int, torch.Tensor] = _execute_step_plans( + tensors=tensors, step_plans=[step_plan] + ) + + assert result == {} + + def test_single_step_passthrough(self) -> None: + tensor: torch.Tensor = torch.tensor([1.0, 2.0]) + + step_plan = AlignerPerStepPlan( + step=5, + input_object_indices=[0], + sub_plans=[], + ) + + result: dict[int, torch.Tensor] = _execute_step_plans( + tensors=[tensor], step_plans=[step_plan] + ) + + assert 5 in result + assert torch.equal(result[5], tensor) + + +class TestExecuteAlignerPlan: + def _make_step_plan(self, *, step: int, indices: list[int]) -> AlignerPerStepPlan: + return AlignerPerStepPlan(step=step, input_object_indices=indices, sub_plans=[]) + + def test_x_side_empty_returns_failed_x(self) -> None: + plan = AlignerPlan( + per_step_plans=Pair( + x=[self._make_step_plan(step=0, indices=[0, 1])], + y=[self._make_step_plan(step=0, indices=[0])], + ), + ) + + tensors_pair: Pair[list[torch.Tensor]] = Pair( + x=[torch.tensor([1.0]), torch.tensor([2.0])], + y=[torch.tensor([3.0])], + ) + + result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + assert result.tensors is None + assert result.failed_side_xy == "x" + + def test_y_side_empty_returns_failed_y(self) -> None: + plan = AlignerPlan( + per_step_plans=Pair( + x=[self._make_step_plan(step=0, indices=[0])], + y=[self._make_step_plan(step=0, indices=[0, 1])], + ), + ) + + tensors_pair: Pair[list[torch.Tensor]] = Pair( + x=[torch.tensor([1.0])], + y=[torch.tensor([2.0]), torch.tensor([3.0])], + ) + + result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + assert result.tensors is None + assert result.failed_side_xy == "y" + + def test_single_step(self) -> None: + plan = AlignerPlan( + per_step_plans=Pair( + x=[self._make_step_plan(step=0, indices=[0])], + y=[self._make_step_plan(step=0, indices=[0])], + ), + ) + + t_x: torch.Tensor = torch.tensor([1.0, 2.0]) + t_y: torch.Tensor = torch.tensor([3.0, 4.0]) + tensors_pair: Pair[list[torch.Tensor]] = Pair(x=[t_x], y=[t_y]) + + result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + assert result.tensors is not None + assert result.failed_side_xy is None + assert torch.equal(result.tensors.x, t_x) + assert torch.equal(result.tensors.y, t_y) + + def test_success_returns_none_failed_side(self) -> None: + plan = AlignerPlan( + per_step_plans=Pair( + x=[self._make_step_plan(step=0, indices=[0])], + y=[self._make_step_plan(step=0, indices=[0])], + ), + ) + + tensors_pair: Pair[list[torch.Tensor]] = Pair( + x=[torch.tensor([10.0])], + y=[torch.tensor([20.0])], + ) + + result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + assert result.failed_side_xy is None + assert result.tensors is not None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py new file mode 100644 index 000000000..3d8f39cc3 --- /dev/null +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py @@ -0,0 +1,146 @@ +import sys +from typing import Any, Optional + +import pytest + +from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import ( + _compute_per_step_plans, + compute_aligner_plan, + compute_per_step_sub_plans, +) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepPlan, + AlignerPerStepSubPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan +from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan +from sglang.srt.debug_utils.comparator.utils import Pair +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="default", nightly=True) + + +def _make_meta( + *, + step: int = 0, + dims: Optional[str] = None, + tp_rank: int = 0, + tp_size: int = 1, + cp_rank: int = 0, + cp_size: int = 1, +) -> dict[str, Any]: + meta: dict[str, Any] = {"step": step} + if dims is not None: + meta["dims"] = dims + meta["sglang_parallel_info"] = { + "tp_rank": tp_rank, + "tp_size": tp_size, + "cp_rank": cp_rank, + "cp_size": cp_size, + } + return meta + + +class TestComputePerStepSubPlans: + def test_empty_metas(self) -> None: + result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans(metas=[]) + assert result == [] + + def test_single_meta(self) -> None: + result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( + metas=[_make_meta(dims="b h(tp)", tp_size=2)] + ) + assert result == [] + + def test_dims_none(self) -> None: + result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( + metas=[ + _make_meta(tp_rank=0, tp_size=2), + _make_meta(tp_rank=1, tp_size=2), + ] + ) + assert result == [] + + def test_tp_sharded_returns_unsharder_plan(self) -> None: + result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( + metas=[ + _make_meta(dims="b h(tp)", tp_rank=0, tp_size=2), + _make_meta(dims="b h(tp)", tp_rank=1, tp_size=2), + ] + ) + assert len(result) >= 1 + unsharder_plans: list[UnsharderPlan] = [ + p for p in result if isinstance(p, UnsharderPlan) + ] + assert len(unsharder_plans) >= 1 + + def test_zigzag_returns_both_plans(self) -> None: + result: list[AlignerPerStepSubPlan] = compute_per_step_sub_plans( + metas=[ + _make_meta(dims="b s(cp,zigzag) h", cp_rank=0, cp_size=2), + _make_meta(dims="b s(cp,zigzag) h", cp_rank=1, cp_size=2), + ] + ) + unsharder_plans: list[UnsharderPlan] = [ + p for p in result if isinstance(p, UnsharderPlan) + ] + reorderer_plans: list[ReordererPlan] = [ + p for p in result if isinstance(p, ReordererPlan) + ] + assert len(unsharder_plans) >= 1 + assert len(reorderer_plans) >= 1 + + +class TestComputePerStepPlans: + def test_groups_by_step(self) -> None: + metas: list[dict[str, Any]] = [ + _make_meta(step=0, tp_rank=0, tp_size=2), + _make_meta(step=0, tp_rank=1, tp_size=2), + _make_meta(step=1, tp_rank=0, tp_size=1), + ] + result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas) + + assert len(result) == 2 + assert result[0].step == 0 + assert result[0].input_object_indices == [0, 1] + assert result[1].step == 1 + assert result[1].input_object_indices == [2] + + def test_sorted_by_step(self) -> None: + metas: list[dict[str, Any]] = [ + _make_meta(step=2), + _make_meta(step=0), + _make_meta(step=1), + ] + result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas) + + steps: list[int] = [p.step for p in result] + assert steps == [0, 1, 2] + + def test_single_meta_per_step_empty_sub_plans(self) -> None: + metas: list[dict[str, Any]] = [ + _make_meta(step=0), + _make_meta(step=1), + ] + result: list[AlignerPerStepPlan] = _compute_per_step_plans(metas=metas) + + assert len(result) == 2 + assert all(plan.sub_plans == [] for plan in result) + + +class TestComputeAlignerPlan: + def test_wraps_both_sides(self) -> None: + metas_x: list[dict[str, Any]] = [_make_meta(step=0)] + metas_y: list[dict[str, Any]] = [_make_meta(step=0)] + + plan: AlignerPlan = compute_aligner_plan( + metas_pair=Pair(x=metas_x, y=metas_y), + ) + + assert len(plan.per_step_plans.x) == 1 + assert len(plan.per_step_plans.y) == 1 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_bundle_matcher.py b/test/registered/debug_utils/comparator/test_bundle_matcher.py new file mode 100644 index 000000000..f34b9a63b --- /dev/null +++ b/test/registered/debug_utils/comparator/test_bundle_matcher.py @@ -0,0 +1,150 @@ +import sys +from typing import Any + +import polars as pl +import pytest + +from sglang.srt.debug_utils.comparator.bundle_matcher import ( + TensorBundleInfo, + TensorFileInfo, + _rows_to_tensor_infos, + match_bundles, +) +from sglang.srt.debug_utils.comparator.utils import Pair +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="default", nightly=True) + + +def _make_row( + *, name: str, step: int = 0, rank: int = 0, filename: str | None = None +) -> dict[str, Any]: + if filename is None: + filename = f"name={name}___step={step}___rank={rank}.pt" + return {"name": name, "step": step, "rank": rank, "filename": filename} + + +def _make_df(rows: list[dict[str, Any]]) -> pl.DataFrame: + return pl.DataFrame(rows) + + +class TestMatchBundles: + def test_single_tensor_single_step(self) -> None: + target_df: pl.DataFrame = _make_df([_make_row(name="t_a")]) + baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys={"filename"}, + ) + + assert len(results) == 1 + assert len(results[0].x) == 1 + assert len(results[0].y) == 1 + assert results[0].y[0].name == "t_a" + + def test_multiple_names_separate_bundles(self) -> None: + target_df: pl.DataFrame = _make_df([ + _make_row(name="t_a"), + _make_row(name="t_b"), + ]) + baseline_df: pl.DataFrame = _make_df([ + _make_row(name="t_a"), + _make_row(name="t_b"), + ]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys={"filename"}, + ) + + assert len(results) == 2 + result_names: list[str] = [r.y[0].name for r in results] + assert "t_a" in result_names + assert "t_b" in result_names + + def test_skip_rank_groups_across_ranks(self) -> None: + target_df: pl.DataFrame = _make_df([ + _make_row(name="t_a", rank=0), + _make_row(name="t_a", rank=1), + ]) + baseline_df: pl.DataFrame = _make_df([ + _make_row(name="t_a", rank=0), + _make_row(name="t_a", rank=1), + ]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys={"filename", "rank"}, + ) + + assert len(results) == 1 + assert len(results[0].y) == 2 + + def test_baseline_missing_tensor(self) -> None: + target_df: pl.DataFrame = _make_df([ + _make_row(name="t_a"), + _make_row(name="t_extra"), + ]) + baseline_df: pl.DataFrame = _make_df([ + _make_row(name="t_a"), + ]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys={"filename"}, + ) + + assert len(results) == 2 + extra_pair: Pair[TensorBundleInfo] = [ + r for r in results if r.y[0].name == "t_extra" + ][0] + assert extra_pair.x == [] + + def test_empty_target_returns_empty(self) -> None: + target_df: pl.DataFrame = _make_df([]) + baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys={"filename"}, + ) + + assert results == [] + + def test_skip_step_groups_across_steps(self) -> None: + target_df: pl.DataFrame = _make_df([ + _make_row(name="t_a", step=0), + _make_row(name="t_a", step=1), + ]) + baseline_df: pl.DataFrame = _make_df([ + _make_row(name="t_a", step=0), + _make_row(name="t_a", step=1), + ]) + + results: list[Pair[TensorBundleInfo]] = match_bundles( + dfs=Pair(x=baseline_df, y=target_df), + skip_keys={"filename", "step"}, + ) + + assert len(results) == 1 + assert len(results[0].y) == 2 + + +class TestRowsToTensorInfos: + def test_filters_extra_columns(self) -> None: + rows: list[dict[str, Any]] = [ + {"filename": "a.pt", "name": "t_a", "step": 0, "rank": 7} + ] + infos: list[TensorFileInfo] = _rows_to_tensor_infos(rows) + + assert len(infos) == 1 + assert infos[0] == TensorFileInfo(filename="a.pt", name="t_a", step=0) + + def test_empty_rows(self) -> None: + infos: list[TensorFileInfo] = _rows_to_tensor_infos([]) + assert infos == [] + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))