Add aligner entrypoint and bundle handler in dump comparator (#19375)

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=}")

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]

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]]

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)
+    ]

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
+    ]

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,
     )
 
 

python/sglang/srt/debug_utils/comparator/pipeline.py

@@ -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=}")