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

Enhance dumper comparator with tensor unifier and location finder (#12623)

Browse Source
This commit is contained in:
fzyzcjy
2025-11-14 17:34:08 +08:00
committed by GitHub
parent ace27c0c01
commit 821fb060c3
2 changed files with 167 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

182
python/sglang/srt/debug_utils/dump_comparator.py
View File

@@ -1,7 +1,11 @@
import argparse
import functools
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Dict, List, Optional
import einops
import polars as pl
import torch
@@ -25,15 +29,38 @@ def main(args):
print("df_target", df_target)
print("df_baseline", df_baseline)
location_info_of_target_pass_id = _get_location_info_of_target_pass_id()
tensor_dim_descs = _get_tensor_dim_descs()
for row in df_target.iter_rows(named=True):
path_target = Path(args.target_path) / row["filename"]
if location_info_of_target_pass_id is not None:
location_info = location_info_of_target_pass_id.get(row["forward_pass_id"])
if location_info is None:
continue
baseline_forward_pass_id = location_info.baseline_forward_pass_id
baseline_token_slice = location_info.baseline_token_slice
else:
baseline_forward_pass_id = (
row["forward_pass_id"] - args.start_id + args.baseline_start_id
)
baseline_token_slice = None
tensor_dim_desc = None
if tensor_dim_descs is not None:
tensor_dim_descs_filtered = [
desc
for desc in tensor_dim_descs
if re.search(desc["pattern"], row["filename"]) is not None
]
if tensor_dim_descs_filtered:
tensor_dim_desc = tensor_dim_descs_filtered[0]
row_baseline = find_row(
df_baseline,
conditions=dict(
forward_pass_id=row["forward_pass_id"]
- args.start_id
+ args.baseline_start_id,
forward_pass_id=baseline_forward_pass_id,
**{
k: v
for k, v in row.items()
@@ -52,21 +79,56 @@ def main(args):
path_baseline = Path(args.baseline_path) / row_baseline["filename"]
print(f"Check: target={str(path_target)} baseline={str(path_baseline)}")
check_tensor_pair(
path_baseline=path_baseline, path_target=path_target, name=row["name"]
path_baseline=path_baseline,
path_target=path_target,
diff_threshold=args.diff_threshold,
name=row["name"],
baseline_token_slice=baseline_token_slice,
tensor_dim_desc=tensor_dim_desc,
)
print()
def check_tensor_pair(path_baseline, path_target, name=""):
def _split_einops_pattern(pattern):
return re.findall(r"\([^()]*\)|\S+", pattern)
def _get_einops_dim_index(pattern: str, dim_name: str):
pattern_list = _split_einops_pattern(pattern)
return pattern_list.index(dim_name)
def check_tensor_pair(
path_baseline,
path_target,
diff_threshold: float = 1e-3,
name="",
baseline_token_slice=None,
tensor_dim_desc: Optional["TensorDimDesc"] = None,
):
x_baseline = _load_object(path_baseline)
x_target = _load_object(path_target)
print(
f"Raw "
f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
f"[{'' if x_baseline.dtype == x_target.dtype else '🟠'}dtype] {x_baseline.dtype} vs {x_target.dtype}"
)
if tensor_dim_desc is not None:
if (s := baseline_token_slice) is not None:
dim = _get_einops_dim_index(tensor_dim_desc.baseline_desc, "num_tokens")
x_baseline = x_baseline.narrow(
dim=dim, start=s.start, length=s.stop - s.start
)
x_baseline = einops.rearrange(
x_baseline,
tensor_dim_desc.baseline_desc + " -> " + tensor_dim_desc.target_desc,
)
if (f := tensor_dim_desc.baseline_cropper) is not None:
print("Apply baseline_cropper")
x_baseline = f(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)
@@ -76,19 +138,28 @@ def check_tensor_pair(path_baseline, path_target, name=""):
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 (
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)),
):
*(
[
("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(
@@ -99,17 +170,46 @@ def check_tensor_pair(path_baseline, path_target, name=""):
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)
needs_print = max_abs_diff > 1e-3
print(
"\t".join(
f"{'❌' if value > 1e-3 else '✅'} {name}={value}"
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),
@@ -118,9 +218,15 @@ def check_tensor_pair(path_baseline, path_target, name=""):
)
)
if needs_print:
print(f"x_baseline(sample)={get_truncated_value(x_baseline)}")
print(f"x_target(sample)={get_truncated_value(x_target)}")
return dict(max_abs_diff=max_abs_diff)
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):
@@ -144,25 +250,51 @@ def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
return 1 - sim
def _comparison_preprocessor(x_baseline, x_target, name):
# can insert arbitrary adhoc postprocessing logic here
return x_baseline, x_target
def _load_object(path):
x = torch.load(path, weights_only=False)
if not isinstance(x, torch.Tensor):
print(f"Skip load {path} since {type(x)=} is not a Tensor")
print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
return None
return x.cuda()
# TODO may make customization endpoints configurable via args pointing to code file
def _comparison_preprocessor(x_baseline, x_target, name):
"""Customization endpoint. Can insert arbitrary adhoc postprocessing logic here."""
return x_baseline, x_target
@dataclass
class LocationInfo:
baseline_forward_pass_id: int
baseline_token_slice: slice
def _get_location_info_of_target_pass_id() -> Dict[int, LocationInfo]:
"""Customization endpoint."""
return {}
@dataclass
class TensorDimDesc:
baseline_desc: str
target_desc: str
baseline_cropper: Optional[Callable[[torch.Tensor], torch.Tensor]]
def _get_tensor_dim_descs() -> List[TensorDimDesc]:
"""Customization endpoint."""
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-id", type=int, default=0)
parser.add_argument("--end-id", type=int, default=1000000)
parser.add_argument("--baseline-start-id", type=int, default=0)
parser.add_argument("--diff-threshold", type=float, default=1e-3)
args = parser.parse_args()
main(args)

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

@@ -72,12 +72,20 @@ def find_row(df, conditions: Dict[str, Any]):
functools.reduce(
lambda a, b: a & b,
[
pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
(
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
],
)
)
assert len(df_sub) <= 1
if len(df_sub) > 1:
print(f"find_row find ambiguous results: {df_sub=}")
return None
return df_sub.to_dicts()[0] if len(df_sub) > 0 else None