diff --git a/python/sglang/srt/debug_utils/dump_comparator.py b/python/sglang/srt/debug_utils/dump_comparator.py index aca9c3b7a..77ce5c5b7 100644 --- a/python/sglang/srt/debug_utils/dump_comparator.py +++ b/python/sglang/srt/debug_utils/dump_comparator.py @@ -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) diff --git a/python/sglang/srt/debug_utils/dump_loader.py b/python/sglang/srt/debug_utils/dump_loader.py index 8e6f2c79b..fe54cc2ab 100644 --- a/python/sglang/srt/debug_utils/dump_loader.py +++ b/python/sglang/srt/debug_utils/dump_loader.py @@ -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