Support replication axis in dump comparator (#19282)

python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py

@@ -2,30 +2,80 @@ import torch
 
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
     ConcatParams,
+    PickParams,
     UnshardParams,
     UnshardPlan,
 )
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.output_types import (
+    AlignWarning,
+    ReplicatedMismatchWarning,
+)
 
 
 def execute_unshard_plan(
     plan: UnshardPlan,
     tensors: list[torch.Tensor],
-) -> list[torch.Tensor]:
+) -> tuple[list[torch.Tensor], list[AlignWarning]]:
+    all_warnings: list[AlignWarning] = []
     result: list[torch.Tensor] = []
-    for group in plan.groups:
+
+    for group_idx, group in enumerate(plan.groups):
         group_tensors = [tensors[i] for i in group]
-        result.append(_apply_unshard(plan.params, group_tensors))
-    return result
+        tensor, warnings = _apply_unshard(
+            plan.params,
+            group_tensors,
+            axis=plan.axis,
+            group_index=group_idx,
+        )
+        result.append(tensor)
+        all_warnings.extend(warnings)
+
+    return result, all_warnings
 
 
 def _apply_unshard(
-    params: UnshardParams, ordered_tensors: list[torch.Tensor]
-) -> torch.Tensor:
+    params: UnshardParams,
+    ordered_tensors: list[torch.Tensor],
+    *,
+    axis: ParallelAxis,
+    group_index: int,
+) -> tuple[torch.Tensor, list[AlignWarning]]:
+    if isinstance(params, PickParams):
+        warnings = _verify_replicated_group(
+            ordered_tensors,
+            axis=axis,
+            group_index=group_index,
+        )
+        return ordered_tensors[0], warnings
+
     if isinstance(params, ConcatParams):
-        return _unshard_concat(ordered_tensors, dim=params.dim)
+        return torch.cat(ordered_tensors, dim=params.dim), []
+
     # Phase 2: ReduceSumParams, CpZigzagParams
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
 
 
-def _unshard_concat(tensors: list[torch.Tensor], dim: int) -> torch.Tensor:
-    return torch.cat(tensors, dim=dim)
+def _verify_replicated_group(
+    ordered_tensors: list[torch.Tensor],
+    *,
+    axis: ParallelAxis,
+    group_index: int,
+) -> list[ReplicatedMismatchWarning]:
+    warnings: list[ReplicatedMismatchWarning] = []
+    baseline = ordered_tensors[0]
+
+    for i in range(1, len(ordered_tensors)):
+        other = ordered_tensors[i]
+        if not torch.allclose(baseline, other, atol=1e-6):
+            warnings.append(
+                ReplicatedMismatchWarning(
+                    axis=axis.value,
+                    group_index=group_index,
+                    differing_index=i,
+                    baseline_index=0,
+                    max_abs_diff=(baseline - other).abs().max().item(),
+                )
+            )
+
+    return warnings

python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py

@@ -4,6 +4,7 @@ from typing import NamedTuple
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
     AxisInfo,
     ConcatParams,
+    PickParams,
     UnshardParams,
     UnshardPlan,
 )
@@ -32,31 +33,38 @@ def compute_unshard_plan(
         for dim_idx, spec in enumerate(dim_specs)
         if spec.parallel is not None
     }
-    if not sharded_axis_infos:
+    sharded_axes: set[ParallelAxis] = set(sharded_axis_infos)
+
+    all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info}
+    replicated_axes: set[ParallelAxis] = all_axes - sharded_axes
+
+    if not sharded_axes and not replicated_axes:
         return []
 
-    _validate(sharded_axes=set(sharded_axis_infos), parallel_infos=parallel_infos)
+    _validate(
+        axes_to_validate=sharded_axes | replicated_axes,
+        parallel_infos=parallel_infos,
+    )
 
     current_coords: _CoordsList = [
-        {axis: info[axis].axis_rank for axis in sharded_axis_infos}
+        {axis: info[axis].axis_rank for axis in sharded_axes | replicated_axes}
         for info in parallel_infos
     ]
 
+    axis_and_params: list[tuple[ParallelAxis, UnshardParams]] = [
+        (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
+    ] + [
+        (axis, _resolve_unshard_params(spec=spec, dim_index=dim_index))
+        for axis, (dim_index, spec) in sharded_axis_infos.items()
+    ]
+
     plans: list[UnshardPlan] = []
-    for axis, (dim_index, spec) in sharded_axis_infos.items():
+    for axis, params in axis_and_params:
         result = _group_and_project(
             current_coords=current_coords,
             target_axis=axis,
         )
-
-        plans.append(
-            UnshardPlan(
-                axis=axis,
-                params=_resolve_unshard_params(spec=spec, dim_index=dim_index),
-                groups=result.groups,
-            )
-        )
-
+        plans.append(UnshardPlan(axis=axis, params=params, groups=result.groups))
         current_coords = result.projected_coords
 
     return plans
@@ -64,18 +72,18 @@ def compute_unshard_plan(
 
 def _validate(
     *,
-    sharded_axes: set[ParallelAxis],
+    axes_to_validate: set[ParallelAxis],
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
 ) -> None:
-    """Check that every rank has all sharded axes, sizes are consistent, and ranks are complete."""
+    """Check that every rank has all axes, sizes are consistent, and ranks are complete."""
     axis_sizes: dict[ParallelAxis, int] = {}
 
     for world_rank, parallel_info in enumerate(parallel_infos):
-        for axis in sharded_axes:
+        for axis in axes_to_validate:
             if axis not in parallel_info:
                 raise ValueError(
                     f"world_rank={world_rank} missing parallel_info for "
-                    f"sharded axis {axis.value!r}"
+                    f"axis {axis.value!r}"
                 )
 
             axis_info = parallel_info[axis]

python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from typing import Literal
+from typing import Annotated, Literal, Union
+
+from pydantic import Field
 
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 from sglang.srt.debug_utils.comparator.utils import _FrozenBase
@@ -16,7 +18,14 @@ class ConcatParams(_FrozenBase):
     dim: int
 
 
-UnshardParams = ConcatParams
+class PickParams(_FrozenBase):
+    op: Literal["pick"] = "pick"
+
+
+UnshardParams = Annotated[
+    Union[ConcatParams, PickParams],
+    Field(discriminator="op"),
+]
 
 
 class UnshardPlan(_FrozenBase):

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

@@ -1,7 +1,7 @@
 from abc import abstractmethod
 from typing import Annotated, Literal, Union
 
-from pydantic import Discriminator, TypeAdapter
+from pydantic import Discriminator, Field, TypeAdapter
 
 from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
     format_comparison,
@@ -12,9 +12,38 @@ from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
 from sglang.srt.debug_utils.comparator.utils import _StrictBase
 
 
+class ReplicatedMismatchWarning(_StrictBase):
+    kind: Literal["replicated_mismatch"] = "replicated_mismatch"
+    axis: str
+    group_index: int
+    differing_index: int
+    baseline_index: int
+    max_abs_diff: float
+
+    def to_text(self) -> str:
+        return (
+            f"Replicated along {self.axis}: group {self.group_index}, "
+            f"index {self.differing_index} differs from {self.baseline_index} "
+            f"(max_abs_diff={self.max_abs_diff:.6e})"
+        )
+
+
+AlignWarning = (
+    ReplicatedMismatchWarning  # future: Annotated[Union[...], Discriminator("kind")]
+)
+
+
 class _OutputRecord(_StrictBase):
+    align_warnings: list[AlignWarning] = Field(default_factory=list)
+
     @abstractmethod
-    def to_text(self) -> str: ...
+    def _format_body(self) -> str: ...
+
+    def to_text(self) -> str:
+        body = self._format_body()
+        if self.align_warnings:
+            body += "\n" + "\n".join(f"  ⚠ {w.to_text()}" for w in self.align_warnings)
+        return body
 
 
 class ConfigRecord(_OutputRecord):
@@ -25,7 +54,7 @@ class ConfigRecord(_OutputRecord):
     start_step: int
     end_step: int
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return (
             f"Config: baseline={self.baseline_path} target={self.target_path}\n"
             f"diff_threshold={self.diff_threshold} "
@@ -39,10 +68,12 @@ class SkipRecord(_OutputRecord):
     reason: str
 
     @property
-    def category(self):
+    def category(self) -> str:
+        if self.align_warnings:
+            return "failed"
         return "skipped"
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return f"Skip: {self.name} ({self.reason})"
 
 
@@ -50,10 +81,12 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
     type: Literal["comparison"] = "comparison"
 
     @property
-    def category(self):
+    def category(self) -> str:
+        if self.align_warnings:
+            return "failed"
         return "passed" if self.diff is not None and self.diff.passed else "failed"
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return format_comparison(self)
 
 
@@ -64,7 +97,7 @@ class SummaryRecord(_OutputRecord):
     failed: int
     skipped: int
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return (
             f"Summary: {self.passed} passed, {self.failed} failed, "
             f"{self.skipped} skipped (total {self.total})"

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

@@ -20,6 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
 from sglang.srt.debug_utils.comparator.dims import parse_dims
 from sglang.srt.debug_utils.comparator.output_types import (
+    AlignWarning,
     ComparisonRecord,
     SkipRecord,
 )
@@ -50,12 +51,13 @@ def process_tensor_group(
     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)
+    b_tensor, b_warns = _execute_plans(b_extracted, b_plans)
+    t_tensor, t_warns = _execute_plans(t_extracted, t_plans)
+    all_warnings: list[AlignWarning] = b_warns + t_warns
 
     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)
+        return SkipRecord(name=name, reason=reason, align_warnings=all_warnings)
 
     info = compare_tensors(
         x_baseline=b_tensor,
@@ -64,7 +66,7 @@ def process_tensor_group(
         diff_threshold=diff_threshold,
     )
 
-    return ComparisonRecord(**info.model_dump())
+    return ComparisonRecord(**info.model_dump(), align_warnings=all_warnings)
 
 
 def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
@@ -112,27 +114,32 @@ def _extract_tensors(
 def _execute_plans(
     tensors: list[torch.Tensor],
     plans: list[Plan],
-) -> Optional[torch.Tensor]:
+) -> tuple[Optional[torch.Tensor], list[AlignWarning]]:
     if not tensors:
-        return None
+        return None, []
 
     if not plans:
         if len(tensors) != 1:
-            return None
-        return tensors[0]
+            return None, []
+        return tensors[0], []
 
+    warnings: list[AlignWarning] = []
     current = tensors
     for plan in plans:
-        current = _execute_plan(current, plan)
+        current, new_warnings = _execute_plan(current, plan)
+        warnings.extend(new_warnings)
 
     assert len(current) == 1
-    return current[0]
+    return current[0], warnings
 
 
-def _execute_plan(tensors, plan):
+def _execute_plan(
+    tensors: list[torch.Tensor],
+    plan: Plan,
+) -> tuple[list[torch.Tensor], list[AlignWarning]]:
     if isinstance(plan, UnshardPlan):
         return execute_unshard_plan(plan, tensors)
     elif isinstance(plan, ReorderPlan):
-        return execute_reorder_plan(plan, tensors)
+        return execute_reorder_plan(plan, tensors), []
     else:
         raise NotImplementedError(f"Unknown {plan=}")