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)
|
||||
Reference in New Issue
Block a user