chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support dims annotation and enhance dump loader in dumper (#19276)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-25 09:41:48 +08:00
committed by GitHub
parent 8b1ab4aaf9
commit 02ca107b2c
9 changed files with 437 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
python/sglang/srt/debug_utils/comparator/dims.py Normal file
View File

@@ -0,0 +1,80 @@
import re
from dataclasses import dataclass
from enum import Enum
from typing import Optional
class ParallelAxis(Enum):
TP = "tp"
CP = "cp"
EP = "ep"
SP = "sp"
class Ordering(Enum):
ZIGZAG = "zigzag"
NATURAL = "natural"
class Reduction(Enum):
PARTIAL = "partial"
@dataclass(frozen=True)
class DimSpec:
name: str
parallel: Optional[ParallelAxis] = None
ordering: Optional[Ordering] = None
reduction: Optional[Reduction] = None
_DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\((?P<modifiers>[^)]+)\))?$")
_MODIFIER_FIELDS: list[tuple[type[Enum], str]] = [
(ParallelAxis, "parallel"),
(Ordering, "ordering"),
(Reduction, "reduction"),
]
_MODIFIER_LOOKUP: dict[str, tuple[str, Enum]] = {}
for _enum_cls, _field in _MODIFIER_FIELDS:
for _member in _enum_cls:
_MODIFIER_LOOKUP[_member.value] = (_field, _member)
def parse_dim(token: str) -> DimSpec:
match = _DIM_PATTERN.match(token)
if match is None:
raise ValueError(f"Invalid dim token: {token!r}")
name = match.group("name")
modifiers_str = match.group("modifiers")
if modifiers_str is None:
return DimSpec(name=name)
fields: dict[str, Enum] = {}
for part in (p.strip() for p in modifiers_str.split(",")):
if part not in _MODIFIER_LOOKUP:
raise ValueError(f"Unknown modifier {part!r} in dim spec: {token!r}")
field_name, enum_value = _MODIFIER_LOOKUP[part]
if field_name in fields:
raise ValueError(f"Multiple {field_name} values in dim token: {token!r}")
fields[field_name] = enum_value
return DimSpec(name=name, **fields)
def parse_dims(dims_str: str) -> list[DimSpec]:
"""Parse 'b s(cp,zigzag) h(tp) d' -> list[DimSpec]."""
if not dims_str.strip():
raise ValueError("dims string must not be empty")
result = [parse_dim(token) for token in dims_str.strip().split()]
names = [spec.name for spec in result]
if len(names) != len(set(names)):
duplicates = sorted({n for n in names if names.count(n) > 1})
raise ValueError(f"Duplicate dim names: {duplicates}")
return result

16
python/sglang/srt/debug_utils/comparator/entrypoint.py
View File

@@ -1,7 +1,9 @@
import argparse
from pathlib import Path
from typing import Optional
import polars as pl
import torch
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
@@ -11,8 +13,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
print_record,
)
from sglang.srt.debug_utils.comparator.tensor_comparison import compare_tensors
from sglang.srt.debug_utils.comparator.utils import load_object
from sglang.srt.debug_utils.dump_loader import find_row, read_meta
from sglang.srt.debug_utils.dump_loader import ValueWithMeta, find_row, read_meta
def main() -> None:
@@ -70,8 +71,8 @@ def run(args: argparse.Namespace) -> None:
path_baseline = Path(args.baseline_path) / row_baseline["filename"]
x_baseline = load_object(path_baseline)
x_target = load_object(path_target)
x_baseline = _load_tensor(path_baseline)
x_target = _load_tensor(path_target)
if x_baseline is None or x_target is None:
counts["skipped"] += 1
@@ -104,6 +105,13 @@ def run(args: argparse.Namespace) -> None:
)
def _load_tensor(path: Path) -> Optional[torch.Tensor]:
loaded = ValueWithMeta.load(path)
if not isinstance(loaded.value, torch.Tensor):
return None
return loaded.value
def _parse_args() -> argparse.Namespace:
# python -m sglang.srt.debug_utils.comparator --baseline-path ... --target-path ...
parser = argparse.ArgumentParser()

17
python/sglang/srt/debug_utils/comparator/utils.py
View File

@@ -1,5 +1,4 @@
import functools
from pathlib import Path
from typing import Optional, Tuple
import torch
@@ -42,19 +41,3 @@ def calc_rel_diff(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
def load_object(path: Path) -> Optional[torch.Tensor]:
try:
x = torch.load(path, weights_only=False)
except Exception as e:
print(f"Skip load {path} since error {e}")
return None
if isinstance(x, dict) and "value" in x:
x = x["value"]
if not isinstance(x, torch.Tensor):
print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
return None
return x.cuda()

93
python/sglang/srt/debug_utils/dump_loader.py
View File

@@ -1,11 +1,60 @@
import functools
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict
from typing import Any, Dict, Tuple
import polars as pl
import torch
_TYPED_FIELDS: list[tuple[str, type]] = [("rank", int)]
def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
stem = Path(path).stem
result: Dict[str, Any] = {}
for kv in stem.split("___"):
if "=" in kv:
k, v = kv.split("=", 1)
result[k] = v
for field, converter in _TYPED_FIELDS:
if field in result:
result[field] = converter(result[field])
return result
@dataclass
class ValueWithMeta:
value: Any
meta: Dict[str, Any]
@staticmethod
def load(path: Path) -> "ValueWithMeta":
path = Path(path)
meta_from_filename = parse_meta_from_filename(path)
try:
raw = torch.load(path, weights_only=False, map_location="cpu")
except Exception as e:
print(f"Skip load {path} since error {e}")
return ValueWithMeta(
value=None, meta={**meta_from_filename, "filename": path.name}
)
value, meta_from_embedded = _unwrap_dict_format(raw)
return ValueWithMeta(
value=value,
meta={**meta_from_filename, **meta_from_embedded, "filename": path.name},
)
def _unwrap_dict_format(obj: Any) -> Tuple[Any, Dict[str, Any]]:
if isinstance(obj, dict) and "value" in obj:
meta = obj.get("meta", {})
assert isinstance(meta, dict), f"Expected meta to be dict, got {type(meta)}"
return obj["value"], meta
return obj, {}
class DumpLoader:
def __init__(self):
@@ -50,10 +99,7 @@ def read_meta(directory):
rows = []
for p in directory.glob("*.pt"):
try:
full_kwargs = {}
for kv in p.stem.split("___"):
k, v = kv.split("=")
full_kwargs[k] = v
full_kwargs = parse_meta_from_filename(p)
rows.append(
{
"filename": str(p.name),
@@ -83,26 +129,27 @@ def _add_duplicate_index(df: pl.DataFrame) -> pl.DataFrame:
return df
def find_row(df, conditions: Dict[str, Any]):
df_sub = df.filter(
functools.reduce(
lambda a, b: a & b,
[
(
pl.col(col)
== _cast_to_polars_dtype(conditions[col], df.schema[col])
if conditions[col] is not None
else pl.col(col).is_null()
)
for col in conditions.keys()
if col in df.columns
],
def filter_rows(df: pl.DataFrame, conditions: Dict[str, Any]) -> list[dict]:
filter_exprs = [
(
pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
if conditions[col] is not None
else pl.col(col).is_null()
)
)
if len(df_sub) > 1:
print(f"find_row find ambiguous results: {df_sub=}")
for col in conditions
if col in df.columns
]
if not filter_exprs:
return []
return df.filter(functools.reduce(lambda a, b: a & b, filter_exprs)).to_dicts()
def find_row(df: pl.DataFrame, conditions: Dict[str, Any]):
rows = filter_rows(df, conditions)
if len(rows) > 1:
print(f"find_row find ambiguous results: {rows=}")
return None
return df_sub.to_dicts()[0] if len(df_sub) > 0 else None
return rows[0] if rows else None
def _cast_to_polars_dtype(value, target_dtype):

36
python/sglang/srt/debug_utils/dumper.py
View File

@@ -223,7 +223,24 @@ class _Dumper:
self._state.step += 1
print(f"[Dumper] [{time.time()}] step={self._state.step}")
def dump(self, name: str, value, save: bool = True, **kwargs) -> None:
def dump(
self,
name: str,
value,
save: bool = True,
dims: Optional[str] = None,
dims_grad: Optional[str] = None,
**kwargs,
) -> None:
value_meta: dict = {}
grad_meta: dict = {}
if dims is not None:
value_meta["dims"] = dims
grad_meta["dims"] = dims
if dims_grad is not None:
value_meta["dims_grad"] = dims_grad
grad_meta["dims"] = dims_grad
self._dump_inner(
name=name,
value=value,
@@ -234,6 +251,8 @@ class _Dumper:
enable_future_grad=self._config.enable_grad,
value_tag="Dumper.Value",
grad_tag="Dumper.Grad",
value_meta_only_fields=value_meta,
grad_meta_only_fields=grad_meta,
)
def dump_model(
@@ -336,6 +355,8 @@ class _Dumper:
enable_future_grad: bool,
value_tag: str,
grad_tag: str,
value_meta_only_fields: Optional[dict] = None,
grad_meta_only_fields: Optional[dict] = None,
) -> None:
self._http_manager # noqa: B018
@@ -359,6 +380,7 @@ class _Dumper:
tags=tags,
value=value,
save=save,
meta_only_fields=value_meta_only_fields or {},
)
if (
@@ -371,6 +393,7 @@ class _Dumper:
tags={**tags, "name": f"grad__{name}"},
value=g,
save=save,
meta_only_fields=grad_meta_only_fields or {},
)
if enable_future_grad:
@@ -379,6 +402,7 @@ class _Dumper:
tensor=value,
extra_kwargs=extra_kwargs,
save=save,
meta_only_fields=grad_meta_only_fields or {},
)
def _register_dump_grad_hook(
@@ -388,6 +412,7 @@ class _Dumper:
tensor,
extra_kwargs: dict,
save: bool,
meta_only_fields: Optional[dict] = None,
) -> None:
if not isinstance(tensor, torch.Tensor):
return
@@ -396,6 +421,7 @@ class _Dumper:
captured_step = self._state.step
captured_tags = dict(name=f"grad__{name}", **deepcopy(extra_kwargs))
captured_meta_only = meta_only_fields or {}
def grad_hook(grad: torch.Tensor) -> None:
self._dump_single(
@@ -404,6 +430,7 @@ class _Dumper:
value=grad,
save=save,
step=captured_step,
meta_only_fields=captured_meta_only,
)
tensor.register_hook(grad_hook)
@@ -416,6 +443,7 @@ class _Dumper:
value,
save: bool,
step: Optional[int] = None,
meta_only_fields: Optional[dict] = None,
) -> None:
self._ensure_exp_name()
self._state.dump_index += 1
@@ -445,7 +473,11 @@ class _Dumper:
if save and (self._config.enable_output_file or capturing):
output_data = {
"value": value,
"meta": dict(**full_kwargs, **self._static_meta),
"meta": dict(
**full_kwargs,
**self._static_meta,
**(meta_only_fields or {}),
),
}
if capturing: