diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 07ca7d74d..f291343fd 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -885,7 +885,7 @@ class _Grafter: if is_send: _log( f"[Grafter] send role={role.value} dir={direction.value} " - f"tags={tags} extras={extras}" + f"tags={tags} extras={extras} local={get_tensor_info(value)}" ) return @@ -902,6 +902,7 @@ class _Grafter: # crash the whole training/inference run. On error we log the full # traceback and skip this graft point; downstream sees the recv # side's original tensor unchanged. + info_before_overridden = get_tensor_info(value) try: value_to_override = self._apply_transform( tags=tags, @@ -909,10 +910,14 @@ class _Grafter: received_extras_list=sender_extras, target=value, ) + diff = _compare_tensors_quick(value, value_to_override) _log( f"[Grafter] recv role={role.value} dir={direction.value} " f"tags={tags} n_senders={len(sender_tensors)} " - f"sender_extras={sender_extras}" + f"sender_extras={sender_extras} " + f"before_overridden={info_before_overridden} " + f"to_override={get_tensor_info(value_to_override)} " + f"diff_pre_vs_new={diff}" ) value.copy_(value_to_override) except Exception as e: @@ -1176,6 +1181,35 @@ def _log(msg: str) -> None: print(f"[Dumper, rank={_get_rank()}, t={time.time():.3f}] {msg}", flush=True) +def _compare_tensors_quick(a: "torch.Tensor", b: "torch.Tensor") -> str: + """One-line summary of how close two tensors are. Inspired by + sglang.srt.debug_utils.dump_comparator._compute_and_print_diff; + intentionally inlined here to keep dumper.py free of cross-file imports. + + Different dtypes are fine -- we unify by casting both to fp32, which is + enough for the order-of-magnitude diff summary we log.""" + if a.shape != b.shape: + return f"shape mismatch (a={tuple(a.shape)} vs b={tuple(b.shape)})" + if a.numel() == 0: + return "empty" + a_float = a.detach().to(torch.float32) + b_float = b.detach().to(torch.float32) + raw_abs = (a_float - b_float).abs() + max_abs = raw_abs.max().item() + mean_abs = raw_abs.mean().item() + rel_diff = _calc_rel_diff(a_float, b_float).item() + return f"rel_diff={rel_diff:.6g} max_abs={max_abs:.6g} mean_abs={mean_abs:.6g}" + + +# Copied verbatim from sglang.srt.debug_utils.dump_comparator (originally from +# DeepGEMM). Kept inline here so dumper.py has no cross-file imports. +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 _obj_to_dict(obj): if isinstance(obj, dict): return obj diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 791031b75..208e5aa1f 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -17,6 +17,7 @@ import torch.distributed as dist from sglang.srt.debug_utils.dumper import ( DumperConfig, _collective_with_timeout, + _compare_tensors_quick, _deepcopy_or_clone, _detect_recompute_status, _Dumper, @@ -397,6 +398,37 @@ class TestLog: assert ", t=" in out, out +class TestCompareTensorsQuick: + def test_identical(self): + a = torch.tensor([1.0, 2.0, 3.0]) + s = _compare_tensors_quick(a, a.clone()) + assert "rel_diff=0" in s, s + assert "max_abs=0" in s, s + + def test_diverged(self): + a = torch.tensor([1.0, 2.0, 3.0]) + b = torch.tensor([1.0, 2.0, 4.0]) # last element differs by 1 + s = _compare_tensors_quick(a, b) + assert "max_abs=1" in s, s + assert "rel_diff=" in s, s + + def test_shape_mismatch(self): + s = _compare_tensors_quick(torch.zeros(3), torch.zeros(4)) + assert "shape mismatch" in s, s + + def test_dtype_unified(self): + s = _compare_tensors_quick( + torch.zeros(3, dtype=torch.float32), + torch.zeros(3, dtype=torch.float64), + ) + assert "rel_diff=" in s, s + assert "max_abs=" in s, s + + def test_empty(self): + s = _compare_tensors_quick(torch.zeros(0), torch.zeros(0)) + assert s == "empty" + + class TestCollectiveTimeout: def test_watchdog_fires_on_timeout(self): block_event = threading.Event()