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