Implement simplest dump comparator v2 (#19274)
This commit is contained in:
fzyzcjy
4
python/sglang/srt/debug_utils/comparator/__main__.py
Normal file
4
python/sglang/srt/debug_utils/comparator/__main__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
from sglang.srt.debug_utils.comparator.entrypoint import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
93
python/sglang/srt/debug_utils/comparator/entrypoint.py
Normal file
93
python/sglang/srt/debug_utils/comparator/entrypoint.py
Normal file
@@ -0,0 +1,93 @@
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison import (
|
||||
compare_tensors,
|
||||
print_comparison,
|
||||
)
|
||||
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.dumper import get_truncated_value
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
run(args)
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> None:
|
||||
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)
|
||||
|
||||
for row in df_target.iter_rows(named=True):
|
||||
path_target = Path(args.target_path) / row["filename"]
|
||||
baseline_step = row["step"]
|
||||
|
||||
row_baseline = find_row(
|
||||
df_baseline,
|
||||
conditions=dict(
|
||||
step=baseline_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']})"
|
||||
)
|
||||
|
||||
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: "
|
||||
f"x_baseline={x_baseline}, x_target={x_target}"
|
||||
)
|
||||
continue
|
||||
|
||||
info = compare_tensors(
|
||||
x_baseline=x_baseline,
|
||||
x_target=x_target,
|
||||
name=row["name"],
|
||||
)
|
||||
print_comparison(info=info, diff_threshold=args.diff_threshold)
|
||||
print()
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
# python -m sglang.srt.debug_utils.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"
|
||||
)
|
||||
return parser.parse_args()
|
||||
@@ -0,0 +1,4 @@
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.printer import (
|
||||
print_comparison,
|
||||
)
|
||||
@@ -0,0 +1,115 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
|
||||
DiffInfo,
|
||||
TensorComparisonInfo,
|
||||
TensorInfo,
|
||||
TensorStats,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.utils import (
|
||||
argmax_coord,
|
||||
calc_rel_diff,
|
||||
compute_smaller_dtype,
|
||||
try_unify_shape,
|
||||
)
|
||||
from sglang.srt.debug_utils.dumper import get_truncated_value
|
||||
|
||||
QUANTILE_NUMEL_THRESHOLD = 10_000_000
|
||||
SAMPLE_DIFF_THRESHOLD = 1e-3
|
||||
|
||||
|
||||
def compare_tensors(
|
||||
x_baseline: torch.Tensor,
|
||||
x_target: torch.Tensor,
|
||||
name: str = "",
|
||||
) -> TensorComparisonInfo:
|
||||
baseline_info = TensorInfo(
|
||||
shape=x_baseline.shape,
|
||||
dtype=x_baseline.dtype,
|
||||
stats=_compute_tensor_stats(x_baseline.float()),
|
||||
sample=None,
|
||||
)
|
||||
target_info = TensorInfo(
|
||||
shape=x_target.shape,
|
||||
dtype=x_target.dtype,
|
||||
stats=_compute_tensor_stats(x_target.float()),
|
||||
sample=None,
|
||||
)
|
||||
|
||||
x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape)
|
||||
unified_shape = x_baseline.shape
|
||||
|
||||
baseline_original_dtype = x_baseline.dtype
|
||||
target_original_dtype = x_target.dtype
|
||||
|
||||
x_baseline_f = x_baseline.float()
|
||||
x_target_f = x_target.float()
|
||||
|
||||
shape_mismatch = x_baseline_f.shape != x_target_f.shape
|
||||
|
||||
diff: Optional[DiffInfo] = None
|
||||
diff_downcast: Optional[DiffInfo] = None
|
||||
downcast_dtype: Optional[torch.dtype] = None
|
||||
|
||||
if not shape_mismatch:
|
||||
diff = _compute_diff(x_baseline=x_baseline_f, x_target=x_target_f)
|
||||
|
||||
needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD
|
||||
if needs_sample:
|
||||
baseline_info.sample = get_truncated_value(x_baseline_f)
|
||||
target_info.sample = get_truncated_value(x_target_f)
|
||||
|
||||
if baseline_original_dtype != target_original_dtype:
|
||||
downcast_dtype = compute_smaller_dtype(
|
||||
baseline_original_dtype, target_original_dtype
|
||||
)
|
||||
if downcast_dtype is not None:
|
||||
diff_downcast = _compute_diff(
|
||||
x_baseline=x_baseline_f.to(downcast_dtype),
|
||||
x_target=x_target_f.to(downcast_dtype),
|
||||
)
|
||||
|
||||
return TensorComparisonInfo(
|
||||
name=name,
|
||||
baseline=baseline_info,
|
||||
target=target_info,
|
||||
unified_shape=unified_shape,
|
||||
shape_mismatch=shape_mismatch,
|
||||
diff=diff,
|
||||
diff_downcast=diff_downcast,
|
||||
downcast_dtype=downcast_dtype,
|
||||
)
|
||||
|
||||
|
||||
def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
|
||||
include_quantiles = x.numel() < QUANTILE_NUMEL_THRESHOLD
|
||||
return TensorStats(
|
||||
mean=torch.mean(x).item(),
|
||||
std=torch.std(x).item(),
|
||||
min=torch.min(x).item(),
|
||||
max=torch.max(x).item(),
|
||||
p1=_quantile_or_none(x, q=0.01, include=include_quantiles),
|
||||
p5=_quantile_or_none(x, q=0.05, include=include_quantiles),
|
||||
p95=_quantile_or_none(x, q=0.95, include=include_quantiles),
|
||||
p99=_quantile_or_none(x, q=0.99, include=include_quantiles),
|
||||
)
|
||||
|
||||
|
||||
def _quantile_or_none(x: torch.Tensor, *, q: float, include: bool) -> Optional[float]:
|
||||
return torch.quantile(x, q).item() if include else None
|
||||
|
||||
|
||||
def _compute_diff(x_baseline: torch.Tensor, x_target: torch.Tensor) -> DiffInfo:
|
||||
raw_abs_diff = (x_target - x_baseline).abs()
|
||||
max_diff_coord = argmax_coord(raw_abs_diff)
|
||||
|
||||
return DiffInfo(
|
||||
rel_diff=calc_rel_diff(x_target, x_baseline).item(),
|
||||
max_abs_diff=raw_abs_diff.max().item(),
|
||||
mean_abs_diff=raw_abs_diff.mean().item(),
|
||||
max_diff_coord=max_diff_coord,
|
||||
baseline_at_max=x_baseline[max_diff_coord].item(),
|
||||
target_at_max=x_target[max_diff_coord].item(),
|
||||
)
|
||||
@@ -0,0 +1,88 @@
|
||||
from dataclasses import fields
|
||||
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
|
||||
DiffInfo,
|
||||
TensorComparisonInfo,
|
||||
TensorStats,
|
||||
)
|
||||
|
||||
|
||||
def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None:
|
||||
baseline = info.baseline
|
||||
target = info.target
|
||||
|
||||
dtype_marker = "" if baseline.dtype == target.dtype else "🟠"
|
||||
print(
|
||||
f"Raw "
|
||||
f"[shape] {baseline.shape} vs {target.shape}\t"
|
||||
f"[{dtype_marker}dtype] {baseline.dtype} vs {target.dtype}"
|
||||
)
|
||||
|
||||
if info.unified_shape != baseline.shape:
|
||||
print(
|
||||
f"Unify shape: {baseline.shape} -> {info.unified_shape} "
|
||||
f"(to match {target.shape})"
|
||||
)
|
||||
|
||||
print(
|
||||
f"After unify "
|
||||
f"[shape] {info.unified_shape} vs {target.shape}\t"
|
||||
f"[dtype] {baseline.dtype} vs {target.dtype}"
|
||||
)
|
||||
|
||||
_print_stats_comparison(baseline=baseline.stats, target=target.stats)
|
||||
|
||||
if info.shape_mismatch:
|
||||
print("⚠️ Shape mismatch")
|
||||
return
|
||||
|
||||
if info.diff is not None:
|
||||
_print_diff(
|
||||
diff=info.diff,
|
||||
diff_threshold=diff_threshold,
|
||||
)
|
||||
|
||||
if info.diff_downcast is not None and info.downcast_dtype is not None:
|
||||
_print_diff(
|
||||
diff=info.diff_downcast,
|
||||
diff_threshold=diff_threshold,
|
||||
prefix_text=f"When downcast to {info.downcast_dtype}: ",
|
||||
)
|
||||
|
||||
if baseline.sample is not None:
|
||||
print(f"x_baseline(sample)={baseline.sample}")
|
||||
if target.sample is not None:
|
||||
print(f"x_target(sample)={target.sample}")
|
||||
|
||||
|
||||
def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None:
|
||||
stat_names = [f.name for f in fields(TensorStats)]
|
||||
for stat_name in stat_names:
|
||||
value_baseline = getattr(baseline, stat_name)
|
||||
value_target = getattr(target, stat_name)
|
||||
if value_baseline is None or value_target is None:
|
||||
continue
|
||||
print(
|
||||
f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
|
||||
f"(diff: {value_target - value_baseline:.4f})"
|
||||
)
|
||||
|
||||
|
||||
def _print_diff(diff: DiffInfo, diff_threshold: float, prefix_text: str = "") -> None:
|
||||
print(
|
||||
prefix_text
|
||||
+ "\t".join(
|
||||
f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
|
||||
for name, value in [
|
||||
("rel_diff", diff.rel_diff),
|
||||
("max_abs_diff", diff.max_abs_diff),
|
||||
("mean_abs_diff", diff.mean_abs_diff),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
print(
|
||||
f"max_abs_diff happens at coord={diff.max_diff_coord} with "
|
||||
f"baseline={diff.baseline_at_max} "
|
||||
f"target={diff.target_at_max}"
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class TensorStats:
|
||||
mean: float
|
||||
std: float
|
||||
min: float
|
||||
max: float
|
||||
p1: Optional[float] = None
|
||||
p5: Optional[float] = None
|
||||
p95: Optional[float] = None
|
||||
p99: Optional[float] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class TensorInfo:
|
||||
shape: torch.Size
|
||||
dtype: torch.dtype
|
||||
stats: TensorStats
|
||||
sample: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffInfo:
|
||||
rel_diff: float
|
||||
max_abs_diff: float
|
||||
mean_abs_diff: float
|
||||
max_diff_coord: Tuple[int, ...]
|
||||
baseline_at_max: float
|
||||
target_at_max: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class TensorComparisonInfo:
|
||||
name: str
|
||||
baseline: TensorInfo
|
||||
target: TensorInfo
|
||||
unified_shape: Optional[torch.Size]
|
||||
shape_mismatch: bool
|
||||
diff: Optional[DiffInfo] = None
|
||||
diff_downcast: Optional[DiffInfo] = None
|
||||
downcast_dtype: Optional[torch.dtype] = None
|
||||
55
python/sglang/srt/debug_utils/comparator/utils.py
Normal file
55
python/sglang/srt/debug_utils/comparator/utils.py
Normal file
@@ -0,0 +1,55 @@
|
||||
import functools
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]:
|
||||
flat_idx = x.argmax()
|
||||
return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))
|
||||
|
||||
|
||||
def compute_smaller_dtype(
|
||||
dtype_a: torch.dtype, dtype_b: torch.dtype
|
||||
) -> Optional[torch.dtype]:
|
||||
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: torch.Size) -> torch.Tensor:
|
||||
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]
|
||||
):
|
||||
return functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# Copied from DeepGEMM
|
||||
def calc_rel_diff(x: torch.Tensor, y: torch.Tensor) -> 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: 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()
|
||||
@@ -1,328 +0,0 @@
|
||||
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
|
||||
|
||||
from sglang.srt.debug_utils.dump_loader import find_row, read_meta
|
||||
from sglang.srt.debug_utils.dumper import get_truncated_value
|
||||
|
||||
|
||||
def main(args):
|
||||
df_target = read_meta(args.target_path)
|
||||
df_target = df_target.filter(
|
||||
(pl.col("step") >= args.start_id) & (pl.col("step") <= args.end_id)
|
||||
)
|
||||
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)
|
||||
|
||||
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["step"])
|
||||
if location_info is None:
|
||||
continue
|
||||
baseline_step = location_info.baseline_step
|
||||
baseline_token_slice = location_info.baseline_token_slice
|
||||
else:
|
||||
baseline_step = row["step"] - 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(
|
||||
step=baseline_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"],
|
||||
baseline_token_slice=baseline_token_slice,
|
||||
tensor_dim_desc=tensor_dim_desc,
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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:
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
|
||||
# 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_step: int
|
||||
baseline_token_slice: slice
|
||||
|
||||
|
||||
def _get_location_info_of_target_pass_id() -> Optional[Dict[int, LocationInfo]]:
|
||||
"""Customization endpoint."""
|
||||
return None
|
||||
|
||||
|
||||
@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)
|
||||
parser.add_argument(
|
||||
"--filter", type=str, default=None, help="Regex to filter filenames"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,150 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.compare import (
|
||||
QUANTILE_NUMEL_THRESHOLD,
|
||||
SAMPLE_DIFF_THRESHOLD,
|
||||
_compute_diff,
|
||||
_compute_tensor_stats,
|
||||
compare_tensors,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=20, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestComputeTensorStats:
|
||||
def test_basic_stats(self):
|
||||
x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
|
||||
stats = _compute_tensor_stats(x)
|
||||
|
||||
assert stats.mean == pytest.approx(3.0, abs=1e-4)
|
||||
assert stats.std == pytest.approx(1.5811, abs=1e-3)
|
||||
assert stats.min == pytest.approx(1.0, abs=1e-4)
|
||||
assert stats.max == pytest.approx(5.0, abs=1e-4)
|
||||
|
||||
def test_quantile_values(self):
|
||||
x = torch.linspace(0.0, 100.0, steps=1000)
|
||||
stats = _compute_tensor_stats(x)
|
||||
|
||||
assert stats.p1 == pytest.approx(1.0, abs=0.5)
|
||||
assert stats.p5 == pytest.approx(5.0, abs=0.5)
|
||||
assert stats.p95 == pytest.approx(95.0, abs=0.5)
|
||||
assert stats.p99 == pytest.approx(99.0, abs=0.5)
|
||||
|
||||
def test_large_tensor_skips_quantiles(self):
|
||||
x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1)
|
||||
stats = _compute_tensor_stats(x)
|
||||
|
||||
assert stats.mean is not None
|
||||
assert stats.p1 is None
|
||||
assert stats.p5 is None
|
||||
assert stats.p95 is None
|
||||
assert stats.p99 is None
|
||||
|
||||
|
||||
class TestComputeDiff:
|
||||
def test_identical_tensors(self):
|
||||
x = torch.ones(10, 10)
|
||||
diff = _compute_diff(x_baseline=x, x_target=x)
|
||||
|
||||
assert diff.rel_diff == pytest.approx(0.0, abs=1e-5)
|
||||
assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
|
||||
assert diff.mean_abs_diff == pytest.approx(0.0, abs=1e-5)
|
||||
|
||||
def test_known_offset(self):
|
||||
x = torch.ones(10, 10)
|
||||
y = x.clone()
|
||||
y[3, 7] = 1.5
|
||||
|
||||
diff = _compute_diff(x_baseline=x, x_target=y)
|
||||
|
||||
assert diff.max_abs_diff == pytest.approx(0.5, abs=1e-4)
|
||||
assert diff.max_diff_coord == (3, 7)
|
||||
assert diff.baseline_at_max == pytest.approx(1.0, abs=1e-4)
|
||||
assert diff.target_at_max == pytest.approx(1.5, abs=1e-4)
|
||||
assert diff.mean_abs_diff == pytest.approx(0.5 / 100, abs=1e-4)
|
||||
|
||||
def test_rel_diff_value(self):
|
||||
x = torch.tensor([1.0, 0.0])
|
||||
y = torch.tensor([0.0, 1.0])
|
||||
diff = _compute_diff(x_baseline=x, x_target=y)
|
||||
|
||||
assert diff.rel_diff == pytest.approx(1.0, abs=1e-5)
|
||||
|
||||
|
||||
class TestCompareTensors:
|
||||
def test_normal(self):
|
||||
x = torch.randn(5, 5)
|
||||
y = x + torch.randn(5, 5) * 0.001
|
||||
|
||||
info = compare_tensors(x_baseline=x, x_target=y, name="test")
|
||||
|
||||
assert info.name == "test"
|
||||
assert info.baseline.shape == torch.Size([5, 5])
|
||||
assert info.target.shape == torch.Size([5, 5])
|
||||
assert info.shape_mismatch is False
|
||||
assert info.diff is not None
|
||||
assert info.diff_downcast is None
|
||||
|
||||
def test_shape_mismatch(self):
|
||||
x = torch.randn(3, 4)
|
||||
y = torch.randn(5, 6)
|
||||
|
||||
info = compare_tensors(x_baseline=x, x_target=y, name="mismatch")
|
||||
|
||||
assert info.shape_mismatch is True
|
||||
assert info.diff is None
|
||||
|
||||
def test_dtype_mismatch(self):
|
||||
x = torch.randn(5, 5, dtype=torch.float32)
|
||||
y = torch.randn(5, 5, dtype=torch.bfloat16)
|
||||
|
||||
info = compare_tensors(x_baseline=x, x_target=y, name="dtype_test")
|
||||
|
||||
assert info.shape_mismatch is False
|
||||
assert info.diff is not None
|
||||
assert info.diff_downcast is not None
|
||||
assert info.downcast_dtype == torch.bfloat16
|
||||
|
||||
def test_shape_unification(self):
|
||||
torch.manual_seed(0)
|
||||
core = torch.randn(4, 8)
|
||||
x = core.unsqueeze(0).unsqueeze(0) # [1, 1, 4, 8]
|
||||
y = core.clone() # [4, 8]
|
||||
|
||||
info = compare_tensors(x_baseline=x, x_target=y, name="unify")
|
||||
|
||||
assert info.baseline.shape == torch.Size([1, 1, 4, 8])
|
||||
assert info.unified_shape == torch.Size([4, 8])
|
||||
assert info.shape_mismatch is False
|
||||
assert info.diff is not None
|
||||
assert info.diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
|
||||
|
||||
def test_sample_generated_when_large_diff(self):
|
||||
x = torch.zeros(5, 5)
|
||||
y = torch.ones(5, 5)
|
||||
|
||||
info = compare_tensors(x_baseline=x, x_target=y, name="big_diff")
|
||||
|
||||
assert info.diff is not None
|
||||
assert info.diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD
|
||||
assert info.baseline.sample is not None
|
||||
assert info.target.sample is not None
|
||||
|
||||
def test_no_sample_when_small_diff(self):
|
||||
x = torch.ones(5, 5)
|
||||
y = x + 1e-5
|
||||
|
||||
info = compare_tensors(x_baseline=x, x_target=y, name="tiny_diff")
|
||||
|
||||
assert info.diff is not None
|
||||
assert info.diff.max_abs_diff < SAMPLE_DIFF_THRESHOLD
|
||||
assert info.baseline.sample is None
|
||||
assert info.target.sample is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,262 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.printer import (
|
||||
print_comparison,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
|
||||
DiffInfo,
|
||||
TensorComparisonInfo,
|
||||
TensorInfo,
|
||||
TensorStats,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
|
||||
def _make_stats(
|
||||
mean: float = 0.0,
|
||||
std: float = 1.0,
|
||||
min: float = -2.0,
|
||||
max: float = 2.0,
|
||||
p1: float | None = -1.8,
|
||||
p5: float | None = -1.5,
|
||||
p95: float | None = 1.5,
|
||||
p99: float | None = 1.8,
|
||||
) -> TensorStats:
|
||||
return TensorStats(
|
||||
mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
|
||||
)
|
||||
|
||||
|
||||
def _make_diff(
|
||||
rel_diff: float = 0.0001,
|
||||
max_abs_diff: float = 0.0005,
|
||||
mean_abs_diff: float = 0.0002,
|
||||
) -> DiffInfo:
|
||||
return DiffInfo(
|
||||
rel_diff=rel_diff,
|
||||
max_abs_diff=max_abs_diff,
|
||||
mean_abs_diff=mean_abs_diff,
|
||||
max_diff_coord=(2, 3),
|
||||
baseline_at_max=1.0,
|
||||
target_at_max=1.0005,
|
||||
)
|
||||
|
||||
|
||||
def _make_tensor_info(
|
||||
shape: torch.Size = torch.Size([4, 8]),
|
||||
dtype: torch.dtype = torch.float32,
|
||||
stats: TensorStats | None = None,
|
||||
sample: str | None = None,
|
||||
) -> TensorInfo:
|
||||
return TensorInfo(
|
||||
shape=shape,
|
||||
dtype=dtype,
|
||||
stats=stats if stats is not None else _make_stats(),
|
||||
sample=sample,
|
||||
)
|
||||
|
||||
|
||||
# Snapshot strings below are intentionally spelled out in full per test.
|
||||
# The shared skeleton (stats block, diff block) looks duplicated, but keeping
|
||||
# each test self-contained makes failures immediately readable without chasing
|
||||
# helper functions. Do not extract common fragments.
|
||||
class TestPrintComparison:
|
||||
def test_normal(self, capsys):
|
||||
info = TensorComparisonInfo(
|
||||
name="test",
|
||||
baseline=_make_tensor_info(
|
||||
stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0),
|
||||
),
|
||||
target=_make_tensor_info(
|
||||
stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001),
|
||||
),
|
||||
unified_shape=torch.Size([4, 8]),
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(),
|
||||
)
|
||||
|
||||
print_comparison(info=info, diff_threshold=1e-3)
|
||||
|
||||
assert capsys.readouterr().out == (
|
||||
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n"
|
||||
"[std] 1.0000 vs 1.0001 (diff: 0.0001)\n"
|
||||
"[min] -2.0000 vs -2.0001 (diff: -0.0001)\n"
|
||||
"[max] 2.0000 vs 2.0001 (diff: 0.0001)\n"
|
||||
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
|
||||
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
|
||||
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
|
||||
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
|
||||
"✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
|
||||
"max_abs_diff happens at coord=(2, 3) with "
|
||||
"baseline=1.0 target=1.0005\n"
|
||||
)
|
||||
|
||||
def test_shape_mismatch(self, capsys):
|
||||
info = TensorComparisonInfo(
|
||||
name="mismatch",
|
||||
baseline=_make_tensor_info(shape=torch.Size([3, 4])),
|
||||
target=_make_tensor_info(shape=torch.Size([5, 6])),
|
||||
unified_shape=torch.Size([3, 4]),
|
||||
shape_mismatch=True,
|
||||
)
|
||||
|
||||
print_comparison(info=info, diff_threshold=1e-3)
|
||||
|
||||
assert capsys.readouterr().out == (
|
||||
"Raw [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"After unify [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
|
||||
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
|
||||
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
|
||||
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
|
||||
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
|
||||
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
|
||||
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
|
||||
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
|
||||
"⚠️ Shape mismatch\n"
|
||||
)
|
||||
|
||||
def test_with_downcast(self, capsys):
|
||||
info = TensorComparisonInfo(
|
||||
name="downcast",
|
||||
baseline=_make_tensor_info(),
|
||||
target=_make_tensor_info(dtype=torch.bfloat16),
|
||||
unified_shape=torch.Size([4, 8]),
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001),
|
||||
diff_downcast=_make_diff(
|
||||
rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002
|
||||
),
|
||||
downcast_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
print_comparison(info=info, diff_threshold=1e-3)
|
||||
|
||||
assert capsys.readouterr().out == (
|
||||
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[🟠dtype] torch.float32 vs torch.bfloat16\n"
|
||||
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.bfloat16\n"
|
||||
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
|
||||
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
|
||||
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
|
||||
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
|
||||
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
|
||||
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
|
||||
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
|
||||
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
|
||||
"❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n"
|
||||
"max_abs_diff happens at coord=(2, 3) with "
|
||||
"baseline=1.0 target=1.0005\n"
|
||||
"When downcast to torch.bfloat16: "
|
||||
"✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
|
||||
"max_abs_diff happens at coord=(2, 3) with "
|
||||
"baseline=1.0 target=1.0005\n"
|
||||
)
|
||||
|
||||
def test_with_shape_unification(self, capsys):
|
||||
info = TensorComparisonInfo(
|
||||
name="unify",
|
||||
baseline=_make_tensor_info(shape=torch.Size([1, 1, 4, 8])),
|
||||
target=_make_tensor_info(),
|
||||
unified_shape=torch.Size([4, 8]),
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(),
|
||||
)
|
||||
|
||||
print_comparison(info=info, diff_threshold=1e-3)
|
||||
|
||||
assert capsys.readouterr().out == (
|
||||
"Raw [shape] torch.Size([1, 1, 4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"Unify shape: torch.Size([1, 1, 4, 8]) -> torch.Size([4, 8]) "
|
||||
"(to match torch.Size([4, 8]))\n"
|
||||
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
|
||||
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
|
||||
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
|
||||
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
|
||||
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
|
||||
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
|
||||
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
|
||||
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
|
||||
"✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
|
||||
"max_abs_diff happens at coord=(2, 3) with "
|
||||
"baseline=1.0 target=1.0005\n"
|
||||
)
|
||||
|
||||
def test_with_samples(self, capsys):
|
||||
info = TensorComparisonInfo(
|
||||
name="samples",
|
||||
baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"),
|
||||
target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"),
|
||||
unified_shape=torch.Size([4, 8]),
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(),
|
||||
)
|
||||
|
||||
print_comparison(info=info, diff_threshold=1e-3)
|
||||
|
||||
assert capsys.readouterr().out == (
|
||||
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
|
||||
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
|
||||
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
|
||||
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
|
||||
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
|
||||
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
|
||||
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
|
||||
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
|
||||
"✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
|
||||
"max_abs_diff happens at coord=(2, 3) with "
|
||||
"baseline=1.0 target=1.0005\n"
|
||||
"x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
|
||||
"x_target(sample)=tensor([0.1, 0.3, ...])\n"
|
||||
)
|
||||
|
||||
def test_none_quantiles(self, capsys):
|
||||
stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None)
|
||||
|
||||
info = TensorComparisonInfo(
|
||||
name="no_quantiles",
|
||||
baseline=_make_tensor_info(stats=stats_no_quantiles),
|
||||
target=_make_tensor_info(stats=stats_no_quantiles),
|
||||
unified_shape=torch.Size([4, 8]),
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(),
|
||||
)
|
||||
|
||||
print_comparison(info=info, diff_threshold=1e-3)
|
||||
|
||||
assert capsys.readouterr().out == (
|
||||
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
|
||||
"[dtype] torch.float32 vs torch.float32\n"
|
||||
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
|
||||
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
|
||||
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
|
||||
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
|
||||
"✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
|
||||
"max_abs_diff happens at coord=(2, 3) with "
|
||||
"baseline=1.0 target=1.0005\n"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
125
test/registered/debug_utils/comparator/test_entrypoint.py
Normal file
125
test/registered/debug_utils/comparator/test_entrypoint.py
Normal file
@@ -0,0 +1,125 @@
|
||||
import sys
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.entrypoint import run
|
||||
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)
|
||||
|
||||
|
||||
def _make_dumper(directory: Path) -> _Dumper:
|
||||
return _Dumper(
|
||||
config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False)
|
||||
)
|
||||
|
||||
|
||||
def _create_dumps(
|
||||
tmp_path: Path,
|
||||
tensor_names: list[str],
|
||||
*,
|
||||
baseline_names: list[str] | None = None,
|
||||
num_steps: int = 1,
|
||||
) -> tuple[Path, Path]:
|
||||
"""Create baseline and target dump directories with given tensor names.
|
||||
|
||||
If baseline_names is None, uses the same names as tensor_names.
|
||||
Each step dumps all names with the same tensor (different per baseline/target).
|
||||
"""
|
||||
if baseline_names is None:
|
||||
baseline_names = tensor_names
|
||||
|
||||
d_baseline = tmp_path / "baseline"
|
||||
d_target = tmp_path / "target"
|
||||
d_baseline.mkdir()
|
||||
d_target.mkdir()
|
||||
|
||||
torch.manual_seed(42)
|
||||
baseline_tensor = torch.randn(10, 10)
|
||||
target_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 = _make_dumper(d)
|
||||
for _ in range(num_steps):
|
||||
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, **overrides) -> Namespace:
|
||||
defaults = dict(
|
||||
baseline_path=str(baseline_path),
|
||||
target_path=str(target_path),
|
||||
start_step=0,
|
||||
end_step=1000000,
|
||||
diff_threshold=1e-3,
|
||||
filter=None,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return Namespace(**defaults)
|
||||
|
||||
|
||||
class TestEntrypoint:
|
||||
def test_run_basic(self, tmp_path, capsys):
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args = _make_args(baseline_path, target_path)
|
||||
|
||||
run(args)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert "df_target" in output
|
||||
assert "df_baseline" in output
|
||||
assert output.count("Check:") == 2
|
||||
assert "tensor_a" in output
|
||||
assert "tensor_b" in output
|
||||
assert "rel_diff" in output
|
||||
assert "Skip" not in output
|
||||
|
||||
def test_filter(self, tmp_path, capsys):
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args = _make_args(baseline_path, target_path, filter="tensor_a")
|
||||
|
||||
run(args)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert output.count("Check:") == 1
|
||||
assert "tensor_a" in output
|
||||
|
||||
def test_no_baseline_skip(self, tmp_path, capsys):
|
||||
baseline_path, target_path = _create_dumps(
|
||||
tmp_path,
|
||||
tensor_names=["tensor_a", "tensor_extra"],
|
||||
baseline_names=["tensor_a"],
|
||||
)
|
||||
args = _make_args(baseline_path, target_path)
|
||||
|
||||
run(args)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert output.count("Check:") == 1
|
||||
assert "Skip:" in output
|
||||
assert "since no baseline" in output
|
||||
|
||||
def test_step_range(self, tmp_path, capsys):
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=3)
|
||||
args = _make_args(baseline_path, target_path, start_step=1, end_step=1)
|
||||
|
||||
run(args)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert output.count("Check:") == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
119
test/registered/debug_utils/comparator/test_utils.py
Normal file
119
test/registered/debug_utils/comparator/test_utils.py
Normal file
@@ -0,0 +1,119 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.debug_utils.comparator.utils import (
|
||||
argmax_coord,
|
||||
calc_rel_diff,
|
||||
compute_smaller_dtype,
|
||||
load_object,
|
||||
try_unify_shape,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestCalcRelDiff:
|
||||
def test_identical_tensors(self):
|
||||
x = torch.randn(10, 10)
|
||||
assert calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
|
||||
|
||||
def test_orthogonal_tensors(self):
|
||||
result = calc_rel_diff(
|
||||
torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])
|
||||
).item()
|
||||
assert result == pytest.approx(1.0, abs=1e-5)
|
||||
|
||||
def test_similar_tensors(self):
|
||||
x = torch.tensor([1.0, 2.0, 3.0])
|
||||
y = torch.tensor([1.01, 2.01, 3.01])
|
||||
result = calc_rel_diff(x, y).item()
|
||||
assert 0.0 < result < 0.01
|
||||
|
||||
def test_negated_tensors(self):
|
||||
x = torch.tensor([1.0, 2.0])
|
||||
result = calc_rel_diff(x, -x).item()
|
||||
assert result == pytest.approx(2.0, abs=1e-5)
|
||||
|
||||
|
||||
class TestArgmaxCoord:
|
||||
def test_1d_tensor(self):
|
||||
x = torch.tensor([0.0, 0.0, 5.0, 0.0])
|
||||
assert argmax_coord(x) == (2,)
|
||||
|
||||
def test_2d_tensor(self):
|
||||
x = torch.zeros(3, 4)
|
||||
x[1, 2] = 10.0
|
||||
assert argmax_coord(x) == (1, 2)
|
||||
|
||||
def test_3d_tensor(self):
|
||||
x = 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):
|
||||
target = torch.Size([3, 4])
|
||||
assert try_unify_shape(torch.randn(1, 1, 3, 4), target).shape == target
|
||||
|
||||
def test_no_squeeze_when_leading_dim_not_one(self):
|
||||
target = torch.Size([3, 4])
|
||||
assert try_unify_shape(torch.randn(2, 3, 4), target).shape == (2, 3, 4)
|
||||
|
||||
def test_same_shape_noop(self):
|
||||
target = torch.Size([3, 4])
|
||||
x = torch.randn(3, 4)
|
||||
result = try_unify_shape(x, target)
|
||||
assert result.shape == target
|
||||
assert result.data_ptr() == x.data_ptr()
|
||||
|
||||
def test_trailing_dims_mismatch(self):
|
||||
target = torch.Size([5, 6])
|
||||
x = torch.randn(1, 3, 4)
|
||||
result = try_unify_shape(x, target)
|
||||
assert result.shape == (1, 3, 4)
|
||||
|
||||
|
||||
class TestComputeSmallerDtype:
|
||||
def test_float32_bfloat16(self):
|
||||
assert compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16
|
||||
|
||||
def test_reverse_order(self):
|
||||
assert compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16
|
||||
|
||||
def test_same_dtype_returns_none(self):
|
||||
assert compute_smaller_dtype(torch.float32, torch.float32) is None
|
||||
|
||||
def test_unknown_pair_returns_none(self):
|
||||
assert compute_smaller_dtype(torch.int32, torch.int64) is None
|
||||
|
||||
|
||||
class TestLoadObject:
|
||||
def test_load_tensor(self, tmp_path):
|
||||
path = tmp_path / "tensor.pt"
|
||||
torch.save(torch.randn(5, 5), path)
|
||||
assert load_object(path).shape == (5, 5)
|
||||
|
||||
def test_load_dict_with_value_key(self, tmp_path):
|
||||
path = tmp_path / "wrapped.pt"
|
||||
tensor = torch.randn(3, 3)
|
||||
torch.save({"value": tensor}, path)
|
||||
result = load_object(path)
|
||||
assert result is not None
|
||||
assert result.shape == (3, 3)
|
||||
|
||||
def test_non_tensor_returns_none(self, tmp_path):
|
||||
path = tmp_path / "tensor.pt"
|
||||
torch.save({"dict": 1}, path)
|
||||
assert load_object(path) is None
|
||||
|
||||
def test_nonexistent_returns_none(self):
|
||||
assert load_object(Path("/nonexistent.pt")) is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -1,139 +0,0 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=60, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestDumpComparator(CustomTestCase):
|
||||
def test_calc_rel_diff(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import _calc_rel_diff
|
||||
|
||||
x = torch.randn(10, 10)
|
||||
self.assertAlmostEqual(_calc_rel_diff(x, x).item(), 0.0, places=5)
|
||||
self.assertAlmostEqual(
|
||||
_calc_rel_diff(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])).item(),
|
||||
1.0,
|
||||
places=5,
|
||||
)
|
||||
|
||||
def test_argmax_coord(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import _argmax_coord
|
||||
|
||||
x = torch.zeros(2, 3, 4)
|
||||
x[1, 2, 3] = 10.0
|
||||
self.assertEqual(_argmax_coord(x), (1, 2, 3))
|
||||
|
||||
def test_try_unify_shape(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import _try_unify_shape
|
||||
|
||||
target = torch.Size([3, 4])
|
||||
self.assertEqual(
|
||||
_try_unify_shape(torch.randn(1, 1, 3, 4), target).shape, target
|
||||
)
|
||||
self.assertEqual(
|
||||
_try_unify_shape(torch.randn(2, 3, 4), target).shape, (2, 3, 4)
|
||||
)
|
||||
|
||||
def test_compute_smaller_dtype(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import _compute_smaller_dtype
|
||||
|
||||
self.assertEqual(
|
||||
_compute_smaller_dtype(torch.float32, torch.bfloat16), torch.bfloat16
|
||||
)
|
||||
self.assertIsNone(_compute_smaller_dtype(torch.float32, torch.float32))
|
||||
|
||||
def test_einops_pattern(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import (
|
||||
_get_einops_dim_index,
|
||||
_split_einops_pattern,
|
||||
)
|
||||
|
||||
self.assertEqual(_split_einops_pattern("a (b c) d"), ["a", "(b c)", "d"])
|
||||
self.assertEqual(_get_einops_dim_index("a b c", "b"), 1)
|
||||
|
||||
def test_load_object(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import _load_object
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = Path(tmpdir) / "tensor.pt"
|
||||
torch.save(torch.randn(5, 5), path)
|
||||
self.assertEqual(_load_object(path).shape, (5, 5))
|
||||
|
||||
torch.save({"dict": 1}, path)
|
||||
self.assertIsNone(_load_object(path))
|
||||
|
||||
self.assertIsNone(_load_object("/nonexistent.pt"))
|
||||
|
||||
def test_compute_and_print_diff(self):
|
||||
from sglang.srt.debug_utils.dump_comparator import _compute_and_print_diff
|
||||
|
||||
x = torch.ones(10, 10)
|
||||
self.assertAlmostEqual(
|
||||
_compute_and_print_diff(x, x, 1e-3)["max_abs_diff"], 0.0, places=5
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
_compute_and_print_diff(x, x + 0.5, 1e-3)["max_abs_diff"], 0.5, places=4
|
||||
)
|
||||
|
||||
|
||||
class TestEndToEnd(CustomTestCase):
|
||||
def test_main(self):
|
||||
from argparse import Namespace
|
||||
|
||||
from sglang.srt.debug_utils.dump_comparator import main
|
||||
from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
|
||||
|
||||
with tempfile.TemporaryDirectory() as d1, tempfile.TemporaryDirectory() as d2:
|
||||
baseline_tensor = torch.randn(10, 10)
|
||||
target_tensor = baseline_tensor + torch.randn(10, 10) * 0.01
|
||||
|
||||
dump_dirs = []
|
||||
for d, tensor in [(d1, baseline_tensor), (d2, target_tensor)]:
|
||||
dumper = _Dumper(
|
||||
config=DumperConfig(
|
||||
enable=True,
|
||||
dir=d,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
dumper.dump("tensor_a", tensor)
|
||||
dumper.step()
|
||||
dumper.dump("tensor_b", tensor * 2)
|
||||
dumper.step()
|
||||
dump_dirs.append(Path(d) / dumper._config.exp_name)
|
||||
|
||||
args = Namespace(
|
||||
baseline_path=str(dump_dirs[0]),
|
||||
target_path=str(dump_dirs[1]),
|
||||
start_id=0,
|
||||
end_id=1,
|
||||
baseline_start_id=0,
|
||||
diff_threshold=1e-3,
|
||||
filter=None,
|
||||
)
|
||||
main(args)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _with_env(name: str, value: str):
|
||||
old = os.environ.get(name)
|
||||
os.environ[name] = value
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if old is None:
|
||||
os.environ.pop(name, None)
|
||||
else:
|
||||
os.environ[name] = old
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user