Support flattened dims in dump comparator (#19678)
This commit is contained in:
fzyzcjy
@@ -6,6 +6,8 @@ import torch
|
||||
from einops import rearrange
|
||||
|
||||
from sglang.srt.debug_utils.comparator.dims import (
|
||||
_FUSED_NAME_SEP,
|
||||
DimSpec,
|
||||
_SingletonDimUtil,
|
||||
parse_dims,
|
||||
)
|
||||
@@ -28,26 +30,19 @@ def compute_axis_aligner_plan(
|
||||
if dims_str_pair.x is None or dims_str_pair.y is None:
|
||||
return None
|
||||
|
||||
# Pair[str] after None check
|
||||
dims_pair: Pair[str] = Pair(x=dims_str_pair.x, y=dims_str_pair.y)
|
||||
specs_pair: Pair[list[DimSpec]] = dims_pair.map(lambda s: parse_dims(s).dims)
|
||||
|
||||
raw_names: Pair[list[str]] = dims_pair.map(
|
||||
lambda s: [spec.name for spec in parse_dims(s).dims]
|
||||
)
|
||||
filtered_names: Pair[list[str]] = dims_pair.map(
|
||||
lambda s: [
|
||||
spec.name for spec in _SingletonDimUtil.filter_out(parse_dims(s).dims)
|
||||
]
|
||||
)
|
||||
|
||||
target_order: Optional[list[str]] = _resolve_target_order(
|
||||
x_names=filtered_names.x, y_names=filtered_names.y
|
||||
)
|
||||
if target_order is None:
|
||||
if not _semantic_names_match(specs_pair):
|
||||
return None
|
||||
|
||||
pattern: Pair[Optional[str]] = raw_names.map(
|
||||
lambda names: _build_pattern(source=names, target=target_order)
|
||||
# Canonical dim order follows y; fused groups stay fused (flatten, not unflatten).
|
||||
canonical_order: Optional[list[str]] = _build_canonical_order(specs_pair)
|
||||
if canonical_order is None:
|
||||
return None
|
||||
|
||||
pattern: Pair[Optional[str]] = specs_pair.map(
|
||||
lambda specs: _build_side_pattern(specs=specs, canonical_order=canonical_order)
|
||||
)
|
||||
|
||||
if pattern.x is None and pattern.y is None:
|
||||
@@ -56,45 +51,126 @@ def compute_axis_aligner_plan(
|
||||
return AxisAlignerPlan(pattern=pattern)
|
||||
|
||||
|
||||
def _resolve_target_order(
|
||||
x_names: list[str], y_names: list[str]
|
||||
) -> Optional[list[str]]:
|
||||
"""Determine the canonical dim order both sides should align to.
|
||||
def _semantic_names_match(specs_pair: Pair[list[DimSpec]]) -> bool:
|
||||
"""Check that both sides share the same semantic name set (ignoring squeeze dims)."""
|
||||
names_pair: Pair[list[str]] = specs_pair.map(_expand_and_skip_squeeze)
|
||||
|
||||
Returns y_names (the target ordering) if name sets match, or None on mismatch.
|
||||
If both sides are identical, returns the shared order.
|
||||
"""
|
||||
if x_names == y_names:
|
||||
return y_names
|
||||
if set(names_pair.x) == set(names_pair.y):
|
||||
return True
|
||||
|
||||
if set(x_names) != set(y_names):
|
||||
# Local import to avoid circular dependency:
|
||||
# output_types -> aligner/entrypoint/types -> axis_aligner -> output_types
|
||||
from sglang.srt.debug_utils.comparator.output_types import ErrorLog
|
||||
# Local import to avoid circular dependency:
|
||||
# output_types -> aligner/entrypoint/types -> axis_aligner -> output_types
|
||||
from sglang.srt.debug_utils.comparator.output_types import ErrorLog
|
||||
|
||||
log_sink.add(
|
||||
ErrorLog(
|
||||
category="axis_aligner_dim_mismatch",
|
||||
message=(
|
||||
f"AxisAligner: dim name sets differ (x={x_names}, y={y_names}), "
|
||||
f"skipping axis swap"
|
||||
),
|
||||
)
|
||||
log_sink.add(
|
||||
ErrorLog(
|
||||
category="axis_aligner_dim_mismatch",
|
||||
message=(
|
||||
f"AxisAligner: dim name sets differ (x={names_pair.x}, y={names_pair.y}), "
|
||||
f"skipping axis swap"
|
||||
),
|
||||
)
|
||||
return None
|
||||
|
||||
return y_names
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _build_pattern(*, source: list[str], target: list[str]) -> Optional[str]:
|
||||
"""Build an einops rearrange pattern from source dim names to target dim names.
|
||||
def _expand_and_skip_squeeze(specs: list[DimSpec]) -> list[str]:
|
||||
"""Expand DimSpecs to flat semantic names, skipping squeeze dims."""
|
||||
return [
|
||||
name
|
||||
for spec in specs
|
||||
if not _SingletonDimUtil.is_squeeze(spec)
|
||||
for name in spec.sub_dims
|
||||
]
|
||||
|
||||
Returns None if source already matches target (no rearrange needed).
|
||||
|
||||
def _build_canonical_order(specs_pair: Pair[list[DimSpec]]) -> Optional[list[str]]:
|
||||
"""Build canonical dim order following y, preferring fused representation.
|
||||
|
||||
Each element is either a plain name (``"c"``) or a fused placeholder (``"a___b"``).
|
||||
Fused groups from *either* side are merged — the separate side must flatten.
|
||||
Squeeze dims are excluded.
|
||||
|
||||
Returns ``None`` if the two sides have overlapping but incompatible fused groups
|
||||
(e.g. x fuses ``(a*b)`` while y fuses ``(b*c)``).
|
||||
"""
|
||||
if source == target:
|
||||
# Map each sub-dim name → (placeholder, siblings) from both sides
|
||||
fused_lookup: dict[str, tuple[str, frozenset[str]]] = {}
|
||||
for spec in (*specs_pair.x, *specs_pair.y):
|
||||
if spec.is_fused:
|
||||
placeholder: str = spec.sanitized_name
|
||||
siblings: frozenset[str] = frozenset(spec.sub_dims)
|
||||
for sub_name in spec.sub_dims:
|
||||
existing: Optional[tuple[str, frozenset[str]]] = fused_lookup.get(
|
||||
sub_name
|
||||
)
|
||||
if existing is not None and existing[1] != siblings:
|
||||
from sglang.srt.debug_utils.comparator.output_types import ErrorLog
|
||||
|
||||
log_sink.add(
|
||||
ErrorLog(
|
||||
category="axis_aligner_fused_conflict",
|
||||
message=(
|
||||
f"AxisAligner: overlapping fused groups for sub-dim {sub_name!r} "
|
||||
f"({existing[0]} vs {placeholder}), skipping axis alignment"
|
||||
),
|
||||
)
|
||||
)
|
||||
return None
|
||||
fused_lookup.setdefault(sub_name, (placeholder, siblings))
|
||||
|
||||
result: list[str] = []
|
||||
consumed: set[str] = set()
|
||||
|
||||
for spec in specs_pair.y:
|
||||
if _SingletonDimUtil.is_squeeze(spec):
|
||||
continue
|
||||
|
||||
names: list[str] = spec.sub_dims
|
||||
if any(n in consumed for n in names):
|
||||
continue
|
||||
|
||||
entry: Optional[tuple[str, frozenset[str]]] = fused_lookup.get(names[0])
|
||||
if entry is not None:
|
||||
fused_placeholder, sibs = entry
|
||||
result.append(fused_placeholder)
|
||||
consumed.update(sibs)
|
||||
else:
|
||||
result.append(spec.name)
|
||||
consumed.update(names)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _build_side_pattern(
|
||||
*, specs: list[DimSpec], canonical_order: list[str]
|
||||
) -> Optional[str]:
|
||||
"""Build an einops pattern for one side to reach ``canonical_order``.
|
||||
|
||||
Fused specs become their placeholder; separate specs that belong to a fused group
|
||||
stay as individual names on the LHS and become ``(a b)`` on the RHS (einops flatten).
|
||||
Squeeze dims (``1``) appear on the LHS but are dropped from the RHS.
|
||||
"""
|
||||
source_tokens: list[str] = [spec.sanitized_name for spec in specs]
|
||||
|
||||
# Build per-side target: replace fused placeholders with ``(a b)`` only if this side
|
||||
# has the sub-dims as separate (non-fused) names in the source
|
||||
fused_placeholders: set[str] = {
|
||||
spec.sanitized_name for spec in specs if spec.is_fused
|
||||
}
|
||||
target_tokens: list[str] = [
|
||||
(
|
||||
f"({t.replace(_FUSED_NAME_SEP, ' ')})"
|
||||
if _FUSED_NAME_SEP in t and t not in fused_placeholders
|
||||
else t
|
||||
)
|
||||
for t in canonical_order
|
||||
]
|
||||
|
||||
if source_tokens == target_tokens:
|
||||
return None
|
||||
|
||||
return f"{' '.join(source)} -> {' '.join(target)}"
|
||||
return f"{' '.join(source_tokens)} -> {' '.join(target_tokens)}"
|
||||
|
||||
|
||||
# --- executor ---
|
||||
@@ -103,6 +179,9 @@ def _build_pattern(*, source: list[str], target: list[str]) -> Optional[str]:
|
||||
def execute_axis_aligner_plan(
|
||||
tensor: torch.Tensor, plan: AxisAlignerPlan, *, side: str
|
||||
) -> torch.Tensor:
|
||||
if side not in ("x", "y"):
|
||||
raise ValueError(f"side must be 'x' or 'y', got {side!r}")
|
||||
|
||||
pattern: Optional[str] = plan.pattern.x if side == "x" else plan.pattern.y
|
||||
|
||||
if pattern is not None:
|
||||
|
||||
@@ -122,7 +122,7 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
|
||||
will be mishandled. Callers should set dims explicitly for non-zigzag CP.
|
||||
"""
|
||||
if ndim == 1:
|
||||
return "t(cp:zigzag)"
|
||||
return "t[cp:zigzag]"
|
||||
raise ValueError(
|
||||
f"SGLang: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
|
||||
)
|
||||
@@ -208,9 +208,9 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
|
||||
will be mishandled. Callers should set dims explicitly for non-zigzag CP.
|
||||
"""
|
||||
if ndim == 1:
|
||||
return "t(cp:zigzag)"
|
||||
return "t[cp:zigzag]"
|
||||
if ndim == 2:
|
||||
return "b s(cp:zigzag)"
|
||||
return "b s[cp:zigzag]"
|
||||
raise ValueError(
|
||||
f"Megatron: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
|
||||
)
|
||||
|
||||
@@ -39,7 +39,9 @@ def compute_unsharder_plan(
|
||||
|
||||
# Within each dim spec, reverse modifier order: innermost shard (rightmost) unshards first.
|
||||
reversed_sharded_modifiers: list[tuple[str, ParallelModifier]] = [
|
||||
(spec.name, m) for spec in dim_specs for m in reversed(spec.parallel_modifiers)
|
||||
(spec.sanitized_name, m)
|
||||
for spec in dim_specs
|
||||
for m in reversed(spec.parallel_modifiers)
|
||||
]
|
||||
|
||||
sharded_axes_raw: set[ParallelAxis] = {
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
@@ -40,10 +42,36 @@ class ParallelModifier(_FrozenBase):
|
||||
reduction: Optional[Reduction] = None
|
||||
|
||||
|
||||
_FUSED_NAME_SEP: str = "___"
|
||||
|
||||
|
||||
class DimSpec(_FrozenBase):
|
||||
name: str
|
||||
parallel_modifiers: list[ParallelModifier] = []
|
||||
|
||||
@property
|
||||
def sub_dims(self) -> list[str]:
|
||||
"""Sub-dim names. Fused: ``["num_heads", "head_dim"]``; plain: ``["h"]``."""
|
||||
return self.name.split("*")
|
||||
|
||||
@property
|
||||
def is_fused(self) -> bool:
|
||||
return len(self.sub_dims) > 1
|
||||
|
||||
@property
|
||||
def sanitized_name(self) -> str:
|
||||
"""Name safe for PyTorch named tensors (``*`` → ``___``)."""
|
||||
if self.is_fused:
|
||||
return _FUSED_NAME_SEP.join(self.sub_dims)
|
||||
return self.name
|
||||
|
||||
|
||||
class DimsSpec(_FrozenBase):
|
||||
"""Parsed result of a full dims string like ``"b s h[tp] # dp:=moe_dp"``."""
|
||||
|
||||
dims: list[DimSpec]
|
||||
dp_group_alias: Optional[str] = None
|
||||
|
||||
|
||||
class DimsSpec(_FrozenBase):
|
||||
"""Parsed result of a full dims string like ``"b s h(tp) # dp:=moe_dp"``."""
|
||||
@@ -92,7 +120,11 @@ class _SingletonDimUtil:
|
||||
return result
|
||||
|
||||
|
||||
_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\((?P<modifiers>[^)]+)\))?$")
|
||||
_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\[(?P<modifiers>[^\]]+)\])?$")
|
||||
|
||||
_FUSED_DIM_PATTERN = re.compile(r"^\((?P<inner>[^)]+)\)(?:\[(?P<modifiers>[^\]]+)\])?$")
|
||||
|
||||
_SUB_DIM_NAME_PATTERN = re.compile(r"^[a-zA-Z_]\w*$")
|
||||
|
||||
_AXIS_LOOKUP: dict[str, ParallelAxis] = {m.value: m for m in ParallelAxis}
|
||||
_QUALIFIER_LOOKUP: dict[str, Ordering | Reduction] = {
|
||||
@@ -154,33 +186,74 @@ def parse_dim(token: str) -> DimSpec:
|
||||
if token == SQUEEZE_DIM_NAME:
|
||||
return DimSpec(name=SQUEEZE_DIM_NAME)
|
||||
|
||||
fused_match = _FUSED_DIM_PATTERN.match(token)
|
||||
if fused_match is not None:
|
||||
return _parse_fused_dim(token=token, fused_match=fused_match)
|
||||
|
||||
return _parse_single_dim(token)
|
||||
|
||||
|
||||
def _parse_single_dim(token: str) -> DimSpec:
|
||||
match = _DIM_PATTERN.match(token)
|
||||
if match is None:
|
||||
raise ValueError(f"Invalid dim token: {token!r}")
|
||||
|
||||
name: str = match.group("name")
|
||||
modifiers_str: Optional[str] = match.group("modifiers")
|
||||
modifiers: list[ParallelModifier] = _parse_modifiers(
|
||||
modifiers_str=match.group("modifiers"), dim_token=token
|
||||
)
|
||||
return DimSpec(name=name, parallel_modifiers=modifiers)
|
||||
|
||||
|
||||
def _parse_fused_dim(*, token: str, fused_match: re.Match[str]) -> DimSpec:
|
||||
inner: str = fused_match.group("inner")
|
||||
modifiers_str: Optional[str] = fused_match.group("modifiers")
|
||||
|
||||
sub_names: list[str] = [s.strip() for s in inner.split("*")]
|
||||
for sub_name in sub_names:
|
||||
if not _SUB_DIM_NAME_PATTERN.match(sub_name):
|
||||
raise ValueError(
|
||||
f"Invalid sub-dim {sub_name!r} in fused dim token: {token!r}"
|
||||
)
|
||||
|
||||
if len(sub_names) != len(set(sub_names)):
|
||||
raise ValueError(f"Duplicate sub-dim names in fused dim token: {token!r}")
|
||||
|
||||
if len(sub_names) < 2:
|
||||
raise ValueError(
|
||||
f"Fused dim must have at least 2 sub-dims, got {len(sub_names)} in: {token!r}"
|
||||
)
|
||||
|
||||
fused_name: str = "*".join(sub_names)
|
||||
modifiers: list[ParallelModifier] = _parse_modifiers(
|
||||
modifiers_str=modifiers_str, dim_token=token
|
||||
)
|
||||
return DimSpec(name=fused_name, parallel_modifiers=modifiers)
|
||||
|
||||
|
||||
def _parse_modifiers(
|
||||
*, modifiers_str: Optional[str], dim_token: str
|
||||
) -> list[ParallelModifier]:
|
||||
if modifiers_str is None:
|
||||
return DimSpec(name=name)
|
||||
return []
|
||||
|
||||
modifiers: list[ParallelModifier] = []
|
||||
seen_axes: set[ParallelAxis] = set()
|
||||
|
||||
for modifier_token in (p.strip() for p in modifiers_str.split(",")):
|
||||
modifier: ParallelModifier = _parse_modifier_token(modifier_token, token)
|
||||
modifier: ParallelModifier = _parse_modifier_token(modifier_token, dim_token)
|
||||
if modifier.axis in seen_axes:
|
||||
raise ValueError(
|
||||
f"Duplicate axis {modifier.axis.value!r} in dim spec: {token!r}"
|
||||
f"Duplicate axis {modifier.axis.value!r} in dim spec: {dim_token!r}"
|
||||
)
|
||||
seen_axes.add(modifier.axis)
|
||||
modifiers.append(modifier)
|
||||
|
||||
return DimSpec(name=name, parallel_modifiers=modifiers)
|
||||
return modifiers
|
||||
|
||||
|
||||
def parse_dims(dims_str: str) -> DimsSpec:
|
||||
"""Parse ``"b s(cp:zigzag) h(tp) d # dp:=moe_dp"`` → :class:`DimsSpec`.
|
||||
"""Parse ``"b s[cp:zigzag] h[tp] d # dp:=moe_dp"`` → :class:`DimsSpec`.
|
||||
|
||||
The shape part (before ``#``) produces :pyattr:`DimsSpec.dims`.
|
||||
The declaration part (after ``#``) is scanned for ``dp:=<group>``
|
||||
@@ -194,13 +267,15 @@ def parse_dims(dims_str: str) -> DimsSpec:
|
||||
|
||||
dims: list[DimSpec] = [parse_dim(token) for token in raw.strip().split()]
|
||||
|
||||
non_squeeze_names: list[str] = [
|
||||
spec.name for spec in dims if not _SingletonDimUtil.is_squeeze(spec)
|
||||
]
|
||||
if len(non_squeeze_names) != len(set(non_squeeze_names)):
|
||||
duplicates = sorted(
|
||||
{n for n in non_squeeze_names if non_squeeze_names.count(n) > 1}
|
||||
)
|
||||
# Collect all semantic names (expanding fused sub-dims) for duplicate detection
|
||||
semantic_names: list[str] = []
|
||||
for spec in dims:
|
||||
if _SingletonDimUtil.is_squeeze(spec):
|
||||
continue
|
||||
semantic_names.extend(spec.sub_dims)
|
||||
|
||||
if len(semantic_names) != len(set(semantic_names)):
|
||||
duplicates = sorted({n for n in semantic_names if semantic_names.count(n) > 1})
|
||||
raise ValueError(f"Duplicate dim names: {duplicates}")
|
||||
|
||||
dp_group_alias: Optional[str] = (
|
||||
@@ -212,13 +287,17 @@ def parse_dims(dims_str: str) -> DimsSpec:
|
||||
|
||||
def resolve_dim_names(dims_str: str) -> list[str]:
|
||||
"""Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...)."""
|
||||
names: list[str] = [spec.name for spec in parse_dims(dims_str).dims]
|
||||
specs: list[DimSpec] = parse_dims(dims_str).dims
|
||||
names: list[str] = [spec.sanitized_name for spec in specs]
|
||||
return _SingletonDimUtil.sanitize_names(names)
|
||||
|
||||
|
||||
def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:
|
||||
names: list[str] = [spec.name for spec in dim_specs]
|
||||
return names.index(name) if name in names else None
|
||||
"""Find index by name. Accepts both ``*``-form and ``___``-form for fused dims."""
|
||||
for i, spec in enumerate(dim_specs):
|
||||
if spec.name == name or spec.sanitized_name == name:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int:
|
||||
|
||||
Reference in New Issue
Block a user