Support loading token aligner data in dump comparator (#19376)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,224 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import polars as pl
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
|
||||
execute_sub_plans,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
|
||||
compute_per_step_sub_plans,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
|
||||
AUX_NAMES,
|
||||
_AuxFrameworkPlugin,
|
||||
_plugins,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
|
||||
TokenAlignerGlobalAux,
|
||||
TokenAlignerStepAux,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
|
||||
normalize_parallel_info,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
|
||||
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
|
||||
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
|
||||
from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows
|
||||
|
||||
# re-export for existing callers
|
||||
__all__ = ["AUX_NAMES", "has_aux_tensors", "load_and_normalize_aux"]
|
||||
|
||||
|
||||
def load_and_normalize_aux(
|
||||
dump_path: Path, df: pl.DataFrame
|
||||
) -> Optional[TokenAlignerGlobalAux]:
|
||||
"""Bootstrap: load, unshard, and normalize auxiliary tensors for one side."""
|
||||
plugin: Optional[_AuxFrameworkPlugin] = _detect_plugin(df, dump_path=dump_path)
|
||||
if plugin is None:
|
||||
return None
|
||||
|
||||
available_names: set[str] = set(df["name"].unique().to_list()) & plugin.all_names
|
||||
steps: list[int] = sorted(df["step"].unique().to_list())
|
||||
assert len(steps) == 1, f"Multi-step not yet supported, got {len(steps)} steps"
|
||||
tensor_names: set[str] = available_names & plugin.tensor_names
|
||||
non_tensor_names: set[str] = available_names & plugin.non_tensor_names
|
||||
|
||||
steps_data: dict[int, dict[str, object]] = {}
|
||||
for step in steps:
|
||||
step_data = dict(
|
||||
_load_step_data(
|
||||
step=step,
|
||||
tensor_names=tensor_names,
|
||||
non_tensor_names=non_tensor_names,
|
||||
df=df,
|
||||
dump_path=dump_path,
|
||||
plugin=plugin,
|
||||
)
|
||||
)
|
||||
if step_data:
|
||||
steps_data[step] = step_data
|
||||
|
||||
layout: str = plugin.detect_layout(steps_data)
|
||||
|
||||
step_auxs: dict[int, TokenAlignerStepAux] = {
|
||||
step: plugin.compute_step_aux(step_data, layout=layout, step=step)
|
||||
for step, step_data in steps_data.items()
|
||||
}
|
||||
|
||||
return TokenAlignerGlobalAux(
|
||||
step_auxs=step_auxs, framework=plugin.name, layout=layout
|
||||
)
|
||||
|
||||
|
||||
def has_aux_tensors(df: pl.DataFrame) -> bool:
|
||||
"""Check if the DataFrame contains the minimum auxiliary tensors for alignment."""
|
||||
names: set[str] = set(df["name"].unique().to_list())
|
||||
return any(plugin.has_required_names(names) for plugin in _plugins)
|
||||
|
||||
|
||||
def _detect_plugin(df: pl.DataFrame, dump_path: Path) -> Optional[_AuxFrameworkPlugin]:
|
||||
names: set[str] = set(df["name"].unique().to_list())
|
||||
|
||||
for plugin in _plugins:
|
||||
if names & plugin.discriminating_names:
|
||||
return plugin
|
||||
|
||||
first_row: dict = df.row(0, named=True)
|
||||
value: ValueWithMeta = ValueWithMeta.load(dump_path / first_row["filename"])
|
||||
|
||||
for plugin in _plugins:
|
||||
if f"{plugin.name}_parallel_info" in value.meta:
|
||||
return plugin
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _load_step_data(
|
||||
*,
|
||||
step: int,
|
||||
tensor_names: set[str],
|
||||
non_tensor_names: set[str],
|
||||
df: pl.DataFrame,
|
||||
dump_path: Path,
|
||||
plugin: _AuxFrameworkPlugin,
|
||||
) -> Iterable[Tuple[str, object]]:
|
||||
"""Load all tensor and non-tensor aux values for a single step."""
|
||||
for name in non_tensor_names:
|
||||
value = _load_non_tensor_aux(name=name, step=step, df=df, dump_path=dump_path)
|
||||
if value is not None:
|
||||
yield name, value
|
||||
|
||||
for name in tensor_names:
|
||||
tensor = _load_and_align_aux_tensor(
|
||||
name=name, step=step, df=df, dump_path=dump_path, plugin=plugin
|
||||
)
|
||||
if tensor is not None:
|
||||
yield name, tensor
|
||||
|
||||
|
||||
def _load_non_tensor_aux(
|
||||
*, name: str, step: int, df: pl.DataFrame, dump_path: Path
|
||||
) -> Optional[object]:
|
||||
"""Load a non-tensor auxiliary value for a step, validating consistency across ranks."""
|
||||
rows = filter_rows(df, conditions={"name": name, "step": step})
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
loaded: list[ValueWithMeta] = [
|
||||
ValueWithMeta.load(dump_path / r["filename"]) for r in rows
|
||||
]
|
||||
|
||||
if len(loaded) > 1:
|
||||
first_value = loaded[0].value
|
||||
for i, item in enumerate(loaded[1:], start=1):
|
||||
if item.value != first_value:
|
||||
warning_sink.add(
|
||||
GeneralWarning(
|
||||
category=f"{name}_mismatch",
|
||||
message=(
|
||||
f"{name} mismatch across ranks: rank 0 has {first_value}, "
|
||||
f"rank {i} has {item.value}"
|
||||
),
|
||||
)
|
||||
)
|
||||
break
|
||||
|
||||
return loaded[0].value
|
||||
|
||||
|
||||
def _load_and_align_aux_tensor(
|
||||
*,
|
||||
name: str,
|
||||
step: int,
|
||||
df: pl.DataFrame,
|
||||
dump_path: Path,
|
||||
plugin: _AuxFrameworkPlugin,
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""Load an auxiliary tensor for (name, step), align if needed."""
|
||||
rows = filter_rows(df, conditions={"name": name, "step": step})
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
loaded: list[ValueWithMeta] = [
|
||||
ValueWithMeta.load(dump_path / r["filename"]) for r in rows
|
||||
]
|
||||
|
||||
tensors: list[torch.Tensor] = [
|
||||
item.value for item in loaded if isinstance(item.value, torch.Tensor)
|
||||
]
|
||||
if not tensors:
|
||||
return None
|
||||
|
||||
if len(tensors) == 1:
|
||||
return tensors[0]
|
||||
|
||||
metas: list[dict] = [item.meta for item in loaded]
|
||||
metas = _ensure_dims_in_metas(name=name, plugin=plugin, metas=metas)
|
||||
|
||||
sub_plans = compute_per_step_sub_plans(metas=metas)
|
||||
if sub_plans:
|
||||
result = execute_sub_plans(tensors=tensors, plans=sub_plans)
|
||||
assert result is not None
|
||||
return result
|
||||
|
||||
warning_sink.add(
|
||||
GeneralWarning(
|
||||
category="aux_no_dims",
|
||||
message=(
|
||||
f"aux tensor '{name}' has {len(tensors)} ranks "
|
||||
f"but no dims metadata, using rank 0 only"
|
||||
),
|
||||
)
|
||||
)
|
||||
return tensors[0]
|
||||
|
||||
|
||||
def _ensure_dims_in_metas(
|
||||
*, name: str, plugin: _AuxFrameworkPlugin, metas: list[dict]
|
||||
) -> list[dict]:
|
||||
"""Inject inferred dims into metas if not already present.
|
||||
|
||||
Returns metas unchanged if dims is already set, or a new list with dims
|
||||
injected if inference succeeds. Raises if the tensor is CP-sharded
|
||||
(not yet supported).
|
||||
"""
|
||||
if metas[0].get("dims") is not None:
|
||||
return metas
|
||||
|
||||
parallel_infos = [normalize_parallel_info(m) for m in metas]
|
||||
has_cp: bool = any(ParallelAxis.CP in info for info in parallel_infos)
|
||||
if not has_cp:
|
||||
return metas
|
||||
|
||||
if name in plugin.cp_sharded_names:
|
||||
raise NotImplementedError(
|
||||
f"Aux tensor '{name}' is CP-sharded but reorderer does not yet support "
|
||||
f"zigzag reordering on the 't' dimension. "
|
||||
f"Pass explicit dims= at dump time or wait for t-dim zigzag support."
|
||||
)
|
||||
|
||||
return metas
|
||||
@@ -0,0 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
|
||||
PositionalSeqId,
|
||||
SeqId,
|
||||
SGLangSeqId,
|
||||
TokenAlignerStepAux,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
|
||||
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
|
||||
|
||||
_BSHD_NOT_SUPPORTED_MSG: str = (
|
||||
"BSHD layout is not currently supported. "
|
||||
"Use aux_loader BSHD→THD conversion (planned)."
|
||||
)
|
||||
|
||||
|
||||
# ── plugin ABC ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class _AuxFrameworkPlugin(ABC):
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def tensor_names(self) -> frozenset[str]: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def non_tensor_names(self) -> frozenset[str]: ...
|
||||
|
||||
@property
|
||||
def cp_sharded_names(self) -> frozenset[str]:
|
||||
return frozenset()
|
||||
|
||||
@property
|
||||
def discriminating_names(self) -> frozenset[str]:
|
||||
"""Field names unique to this framework (excluding shared names like input_ids)."""
|
||||
return frozenset()
|
||||
|
||||
@abstractmethod
|
||||
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: ...
|
||||
|
||||
@abstractmethod
|
||||
def compute_step_aux(
|
||||
self, step_data: dict[str, object], *, layout: str, step: int
|
||||
) -> TokenAlignerStepAux: ...
|
||||
|
||||
@abstractmethod
|
||||
def has_required_names(self, names: set[str]) -> bool:
|
||||
"""Whether the minimum set of aux names needed for alignment is present."""
|
||||
...
|
||||
|
||||
@property
|
||||
def all_names(self) -> frozenset[str]:
|
||||
return self.tensor_names | self.non_tensor_names
|
||||
|
||||
|
||||
# ── sglang plugin ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
class _SGLangPlugin(_AuxFrameworkPlugin):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "sglang"
|
||||
|
||||
@property
|
||||
def tensor_names(self) -> frozenset[str]:
|
||||
return frozenset({"input_ids", "positions", "seq_lens", "req_pool_indices"})
|
||||
|
||||
@property
|
||||
def non_tensor_names(self) -> frozenset[str]:
|
||||
return frozenset({"rids"})
|
||||
|
||||
@property
|
||||
def cp_sharded_names(self) -> frozenset[str]:
|
||||
return frozenset({"input_ids", "positions"})
|
||||
|
||||
@property
|
||||
def discriminating_names(self) -> frozenset[str]:
|
||||
return frozenset({"seq_lens", "positions", "req_pool_indices", "rids"})
|
||||
|
||||
def has_required_names(self, names: set[str]) -> bool:
|
||||
return "input_ids" in names and "seq_lens" in names
|
||||
|
||||
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
|
||||
return "thd"
|
||||
|
||||
def compute_step_aux(
|
||||
self, step_data: dict[str, object], *, layout: str, step: int
|
||||
) -> TokenAlignerStepAux:
|
||||
input_ids = step_data["input_ids"]
|
||||
positions = step_data["positions"]
|
||||
seq_lens = step_data["seq_lens"]
|
||||
rids_raw = step_data.get("rids")
|
||||
|
||||
assert isinstance(
|
||||
input_ids, torch.Tensor
|
||||
), f"input_ids: expected Tensor, got {type(input_ids)}"
|
||||
assert isinstance(
|
||||
positions, torch.Tensor
|
||||
), f"positions: expected Tensor, got {type(positions)}"
|
||||
assert isinstance(
|
||||
seq_lens, torch.Tensor
|
||||
), f"seq_lens: expected Tensor, got {type(seq_lens)}"
|
||||
|
||||
seq_lens_list: list[int] = seq_lens.tolist()
|
||||
num_seqs: int = len(seq_lens_list)
|
||||
|
||||
seq_ids: list[SeqId]
|
||||
if rids_raw is not None and isinstance(rids_raw, (list, tuple)):
|
||||
seq_ids = [SGLangSeqId(rid=str(r)) for r in rids_raw]
|
||||
else:
|
||||
seq_ids = [PositionalSeqId(step=step, seq_index=i) for i in range(num_seqs)]
|
||||
|
||||
return TokenAlignerStepAux(
|
||||
input_ids=input_ids.tolist(),
|
||||
positions=positions.tolist(),
|
||||
seq_lens=seq_lens_list,
|
||||
seq_ids=seq_ids,
|
||||
)
|
||||
|
||||
|
||||
# ── megatron plugin ───────────────────────────────────────────────
|
||||
|
||||
|
||||
class _MegatronPlugin(_AuxFrameworkPlugin):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "megatron"
|
||||
|
||||
@property
|
||||
def tensor_names(self) -> frozenset[str]:
|
||||
return frozenset({"input_ids", "position_ids", "cu_seqlens_q", "cu_seqlens_kv"})
|
||||
|
||||
@property
|
||||
def non_tensor_names(self) -> frozenset[str]:
|
||||
return frozenset({"qkv_format"})
|
||||
|
||||
@property
|
||||
def cp_sharded_names(self) -> frozenset[str]:
|
||||
return frozenset({"input_ids", "position_ids"})
|
||||
|
||||
@property
|
||||
def discriminating_names(self) -> frozenset[str]:
|
||||
return frozenset({"cu_seqlens_q", "cu_seqlens_kv", "qkv_format"})
|
||||
|
||||
def has_required_names(self, names: set[str]) -> bool:
|
||||
return "input_ids" in names and "cu_seqlens_q" in names
|
||||
|
||||
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
|
||||
for step_data in raw.values():
|
||||
if (qkv_format := step_data.get("qkv_format")) is not None:
|
||||
fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format)
|
||||
if "bshd" in fmt.lower():
|
||||
raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
|
||||
return "thd"
|
||||
|
||||
input_ids = step_data.get("input_ids")
|
||||
if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2:
|
||||
raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
|
||||
|
||||
warning_sink.add(
|
||||
GeneralWarning(
|
||||
category="layout_detection_fallback",
|
||||
message=(
|
||||
"Megatron layout detection: no qkv_format or 2D input_ids found, "
|
||||
"falling back to thd"
|
||||
),
|
||||
)
|
||||
)
|
||||
return "thd"
|
||||
|
||||
def compute_step_aux(
|
||||
self, step_data: dict[str, object], *, layout: str, step: int
|
||||
) -> TokenAlignerStepAux:
|
||||
input_ids: torch.Tensor = step_data["input_ids"]
|
||||
|
||||
if (cu_seqlens_q := step_data.get("cu_seqlens_q")) is not None:
|
||||
seq_lens: torch.Tensor = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
|
||||
else:
|
||||
seq_lens = torch.tensor([input_ids.shape[0]], dtype=torch.long)
|
||||
|
||||
if (position_ids := step_data.get("position_ids")) is not None:
|
||||
positions: torch.Tensor = position_ids
|
||||
else:
|
||||
positions = _infer_positions(seq_lens=seq_lens)
|
||||
|
||||
seq_lens_list: list[int] = seq_lens.tolist()
|
||||
num_seqs: int = len(seq_lens_list)
|
||||
seq_ids: list[SeqId] = [
|
||||
PositionalSeqId(step=step, seq_index=seq_index)
|
||||
for seq_index in range(num_seqs)
|
||||
]
|
||||
|
||||
return TokenAlignerStepAux(
|
||||
input_ids=input_ids.tolist(),
|
||||
positions=positions.tolist(),
|
||||
seq_lens=seq_lens_list,
|
||||
seq_ids=seq_ids,
|
||||
)
|
||||
|
||||
|
||||
# ── plugin registry ───────────────────────────────────────────────
|
||||
|
||||
_plugins: list[_AuxFrameworkPlugin] = [_SGLangPlugin(), _MegatronPlugin()]
|
||||
|
||||
AUX_NAMES: frozenset[str] = frozenset().union(*(p.all_names for p in _plugins))
|
||||
|
||||
|
||||
# ── helpers ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _infer_positions(*, seq_lens: torch.Tensor) -> torch.Tensor:
|
||||
"""Infer positions when position_ids is missing (THD only)."""
|
||||
return torch.cat([torch.arange(int(slen.item())) for slen in seq_lens])
|
||||
@@ -0,0 +1,120 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import NamedTuple, Union
|
||||
|
||||
from pydantic import model_validator
|
||||
|
||||
from sglang.srt.debug_utils.comparator.utils import (
|
||||
Pair,
|
||||
_check_equal_lengths,
|
||||
_FrozenBase,
|
||||
)
|
||||
|
||||
|
||||
class SGLangSeqId(NamedTuple):
|
||||
rid: str
|
||||
|
||||
|
||||
class PositionalSeqId(NamedTuple):
|
||||
step: int
|
||||
seq_index: int
|
||||
|
||||
|
||||
SeqId = Union[SGLangSeqId, PositionalSeqId]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TokenAlignerStepAux:
|
||||
"""Normalized auxiliary tensors for a single step (framework-agnostic)."""
|
||||
|
||||
input_ids: list[int] # [num_tokens]
|
||||
positions: list[int] # [num_tokens]
|
||||
seq_lens: list[int] # [num_seqs]
|
||||
seq_ids: list[SeqId] # [num_seqs] — sequence identity
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
_check_equal_lengths(input_ids=self.input_ids, positions=self.positions)
|
||||
_check_equal_lengths(seq_lens=self.seq_lens, seq_ids=self.seq_ids)
|
||||
|
||||
token_count: int = sum(self.seq_lens)
|
||||
if token_count != len(self.input_ids):
|
||||
raise ValueError(
|
||||
f"sum(seq_lens)={token_count} != len(input_ids)={len(self.input_ids)}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TokenAlignerGlobalAux:
|
||||
"""Auxiliary tensors for one side across all steps + side-level metadata."""
|
||||
|
||||
step_auxs: dict[int, TokenAlignerStepAux]
|
||||
framework: str # "sglang" | "megatron"
|
||||
layout: str # "thd"
|
||||
|
||||
|
||||
class TokenLocator(_FrozenBase):
|
||||
"""Locates tokens within a single-step tensor.
|
||||
|
||||
token i is at tensor[token_index_in_step[i]].
|
||||
"""
|
||||
|
||||
token_index_in_step: list[int]
|
||||
|
||||
def __add__(self, other: TokenLocator) -> TokenLocator:
|
||||
return TokenLocator(
|
||||
token_index_in_step=self.token_index_in_step + other.token_index_in_step,
|
||||
)
|
||||
|
||||
|
||||
class TokenAlignerSeqInfo(_FrozenBase):
|
||||
"""Information for a sequence, containing information to locate all the tokens inside the sequence."""
|
||||
|
||||
# All these fields are of shape (num_tokens_in_seq,)
|
||||
input_ids: list[int]
|
||||
positions: list[int]
|
||||
locator: TokenLocator
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_fields(self) -> TokenAlignerSeqInfo:
|
||||
n: int = len(self.input_ids)
|
||||
_check_equal_lengths(
|
||||
input_ids=self.input_ids,
|
||||
positions=self.positions,
|
||||
locator_token_index_in_step=self.locator.token_index_in_step,
|
||||
)
|
||||
|
||||
if self.positions != list(range(n)):
|
||||
raise ValueError(
|
||||
f"positions must be [0, 1, ..., {n - 1}], got {self.positions}"
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
def __add__(self, other: TokenAlignerSeqInfo) -> TokenAlignerSeqInfo:
|
||||
return TokenAlignerSeqInfo(
|
||||
input_ids=self.input_ids + other.input_ids,
|
||||
positions=self.positions + other.positions,
|
||||
locator=self.locator + other.locator,
|
||||
)
|
||||
|
||||
|
||||
class TokenAlignerSeqsInfo(_FrozenBase):
|
||||
"""All sequences for one side across all steps."""
|
||||
|
||||
sequences: dict[SeqId, TokenAlignerSeqInfo]
|
||||
layout: str
|
||||
|
||||
|
||||
class TokenAlignerPlan(_FrozenBase):
|
||||
"""Token alignment plan. locators.x[i] and locators.y[i] correspond to the same logical token."""
|
||||
|
||||
locators: Pair[TokenLocator]
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_fields(self) -> TokenAlignerPlan:
|
||||
_check_equal_lengths(
|
||||
locators_x_token_index_in_step=self.locators.x.token_index_in_step,
|
||||
locators_y_token_index_in_step=self.locators.y.token_index_in_step,
|
||||
)
|
||||
return self
|
||||
@@ -0,0 +1,278 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
|
||||
_detect_plugin,
|
||||
_ensure_dims_in_metas,
|
||||
_load_and_align_aux_tensor,
|
||||
_load_non_tensor_aux,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
|
||||
_MegatronPlugin,
|
||||
_SGLangPlugin,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
|
||||
from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
_sglang_plugin = _SGLangPlugin()
|
||||
_megatron_plugin = _MegatronPlugin()
|
||||
|
||||
|
||||
def _save_pt(
|
||||
dump_path: Path,
|
||||
*,
|
||||
name: str,
|
||||
step: int,
|
||||
rank: int,
|
||||
value: object,
|
||||
meta: dict | None = None,
|
||||
) -> str:
|
||||
filename: str = f"name={name}___step={step}___rank={rank}.pt"
|
||||
payload: dict = {"value": value, "meta": meta or {}}
|
||||
torch.save(payload, dump_path / filename)
|
||||
return filename
|
||||
|
||||
|
||||
def _make_df_from_filenames(filenames: list[str]) -> pl.DataFrame:
|
||||
rows: list[dict] = []
|
||||
for fn in filenames:
|
||||
parts: dict = {}
|
||||
stem: str = fn.removesuffix(".pt")
|
||||
for kv in stem.split("___"):
|
||||
if "=" in kv:
|
||||
k, v = kv.split("=", 1)
|
||||
parts[k] = v
|
||||
rows.append(
|
||||
{
|
||||
"filename": fn,
|
||||
"name": parts["name"],
|
||||
"step": int(parts["step"]),
|
||||
"rank": int(parts["rank"]),
|
||||
}
|
||||
)
|
||||
return pl.DataFrame(rows)
|
||||
|
||||
|
||||
class TestEnsureDimsInMetas:
|
||||
"""Tests for _ensure_dims_in_metas."""
|
||||
|
||||
def _make_meta(self, *, cp_size: int = 1, cp_rank: int = 0) -> dict:
|
||||
return {
|
||||
"sglang_parallel_info": {
|
||||
"tp_rank": 0,
|
||||
"tp_size": 1,
|
||||
"cp_rank": cp_rank,
|
||||
"cp_size": cp_size,
|
||||
}
|
||||
}
|
||||
|
||||
def test_no_cp_returns_metas_unchanged(self):
|
||||
"""Without CP parallelism, metas are returned as-is."""
|
||||
metas: list[dict] = [self._make_meta(cp_size=1)]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_sglang_plugin, metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
def test_dims_already_present_returns_metas_unchanged(self):
|
||||
"""If dims is already in meta, metas are returned as-is."""
|
||||
metas: list[dict] = [{**self._make_meta(cp_size=2, cp_rank=0), "dims": "t"}]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_sglang_plugin, metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
def test_cp_sharded_sglang_input_ids_raises(self):
|
||||
"""CP + input_ids in sglang raises NotImplementedError."""
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="CP-sharded"):
|
||||
_ensure_dims_in_metas(name="input_ids", plugin=_sglang_plugin, metas=metas)
|
||||
|
||||
def test_cp_sharded_sglang_positions_raises(self):
|
||||
"""CP + positions in sglang raises NotImplementedError."""
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="CP-sharded"):
|
||||
_ensure_dims_in_metas(name="positions", plugin=_sglang_plugin, metas=metas)
|
||||
|
||||
def test_cp_sharded_megatron_input_ids_raises(self):
|
||||
"""CP + input_ids in megatron raises NotImplementedError."""
|
||||
metas: list[dict] = [
|
||||
{"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
|
||||
{"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
|
||||
]
|
||||
with pytest.raises(NotImplementedError, match="CP-sharded"):
|
||||
_ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_megatron_plugin, metas=metas
|
||||
)
|
||||
|
||||
def test_cp_non_sharded_name_returns_metas_unchanged(self):
|
||||
"""CP + non-sharded tensor name (seq_lens) returns metas as-is."""
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="seq_lens", plugin=_sglang_plugin, metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
def test_unknown_plugin_returns_metas_unchanged(self):
|
||||
"""CP + plugin with empty cp_sharded_names returns metas as-is."""
|
||||
|
||||
class _DummyPlugin(_SGLangPlugin):
|
||||
@property
|
||||
def cp_sharded_names(self) -> frozenset[str]:
|
||||
return frozenset()
|
||||
|
||||
metas: list[dict] = [
|
||||
self._make_meta(cp_size=2, cp_rank=0),
|
||||
self._make_meta(cp_size=2, cp_rank=1),
|
||||
]
|
||||
result = _ensure_dims_in_metas(
|
||||
name="input_ids", plugin=_DummyPlugin(), metas=metas
|
||||
)
|
||||
assert result is metas
|
||||
|
||||
|
||||
class TestDetectPlugin:
|
||||
def test_discriminating_names_sglang(self, tmp_path: Path) -> None:
|
||||
fn: str = _save_pt(
|
||||
tmp_path, name="seq_lens", step=0, rank=0, value=torch.tensor([3])
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn])
|
||||
|
||||
result = _detect_plugin(df, dump_path=tmp_path)
|
||||
|
||||
assert result is not None
|
||||
assert result.name == "sglang"
|
||||
|
||||
def test_fallback_to_meta_based_detection(self, tmp_path: Path) -> None:
|
||||
fn: str = _save_pt(
|
||||
tmp_path,
|
||||
name="input_ids",
|
||||
step=0,
|
||||
rank=0,
|
||||
value=torch.tensor([1, 2, 3]),
|
||||
meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn])
|
||||
|
||||
result = _detect_plugin(df, dump_path=tmp_path)
|
||||
|
||||
assert result is not None
|
||||
assert result.name == "sglang"
|
||||
|
||||
def test_returns_none_no_match(self, tmp_path: Path) -> None:
|
||||
fn: str = _save_pt(
|
||||
tmp_path, name="unrelated_tensor", step=0, rank=0, value=torch.tensor([1])
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn])
|
||||
|
||||
result = _detect_plugin(df, dump_path=tmp_path)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestLoadNonTensorAux:
|
||||
def test_multi_rank_mismatch_warning(self, tmp_path: Path) -> None:
|
||||
fn0: str = _save_pt(tmp_path, name="rids", step=0, rank=0, value=["req_A"])
|
||||
fn1: str = _save_pt(tmp_path, name="rids", step=0, rank=1, value=["req_B"])
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
|
||||
|
||||
sink = WarningSink()
|
||||
with sink.context() as warnings:
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch(
|
||||
"sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
|
||||
sink,
|
||||
):
|
||||
result = _load_non_tensor_aux(
|
||||
name="rids", step=0, df=df, dump_path=tmp_path
|
||||
)
|
||||
|
||||
assert result == ["req_A"]
|
||||
assert len(warnings) == 1
|
||||
assert isinstance(warnings[0], GeneralWarning)
|
||||
assert "rids_mismatch" in warnings[0].category
|
||||
|
||||
def test_no_rows_returns_none(self, tmp_path: Path) -> None:
|
||||
df: pl.DataFrame = _make_df_from_filenames([])
|
||||
|
||||
result = _load_non_tensor_aux(name="rids", step=0, df=df, dump_path=tmp_path)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestLoadAndAlignAuxTensor:
|
||||
def test_multi_rank_no_dims_emits_warning(self, tmp_path: Path) -> None:
|
||||
fn0: str = _save_pt(
|
||||
tmp_path,
|
||||
name="input_ids",
|
||||
step=0,
|
||||
rank=0,
|
||||
value=torch.tensor([1, 2, 3]),
|
||||
meta={
|
||||
"sglang_parallel_info": {
|
||||
"tp_rank": 0,
|
||||
"tp_size": 2,
|
||||
"cp_rank": 0,
|
||||
"cp_size": 1,
|
||||
}
|
||||
},
|
||||
)
|
||||
fn1: str = _save_pt(
|
||||
tmp_path,
|
||||
name="input_ids",
|
||||
step=0,
|
||||
rank=1,
|
||||
value=torch.tensor([4, 5, 6]),
|
||||
meta={
|
||||
"sglang_parallel_info": {
|
||||
"tp_rank": 1,
|
||||
"tp_size": 2,
|
||||
"cp_rank": 0,
|
||||
"cp_size": 1,
|
||||
}
|
||||
},
|
||||
)
|
||||
df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
|
||||
|
||||
sink = WarningSink()
|
||||
with sink.context() as warnings:
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch(
|
||||
"sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
|
||||
sink,
|
||||
):
|
||||
result = _load_and_align_aux_tensor(
|
||||
name="input_ids",
|
||||
step=0,
|
||||
df=df,
|
||||
dump_path=tmp_path,
|
||||
plugin=_sglang_plugin,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert torch.equal(result, torch.tensor([1, 2, 3]))
|
||||
assert len(warnings) == 1
|
||||
assert isinstance(warnings[0], GeneralWarning)
|
||||
assert "aux_no_dims" in warnings[0].category
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,144 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
|
||||
_infer_positions,
|
||||
_MegatronPlugin,
|
||||
_SGLangPlugin,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
|
||||
PositionalSeqId,
|
||||
SGLangSeqId,
|
||||
TokenAlignerStepAux,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=15, suite="default", nightly=True)
|
||||
|
||||
_sglang_plugin = _SGLangPlugin()
|
||||
_megatron_plugin = _MegatronPlugin()
|
||||
|
||||
|
||||
class TestNormalizeSGLang:
|
||||
"""Tests for SGLang aux tensor normalization."""
|
||||
|
||||
def test_with_rids(self):
|
||||
"""SGLang tensors with rids produce string seq_ids."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30]),
|
||||
"positions": torch.tensor([0, 1, 2]),
|
||||
"seq_lens": torch.tensor([3]),
|
||||
"rids": ["A"],
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.input_ids == [10, 20, 30]
|
||||
assert result.positions == [0, 1, 2]
|
||||
assert result.seq_lens == [3]
|
||||
assert result.seq_ids == [SGLangSeqId(rid="A")]
|
||||
|
||||
def test_rids_none_fallback(self):
|
||||
"""Missing rids results in (step, index) fallback seq_ids."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20]),
|
||||
"positions": torch.tensor([0, 1]),
|
||||
"seq_lens": torch.tensor([2]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=3
|
||||
)
|
||||
assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
|
||||
|
||||
def test_multiple_seqs_with_rids(self):
|
||||
"""Multiple sequences with rids."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"positions": torch.tensor([0, 1, 2, 0, 1]),
|
||||
"seq_lens": torch.tensor([3, 2]),
|
||||
"rids": ["A", "B"],
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
|
||||
|
||||
|
||||
class TestNormalizeMegatron:
|
||||
"""Tests for Megatron aux tensor normalization."""
|
||||
|
||||
def test_cu_seqlens_to_seq_lens(self):
|
||||
"""cu_seqlens_q is converted to seq_lens via differencing."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"cu_seqlens_q": torch.tensor([0, 3, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.seq_lens == [3, 2]
|
||||
|
||||
def test_positions_inferred_thd(self):
|
||||
"""Positions inferred from seq_lens in thd layout."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"cu_seqlens_q": torch.tensor([0, 3, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.positions == [0, 1, 2, 0, 1]
|
||||
|
||||
def test_position_ids_passthrough(self):
|
||||
"""Explicit position_ids used directly instead of inference."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"position_ids": torch.tensor([5, 6, 7, 8, 9]),
|
||||
"cu_seqlens_q": torch.tensor([0, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=0
|
||||
)
|
||||
|
||||
assert result.positions == [5, 6, 7, 8, 9]
|
||||
|
||||
def test_seq_ids_are_step_index_tuples(self):
|
||||
"""Megatron seq_ids are (step, seq_index) tuples."""
|
||||
step_data: dict = {
|
||||
"input_ids": torch.tensor([10, 20, 30, 40, 50]),
|
||||
"cu_seqlens_q": torch.tensor([0, 3, 5]),
|
||||
}
|
||||
|
||||
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
|
||||
step_data, layout="thd", step=5
|
||||
)
|
||||
assert result.seq_ids == [
|
||||
PositionalSeqId(step=5, seq_index=0),
|
||||
PositionalSeqId(step=5, seq_index=1),
|
||||
]
|
||||
|
||||
|
||||
class TestInferPositions:
|
||||
"""Tests for position inference helper."""
|
||||
|
||||
def test_thd_multiple_sequences(self):
|
||||
"""thd: positions reset to 0 for each sequence."""
|
||||
result = _infer_positions(
|
||||
seq_lens=torch.tensor([2, 3]),
|
||||
)
|
||||
assert torch.equal(result, torch.tensor([0, 1, 0, 1, 2]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -6,9 +6,11 @@ import pytest
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
ComparisonRecord,
|
||||
ConfigRecord,
|
||||
GeneralWarning,
|
||||
ReplicatedMismatchWarning,
|
||||
SkipRecord,
|
||||
SummaryRecord,
|
||||
WarningRecord,
|
||||
parse_record_json,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
|
||||
@@ -104,7 +106,7 @@ class TestRecordTypes:
|
||||
"diff_threshold": 1e-3,
|
||||
"start_step": 0,
|
||||
"end_step": 100,
|
||||
},
|
||||
}
|
||||
),
|
||||
SkipRecord(name="attn", reason="no_baseline"),
|
||||
ComparisonRecord(
|
||||
@@ -115,6 +117,9 @@ class TestRecordTypes:
|
||||
shape_mismatch=False,
|
||||
),
|
||||
SummaryRecord(total=10, passed=8, failed=1, skipped=1),
|
||||
WarningRecord(
|
||||
warnings=[GeneralWarning(category="test", message="test warning")],
|
||||
),
|
||||
]:
|
||||
restored = parse_record_json(record.model_dump_json())
|
||||
assert type(restored) is type(record)
|
||||
@@ -186,6 +191,36 @@ class TestWarnings:
|
||||
assert restored_warning.baseline_index == 0
|
||||
assert restored_warning.max_abs_diff == pytest.approx(0.42)
|
||||
|
||||
def test_any_warning_discriminated_union_round_trip(self):
|
||||
"""All AnyWarning variants survive JSON round-trip via a WarningRecord."""
|
||||
all_warnings = [
|
||||
ReplicatedMismatchWarning(
|
||||
axis="tp",
|
||||
group_index=0,
|
||||
differing_index=1,
|
||||
baseline_index=0,
|
||||
max_abs_diff=0.1,
|
||||
),
|
||||
GeneralWarning(
|
||||
category="aux_tensors_missing",
|
||||
message="Aux tensors missing, skipping token alignment",
|
||||
),
|
||||
GeneralWarning(
|
||||
category="rids_mismatch",
|
||||
message="rids mismatch across ranks: rank 0 has [1,2,3], "
|
||||
"rank 1 has [4,5,6]",
|
||||
),
|
||||
]
|
||||
|
||||
record = WarningRecord(warnings=all_warnings)
|
||||
restored = parse_record_json(record.model_dump_json())
|
||||
assert isinstance(restored, WarningRecord)
|
||||
assert len(restored.warnings) == len(all_warnings)
|
||||
|
||||
for original, parsed in zip(all_warnings, restored.warnings):
|
||||
assert type(parsed) is type(original)
|
||||
assert parsed == original
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
|
||||
@@ -44,14 +44,18 @@ class TestMatchBundles:
|
||||
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"),
|
||||
])
|
||||
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),
|
||||
@@ -64,14 +68,18 @@ class TestMatchBundles:
|
||||
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),
|
||||
])
|
||||
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),
|
||||
@@ -82,13 +90,17 @@ class TestMatchBundles:
|
||||
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"),
|
||||
])
|
||||
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),
|
||||
@@ -113,14 +125,18 @@ class TestMatchBundles:
|
||||
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),
|
||||
])
|
||||
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),
|
||||
|
||||
@@ -3,6 +3,13 @@ import sys
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
|
||||
PositionalSeqId,
|
||||
TokenAlignerPlan,
|
||||
TokenAlignerSeqInfo,
|
||||
TokenAlignerStepAux,
|
||||
TokenLocator,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
ComparisonRecord,
|
||||
@@ -15,7 +22,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
|
||||
TensorInfo,
|
||||
TensorStats,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.utils import _check_equal_lengths
|
||||
from sglang.srt.debug_utils.comparator.utils import Pair, _check_equal_lengths
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
@@ -33,6 +40,109 @@ class TestCheckEqualLengths:
|
||||
_check_equal_lengths(a=[1, 2], b=[3])
|
||||
|
||||
|
||||
class TestTokenAlignerStepAux:
|
||||
def test_valid(self):
|
||||
aux = TokenAlignerStepAux(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
seq_lens=[2, 1],
|
||||
seq_ids=[
|
||||
PositionalSeqId(step=0, seq_index=0),
|
||||
PositionalSeqId(step=0, seq_index=1),
|
||||
],
|
||||
)
|
||||
assert len(aux.input_ids) == 3
|
||||
|
||||
def test_token_length_mismatch(self):
|
||||
with pytest.raises(ValueError, match="Length mismatch"):
|
||||
TokenAlignerStepAux(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1],
|
||||
seq_lens=[2, 1],
|
||||
seq_ids=[
|
||||
PositionalSeqId(step=0, seq_index=0),
|
||||
PositionalSeqId(step=0, seq_index=1),
|
||||
],
|
||||
)
|
||||
|
||||
def test_seq_length_mismatch(self):
|
||||
with pytest.raises(ValueError, match="Length mismatch"):
|
||||
TokenAlignerStepAux(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
seq_lens=[2, 1],
|
||||
seq_ids=[PositionalSeqId(step=0, seq_index=0)],
|
||||
)
|
||||
|
||||
def test_sum_seq_lens_mismatch(self):
|
||||
with pytest.raises(ValueError, match="sum\\(seq_lens\\)"):
|
||||
TokenAlignerStepAux(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
seq_lens=[1, 1],
|
||||
seq_ids=[
|
||||
PositionalSeqId(step=0, seq_index=0),
|
||||
PositionalSeqId(step=0, seq_index=1),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestTokenAlignerSeqInfo:
|
||||
def test_valid(self):
|
||||
info = TokenAlignerSeqInfo(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
locator=TokenLocator(token_index_in_step=[0, 1, 0]),
|
||||
)
|
||||
assert len(info.input_ids) == 3
|
||||
|
||||
def test_length_mismatch(self):
|
||||
with pytest.raises(ValidationError):
|
||||
TokenAlignerSeqInfo(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 1, 2],
|
||||
locator=TokenLocator(token_index_in_step=[0, 1]),
|
||||
)
|
||||
|
||||
def test_positions_not_sequential(self):
|
||||
with pytest.raises(ValidationError, match="positions must be"):
|
||||
TokenAlignerSeqInfo(
|
||||
input_ids=[10, 20, 30],
|
||||
positions=[0, 2, 1],
|
||||
locator=TokenLocator(token_index_in_step=[0, 1, 0]),
|
||||
)
|
||||
|
||||
|
||||
class TestTokenAlignerPlan:
|
||||
def test_valid(self):
|
||||
plan = TokenAlignerPlan(
|
||||
locators=Pair(
|
||||
x=TokenLocator(token_index_in_step=[0, 1, 0]),
|
||||
y=TokenLocator(token_index_in_step=[0, 0, 1]),
|
||||
),
|
||||
)
|
||||
assert len(plan.locators.x.token_index_in_step) == 3
|
||||
|
||||
def test_length_mismatch(self):
|
||||
with pytest.raises(ValidationError, match="Length mismatch"):
|
||||
TokenAlignerPlan(
|
||||
locators=Pair(
|
||||
x=TokenLocator(token_index_in_step=[0, 1]),
|
||||
y=TokenLocator(token_index_in_step=[0, 0, 1]),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestSummaryRecord:
|
||||
def test_valid(self):
|
||||
record = SummaryRecord(total=10, passed=7, failed=2, skipped=1)
|
||||
assert record.total == 10
|
||||
|
||||
def test_total_mismatch(self):
|
||||
with pytest.raises(ValidationError, match="total=10"):
|
||||
SummaryRecord(total=10, passed=5, failed=2, skipped=1)
|
||||
|
||||
|
||||
class TestAxisInfo:
|
||||
def test_valid(self):
|
||||
info = AxisInfo(axis_rank=0, axis_size=4)
|
||||
@@ -59,16 +169,6 @@ class TestAxisInfo:
|
||||
assert info.axis_rank == 3
|
||||
|
||||
|
||||
class TestSummaryRecord:
|
||||
def test_valid(self):
|
||||
record = SummaryRecord(total=10, passed=7, failed=2, skipped=1)
|
||||
assert record.total == 10
|
||||
|
||||
def test_total_mismatch(self):
|
||||
with pytest.raises(ValidationError, match="total=10"):
|
||||
SummaryRecord(total=10, passed=5, failed=2, skipped=1)
|
||||
|
||||
|
||||
def _make_tensor_info() -> TensorInfo:
|
||||
return TensorInfo(
|
||||
shape=[4, 4],
|
||||
|
||||
Reference in New Issue
Block a user