Support simple unsharding in dumper comparator (#19277)

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

@@ -0,0 +1,32 @@
+import torch
+
+from sglang.srt.debug_utils.comparator.unshard.types import (
+    ConcatParams,
+    UnshardPlan,
+)
+
+
+def execute_unshard_plan(
+    plan: UnshardPlan,
+    tensors_by_world_rank: dict[int, torch.Tensor],
+) -> torch.Tensor:
+    ordered_tensors = [
+        tensors_by_world_rank[world_rank]
+        for world_rank in plan.world_ranks_by_axis_rank
+    ]
+
+    return _apply_unshard(plan, ordered_tensors)
+
+
+def _apply_unshard(
+    plan: UnshardPlan, ordered_tensors: list[torch.Tensor]
+) -> torch.Tensor:
+    params = plan.params
+    if isinstance(params, ConcatParams):
+        return _unshard_concat(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)

python/sglang/srt/debug_utils/comparator/unshard/parallel_info.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
+
+_PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")
+
+
+def normalize_parallel_info(meta: dict) -> dict[ParallelAxis, AxisInfo]:
+    """Extract unified parallel info from dump meta."""
+    info: Optional[dict] = None
+    for key in _PARALLEL_INFO_KEYS:
+        value = meta.get(key)
+        if isinstance(value, dict) and value:
+            if info is not None:
+                raise ValueError(
+                    f"Meta contains multiple parallel_info keys among {_PARALLEL_INFO_KEYS}"
+                )
+            info = value
+
+    if info is None:
+        return {}
+
+    result: dict[ParallelAxis, AxisInfo] = {}
+    for axis in ParallelAxis:
+        axis_rank = info.get(f"{axis.value}_rank")
+        axis_size = info.get(f"{axis.value}_size")
+        if axis_rank is not None and axis_size is not None and axis_size > 1:
+            result[axis] = AxisInfo(
+                axis_rank=axis_rank,
+                axis_size=axis_size,
+            )
+
+    return result

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

@@ -0,0 +1,80 @@
+from typing import Optional
+
+from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis
+from sglang.srt.debug_utils.comparator.unshard.types import (
+    AxisInfo,
+    ConcatParams,
+    UnshardParams,
+    UnshardPlan,
+)
+
+
+def compute_unshard_plan(
+    dim_specs: list[DimSpec],
+    parallel_infos: list[dict[ParallelAxis, AxisInfo]],
+) -> Optional[UnshardPlan]:
+    if not parallel_infos:
+        raise ValueError("parallel_infos must not be empty")
+
+    sharded_axes: dict[ParallelAxis, tuple[int, DimSpec]] = {}
+    for dim_idx, spec in enumerate(dim_specs):
+        if spec.parallel is not None:
+            sharded_axes[spec.parallel] = (dim_idx, spec)
+
+    if len(sharded_axes) > 1:
+        raise NotImplementedError(
+            f"Multi-axis unshard is not supported. "
+            f"Got {len(sharded_axes)} sharded axes: {sorted(a.value for a in sharded_axes)}"
+        )
+
+    if not sharded_axes:
+        return None
+
+    axis_name, (dim_idx, spec) = next(iter(sharded_axes.items()))
+
+    expected_size: Optional[int] = None
+    rank_to_world: dict[int, int] = {}
+
+    for world_rank, pinfo in enumerate(parallel_infos):
+        if axis_name not in pinfo:
+            continue
+
+        ainfo = pinfo[axis_name]
+
+        if expected_size is None:
+            expected_size = ainfo.axis_size
+        elif ainfo.axis_size != expected_size:
+            raise ValueError(
+                f"Inconsistent axis_size for {axis_name.value}: "
+                f"expected {expected_size}, got {ainfo.axis_size} "
+                f"at world_rank={world_rank}"
+            )
+
+        rank_to_world.setdefault(ainfo.axis_rank, world_rank)
+
+    if expected_size is None:
+        raise ValueError(f"No parallel_info found for sharded axis {axis_name.value!r}")
+
+    if set(rank_to_world.keys()) != set(range(expected_size)):
+        raise ValueError(
+            f"axis_rank coverage for {axis_name.value} is incomplete: "
+            f"got {sorted(rank_to_world.keys())}, expected 0..{expected_size - 1}"
+        )
+
+    return UnshardPlan(
+        axis=spec.parallel,
+        params=_resolve_unshard_params(spec=spec, dim_index=dim_idx),
+        world_ranks_by_axis_rank=[rank_to_world[i] for i in range(expected_size)],
+    )
+
+
+def _resolve_unshard_params(*, spec: DimSpec, dim_index: int) -> UnshardParams:
+    if spec.reduction is not None:
+        raise NotImplementedError(
+            f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"
+        )
+    if spec.ordering is not None and spec.ordering != Ordering.NATURAL:
+        raise NotImplementedError(
+            f"Unshard for ordering={spec.ordering} not yet implemented (Phase 2)"
+        )
+    return ConcatParams(dim=dim_index)

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

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.utils import _FrozenBase
+
+
+class AxisInfo(_FrozenBase):
+    axis_rank: int
+    axis_size: int
+
+
+class ConcatParams(_FrozenBase):
+    op: Literal["concat"] = "concat"
+    dim: int
+
+
+# Phase 2: add ReduceSumParams, CpZigzagParams here, then change UnshardParams to:
+#   UnshardParams = Annotated[
+#       Union[ConcatParams, ReduceSumParams, CpZigzagParams],
+#       Field(discriminator="op"),
+#   ]
+UnshardParams = ConcatParams
+
+
+class UnshardPlan(_FrozenBase):
+    axis: ParallelAxis
+    params: UnshardParams
+    world_ranks_by_axis_rank: list[int]
+
+
+# Union of all plan types. Future pipeline components (e.g. reduction,
+# reordering) will add their own plan types here.
+Plan = UnshardPlan

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

@@ -9,6 +9,10 @@ class _StrictBase(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
 
+class _FrozenBase(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+
 def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]:
     flat_idx = x.argmax()
     return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))