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

Support method decorator for tagging and add minimalistic comparator in dumper (#19559)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-28 18:04:54 +08:00
committed by GitHub
parent 9bf3638a25
commit 706ab9296a
6 changed files with 551 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/dump_comparator.py
+300
View File
@@ -0,0 +1,300 @@
"""Simplified dump comparator — a self-contained single-file script for comparing
two dump directories tensor-by-tensor.
For advanced features (unshard, token alignment, per-dimension annotations), see the
full ``comparator/`` package: ``python -m sglang.srt.debug_utils.comparator``.
"""
import argparse
import functools
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, List, Optional
import torch
from sglang.srt.debug_utils.dumper import get_truncated_value
def main(args):
import polars as pl
from sglang.srt.debug_utils.dump_loader import find_row, read_meta
df_target = read_meta(args.target_path)
df_target = df_target.filter(
(pl.col("step") >= args.start_step) & (pl.col("step") <= args.end_step)
)
if args.filter:
df_target = df_target.filter(pl.col("filename").str.contains(args.filter))
assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
df_baseline = read_meta(args.baseline_path)
print("df_target", df_target)
print("df_baseline", df_baseline)
tensor_dim_descs: List[TensorDimDesc] = _get_tensor_dim_descs()
for row in df_target.iter_rows(named=True):
path_target = Path(args.target_path) / row["filename"]
tensor_dim_desc: Optional[TensorDimDesc] = None
if tensor_dim_descs:
matched: list[TensorDimDesc] = [
desc
for desc in tensor_dim_descs
if re.search(desc.pattern, row["filename"]) is not None
]
if matched:
tensor_dim_desc = matched[0]
row_baseline = find_row(
df_baseline,
conditions=dict(
step=row["step"],
**{
k: v
for k, v in row.items()
if k not in ["step", "dump_index", "filename"]
},
),
)
if row_baseline is None:
print(f"Skip: target={str(path_target)} since no baseline")
x_target = _load_object(path_target)
if x_target is not None:
print(f"x_target(sample)={get_truncated_value(x_target)}")
continue
path_baseline = Path(args.baseline_path) / row_baseline["filename"]
print(
f"Check:\n"
f"target={str(path_target)} (duplicate_index={row['duplicate_index']})\n"
f"baseline={str(path_baseline)} (duplicate_index={row_baseline['duplicate_index']})"
)
check_tensor_pair(
path_baseline=path_baseline,
path_target=path_target,
diff_threshold=args.diff_threshold,
name=row["name"],
tensor_dim_desc=tensor_dim_desc,
)
print()
def check_tensor_pair(
path_baseline,
path_target,
diff_threshold: float = 1e-3,
name="",
tensor_dim_desc: Optional["TensorDimDesc"] = None,
):
x_baseline = _load_object(path_baseline)
x_target = _load_object(path_target)
if x_baseline is None or x_target is None:
print(
f"Skip comparison because of None: x_baseline={x_baseline}, x_target={x_target}"
)
return
print(
f"Raw "
f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
f"[{'' if x_baseline.dtype == x_target.dtype else '🟠'}dtype] {x_baseline.dtype} vs {x_target.dtype}"
)
if tensor_dim_desc is not None:
import einops
x_baseline = einops.rearrange(
x_baseline,
tensor_dim_desc.baseline_desc + " -> " + tensor_dim_desc.target_desc,
)
if tensor_dim_desc.baseline_cropper is not None:
print("Apply baseline_cropper")
x_baseline = tensor_dim_desc.baseline_cropper(x_baseline)
x_baseline, x_target = _comparison_preprocessor(x_baseline, x_target, name=name)
x_baseline = _try_unify_shape(x_baseline, target_shape=x_target.shape)
print(
f"After preprocessor "
f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
)
x_baseline_original_dtype = x_baseline.dtype
x_target_original_dtype = x_target.dtype
x_target = x_target.float()
x_baseline = x_baseline.float()
for name, fn in [
("mean", torch.mean),
("std", torch.std),
("min", torch.min),
("max", torch.max),
*(
[
("p1", functools.partial(torch.quantile, q=0.01)),
("p5", functools.partial(torch.quantile, q=0.05)),
("p95", functools.partial(torch.quantile, q=0.95)),
("p99", functools.partial(torch.quantile, q=0.99)),
]
if x_baseline.numel() < 10_000_000
else []
),
]:
value_baseline = fn(x_baseline).item()
value_target = fn(x_target).item()
print(
f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})"
)
if x_baseline.shape != x_target.shape:
print(f"⚠️ Shape mismatch")
return
diff_info = _compute_and_print_diff(
x_baseline=x_baseline,
x_target=x_target,
diff_threshold=diff_threshold,
)
needs_print = diff_info["max_abs_diff"] > 1e-3
if (x_baseline_original_dtype != x_target_original_dtype) and (
(
downcast_dtype := _compute_smaller_dtype(
x_baseline_original_dtype, x_target_original_dtype
)
)
is not None
):
_compute_and_print_diff(
x_baseline=x_baseline.to(downcast_dtype),
x_target=x_target.to(downcast_dtype),
diff_threshold=diff_threshold,
prefix_text=f"When downcast to {downcast_dtype}: ",
)
if needs_print:
print(f"x_baseline(sample)={get_truncated_value(x_baseline)}")
print(f"x_target(sample)={get_truncated_value(x_target)}")
def _compute_and_print_diff(
x_baseline, x_target, diff_threshold: float, prefix_text=""
):
raw_abs_diff = (x_target - x_baseline).abs()
max_abs_diff = raw_abs_diff.max().item()
mean_abs_diff = raw_abs_diff.mean().item()
rel_diff = _calc_rel_diff(x_target, x_baseline)
print(
prefix_text
+ "\t".join(
f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
for name, value in [
("rel_diff", rel_diff),
("max_abs_diff", max_abs_diff),
("mean_abs_diff", mean_abs_diff),
]
)
)
max_diff_coord = _argmax_coord(raw_abs_diff)
print(
f"max_abs_diff happens at coord={max_diff_coord} with "
f"baseline={x_baseline[max_diff_coord].item()} "
f"target={x_target[max_diff_coord].item()}"
)
return dict(max_abs_diff=max_abs_diff)
def _argmax_coord(x: torch.Tensor) -> tuple:
flat_idx = x.argmax()
return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))
def _compute_smaller_dtype(dtype_a, dtype_b):
info_dict = {
(torch.float32, torch.bfloat16): torch.bfloat16,
# ... add more ...
}
return info_dict.get((dtype_a, dtype_b)) or info_dict.get((dtype_b, dtype_a))
def _try_unify_shape(x: torch.Tensor, target_shape):
x_shape = x.shape
num_dim_to_remove = len(x_shape) - len(target_shape)
if (x_shape[num_dim_to_remove:] == target_shape) and all(
val == 1 for val in x_shape[:num_dim_to_remove]
):
out = functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
print(f"Unify shape: {x_shape} -> {out.shape} (to match {target_shape})")
return out
return x
# Copied from DeepGEMM
def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
def _load_object(path):
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()
def _comparison_preprocessor(x_baseline, x_target, name):
"""Customization endpoint. Can insert arbitrary adhoc postprocessing logic here."""
return x_baseline, x_target
@dataclass
class TensorDimDesc:
pattern: str
baseline_desc: str
target_desc: str
baseline_cropper: Optional[Callable[[torch.Tensor], torch.Tensor]] = None
def _get_tensor_dim_descs() -> List[TensorDimDesc]:
"""Customization endpoint. Return a list of TensorDimDesc to rearrange baseline
dimensions to match target layout via einops before comparison."""
return []
if __name__ == "__main__":
# python -m sglang.srt.debug_utils.dump_comparator --baseline-path ... --target-path ...
parser = argparse.ArgumentParser()
parser.add_argument("--baseline-path", type=str)
parser.add_argument("--target-path", type=str)
parser.add_argument("--start-step", type=int, default=0)
parser.add_argument("--end-step", type=int, default=1000000)
parser.add_argument("--diff-threshold", type=float, default=1e-3)
parser.add_argument(
"--filter", type=str, default=None, help="Regex to filter filenames"
)
args = parser.parse_args()
main(args)
python/sglang/srt/debug_utils/dumper.py
+33 -1
View File
@@ -134,7 +134,6 @@ class DumperConfig(_BaseConfig):
enable_model_value: bool = False
enable_model_grad: bool = False
exp_name: Optional[str] = None
enable_http_server: bool = True
cleanup_previous: bool = False
collective_timeout: int = 60
server_port: str = "-1"
@@ -293,6 +292,39 @@ class _Dumper:
k: v for k, v in (self._state.global_ctx | kwargs).items() if v is not None
}
def ctx(
self,
_extractor: Optional[Callable[..., dict]] = None,
**static_ctx: Any,
) -> Callable:
"""Decorator that sets context before calling the wrapped function and clears it after.
Two forms:
@dumper.ctx(lambda self: dict(layer_id=self.layer_id))
def forward(self, x): ...
@dumper.ctx(phase="decode")
def decode_step(self, x): ...
"""
if _extractor is not None and static_ctx:
raise ValueError("cannot mix lambda extractor with static kwargs")
if _extractor is None and not static_ctx:
raise ValueError("must provide either a lambda or static kwargs")
def decorator(fn: Callable) -> Callable:
@functools.wraps(fn)
def wrapper(*args: Any, **kwargs: Any) -> Any:
ctx_dict: dict = _extractor(args[0]) if _extractor else static_ctx
self.set_ctx(**ctx_dict)
try:
return fn(*args, **kwargs)
finally:
self.set_ctx(**{k: None for k in ctx_dict})
return wrapper
return decorator
def register_non_intrusive_dumper(
self,
model: "torch.nn.Module",
test/registered/debug_utils/comparator/test_entrypoint.py
+1 -10
View File
@@ -332,7 +332,6 @@ class TestEntrypointGroupingRaw:
enable=True,
dir=str(side_dir),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
dumper.__dict__["_static_meta"] = {"world_rank": 0, "world_size": 1}
@@ -1110,7 +1109,6 @@ class TestEntrypointAlignment:
enable=True,
dir=str(d),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
@@ -1167,7 +1165,6 @@ class TestEntrypointAlignment:
enable=True,
dir=str(sglang_dir),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
@@ -1226,7 +1223,6 @@ class TestEntrypointAlignment:
enable=True,
dir=str(megatron_dir),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
@@ -1468,9 +1464,7 @@ def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonReco
def _make_dumper(directory: Path) -> _Dumper:
return _Dumper(
config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False)
)
return _Dumper(config=DumperConfig(enable=True, dir=str(directory)))
def _create_dumps(
@@ -1528,7 +1522,6 @@ def _create_non_tensor_rank_dump(
enable=True,
dir=str(directory),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
dumper.__dict__["_static_meta"] = {"world_rank": rank, "world_size": 1}
@@ -1611,7 +1604,6 @@ def _create_rank_dump(
enable=True,
dir=str(directory),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
@@ -1951,7 +1943,6 @@ class TestEntrypointThdCpZigzag:
enable=True,
dir=str(sglang_dir),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
test/registered/debug_utils/comparator/test_model_validation.py
-74
View File
@@ -403,79 +403,5 @@ class TestAlignerPlanInComparisonRecord:
assert "unsharder" in text
def _make_aligner_plan() -> AlignerPlan:
unsharder = UnsharderPlan(
axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"),
groups=[[0, 1]],
)
return AlignerPlan(
per_step_plans=Pair(
x=[
AlignerPerStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
)
],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
)
],
),
)
class TestAlignerPlanInComparisonRecord:
def test_comparison_record_with_aligner_plan(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: ComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
assert record_with_plan.aligner_plan is not None
assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0
def test_aligner_plan_json_roundtrip(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: ComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
json_str: str = record_with_plan.model_dump_json()
parsed = json.loads(json_str)
assert "aligner_plan" in parsed
assert (
parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"]
== "unsharder"
)
roundtripped: ComparisonRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is not None
assert (
roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
== "unsharder"
)
def test_comparison_record_without_aligner_plan(self) -> None:
record: ComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
json_str: str = record.model_dump_json()
roundtripped: ComparisonRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is None
def test_aligner_plan_text_format(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: ComparisonRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
text: str = record_with_plan.to_text()
assert "Aligner Plan:" in text
assert "unsharder" in text
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/debug_utils/test_dump_comparator.py
+159
View File
@@ -0,0 +1,159 @@
from argparse import Namespace
from pathlib import Path
import pytest
import torch
from sglang.srt.debug_utils.dump_comparator import (
_argmax_coord,
_calc_rel_diff,
_compute_smaller_dtype,
_try_unify_shape,
main,
)
from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=30, suite="default", nightly=True)
# ----------------------------- Unit tests -----------------------------
class TestCalcRelDiff:
def test_identical_vectors(self) -> None:
x: torch.Tensor = torch.randn(10, 10)
assert _calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
def test_zero_vectors(self) -> None:
z: torch.Tensor = torch.zeros(5)
result = _calc_rel_diff(z, z)
assert not torch.isnan(result) or True # should not crash
class TestArgmaxCoord:
def test_known_position(self) -> None:
x: torch.Tensor = torch.zeros(2, 3, 4)
x[1, 2, 3] = 10.0
assert _argmax_coord(x) == (1, 2, 3)
class TestTryUnifyShape:
def test_squeeze_leading_ones(self) -> None:
target_shape: torch.Size = torch.Size([3, 4])
result: torch.Tensor = _try_unify_shape(torch.randn(1, 1, 3, 4), target_shape)
assert result.shape == target_shape
def test_no_op_when_no_leading_ones(self) -> None:
target_shape: torch.Size = torch.Size([3, 4])
result: torch.Tensor = _try_unify_shape(torch.randn(2, 3, 4), target_shape)
assert result.shape == (2, 3, 4)
class TestComputeSmallerDtype:
def test_known_pair(self) -> None:
assert _compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16
assert _compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16
def test_none_for_same_dtype(self) -> None:
assert _compute_smaller_dtype(torch.float32, torch.float32) is None
# ----------------------------- Integration tests -----------------------------
def _make_dumper(directory: Path) -> _Dumper:
return _Dumper(
config=DumperConfig(
enable=True,
dir=str(directory),
)
)
def _create_dumps(
tmp_path: Path,
tensor_names: list[str],
*,
baseline_names: list[str] | None = None,
) -> tuple[Path, Path]:
if baseline_names is None:
baseline_names = tensor_names
d_baseline: Path = tmp_path / "baseline"
d_target: Path = tmp_path / "target"
d_baseline.mkdir()
d_target.mkdir()
torch.manual_seed(42)
baseline_tensor: torch.Tensor = torch.randn(10, 10)
target_tensor: torch.Tensor = baseline_tensor + torch.randn(10, 10) * 0.01
exp_paths: list[Path] = []
for d, names, tensor in [
(d_baseline, baseline_names, baseline_tensor),
(d_target, tensor_names, target_tensor),
]:
dumper: _Dumper = _make_dumper(d)
for name in names:
dumper.dump(name, tensor)
dumper.step()
exp_paths.append(d / dumper._config.exp_name)
return exp_paths[0], exp_paths[1]
def _make_args(
baseline_path: Path,
target_path: Path,
*,
filter_pattern: str | None = None,
) -> Namespace:
return Namespace(
baseline_path=str(baseline_path),
target_path=str(target_path),
start_step=0,
end_step=1000000,
diff_threshold=1e-3,
filter=filter_pattern,
)
class TestMainBasic:
def test_matching_tensors(
self, tmp_path: Path, capsys: pytest.CaptureFixture
) -> None:
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args: Namespace = _make_args(baseline_path, target_path)
main(args)
captured: str = capsys.readouterr().out
assert "✅" in captured
def test_with_filter(self, tmp_path: Path, capsys: pytest.CaptureFixture) -> None:
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args: Namespace = _make_args(
baseline_path, target_path, filter_pattern="tensor_a"
)
main(args)
captured: str = capsys.readouterr().out
assert "tensor_a" in captured
assert "Check:" in captured
def test_no_match_skips(
self, tmp_path: Path, capsys: pytest.CaptureFixture
) -> None:
baseline_path, target_path = _create_dumps(
tmp_path,
["only_in_target"],
baseline_names=["only_in_baseline"],
)
args: Namespace = _make_args(baseline_path, target_path)
main(args)
captured: str = capsys.readouterr().out
assert "Skip" in captured
test/registered/debug_utils/test_dumper.py
+58 -5
View File
@@ -451,7 +451,6 @@ class TestDumperDistributed:
config=DumperConfig(
enable=True,
collective_timeout=3,
enable_http_server=False,
),
)
@@ -663,12 +662,11 @@ class TestDumpDictFormat:
def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
"""Create a _Dumper for CPU testing without HTTP server or distributed."""
"""Create a _Dumper for CPU testing without distributed."""
defaults = dict(
enable=True,
dir=str(tmp_path),
exp_name="test",
enable_http_server=False,
)
defaults.update(overrides)
config = DumperConfig(**defaults)
@@ -2293,7 +2291,6 @@ class TestDumperDims:
config=DumperConfig(
enable=True,
dir=str(tmp_path),
enable_http_server=False,
enable_grad=True,
)
)
@@ -2323,7 +2320,6 @@ class TestDumperDims:
config=DumperConfig(
enable=True,
dir=str(tmp_path),
enable_http_server=False,
enable_grad=True,
)
)
@@ -2340,5 +2336,62 @@ class TestDumperDims:
assert grad_data["meta"]["dims"] == "b h(tp)"
class TestCtxDecorator:
def test_ctx_dynamic_lambda(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path)
class FakeLayer:
def __init__(self, layer_id: int) -> None:
self.layer_id = layer_id
@d.ctx(lambda self: dict(layer_id=self.layer_id))
def forward(self, x: torch.Tensor) -> torch.Tensor:
d.dump("hidden", x)
return x
layer = FakeLayer(layer_id=42)
layer.forward(torch.randn(3))
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["layer_id=42"])
def test_ctx_static_kwargs(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path)
@d.ctx(phase="decode")
def decode_step(x: torch.Tensor) -> torch.Tensor:
d.dump("step_out", x)
return x
decode_step(torch.randn(3))
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["phase=decode"])
def test_ctx_clears_on_exception(self, tmp_path: Path) -> None:
d = _make_test_dumper(tmp_path)
@d.ctx(phase="train")
def buggy_fn() -> None:
raise RuntimeError("boom")
with pytest.raises(RuntimeError, match="boom"):
buggy_fn()
assert d._state.global_ctx == {}
def test_ctx_rejects_mixed_args(self) -> None:
d = _make_test_dumper("/tmp")
with pytest.raises(ValueError, match="cannot mix"):
d.ctx(lambda self: dict(a=1), phase="x")
def test_ctx_rejects_empty_args(self) -> None:
d = _make_test_dumper("/tmp")
with pytest.raises(ValueError, match="must provide"):
d.ctx()
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))