chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support dumping rid for correlation across passes in dump comparator (#19372)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-26 09:57:57 +08:00
committed by GitHub
parent 7c9e8e2def
commit 46321ee70e
5 changed files with 69 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py
View File

@@ -1,5 +1,3 @@
from dataclasses import fields
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
DiffInfo,
TensorComparisonInfo,
@@ -56,7 +54,7 @@ def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None:
def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None:
stat_names = [f.name for f in fields(TensorStats)]
stat_names = list(TensorStats.model_fields.keys())
for stat_name in stat_names:
value_baseline = getattr(baseline, stat_name)
value_target = getattr(target, stat_name)

10
python/sglang/srt/debug_utils/dumper.py
View File

@@ -1165,11 +1165,15 @@ class _SGLangPlugin(_FrameworkPlugin):
if isinstance(value, self.ForwardBatch):
if skip_forward_batch:
return {}
return {
result = {
"input_ids": value.input_ids,
"seq_lens": value.seq_lens,
"positions": value.positions,
"req_pool_indices": value.req_pool_indices,
}
if value.rids is not None:
result["rids"] = value.rids
return result
if isinstance(value, self.PPProxyTensors):
return {k: v for k, v in value.tensors.items()}
@@ -1181,7 +1185,9 @@ class _SGLangPlugin(_FrameworkPlugin):
return None
def core_fields(self) -> frozenset[str]:
return frozenset({"input_ids", "positions", "seq_lens"})
return frozenset(
{"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"}
)
class _MegatronPlugin(_FrameworkPlugin):

4
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -373,6 +373,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# For hidden states before normal
return_hidden_states_before_norm: bool = False
# For dumper: request IDs for cross-step sequence tracking
rids: Optional[List[str]] = None
@classmethod
def init_new(
cls,
@@ -417,6 +420,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
tbo_split_seq_index=batch.tbo_split_seq_index,
dimensions=batch.dimensions,
return_hidden_states_before_norm=batch.return_hidden_states_before_norm,
rids=[req.rid for req in batch.reqs],
)
device = model_runner.device

81
test/registered/debug_utils/comparator/tensor_comparison/test_printer.py
View File

@@ -1,7 +1,6 @@
import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.tensor_comparison.printer import (
print_comparison,
@@ -36,25 +35,27 @@ def _make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_diff,
max_diff_coord=(2, 3),
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
passed=passed,
)
def _make_tensor_info(
shape: torch.Size = torch.Size([4, 8]),
dtype: torch.dtype = torch.float32,
shape: list[int] | None = None,
dtype: str = "torch.float32",
stats: TensorStats | None = None,
sample: str | None = None,
) -> TensorInfo:
return TensorInfo(
shape=shape,
shape=shape if shape is not None else [4, 8],
dtype=dtype,
stats=stats if stats is not None else _make_stats(),
sample=sample,
@@ -75,7 +76,7 @@ class TestPrintComparison:
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]),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
@@ -83,9 +84,9 @@ class TestPrintComparison:
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"After unify [shape] [4, 8] vs [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"
@@ -96,25 +97,25 @@ class TestPrintComparison:
"[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 "
"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]),
baseline=_make_tensor_info(shape=[3, 4]),
target=_make_tensor_info(shape=[5, 6]),
unified_shape=[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"
"Raw [shape] [3, 4] vs [5, 6]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t"
"After unify [shape] [3, 4] vs [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"
@@ -131,22 +132,24 @@ class TestPrintComparison:
info = TensorComparisonInfo(
name="downcast",
baseline=_make_tensor_info(),
target=_make_tensor_info(dtype=torch.bfloat16),
unified_shape=torch.Size([4, 8]),
target=_make_tensor_info(dtype="torch.bfloat16"),
unified_shape=[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
diff=_make_diff(
rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001, passed=False
),
downcast_dtype=torch.bfloat16,
diff_downcast=_make_diff(
rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002, passed=True
),
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"
"Raw [shape] [4, 8] vs [4, 8]\t"
"[🟠dtype] torch.float32 vs torch.bfloat16\n"
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"After unify [shape] [4, 8] vs [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"
@@ -157,20 +160,20 @@ class TestPrintComparison:
"[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 "
"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 "
"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])),
baseline=_make_tensor_info(shape=[1, 1, 4, 8]),
target=_make_tensor_info(),
unified_shape=torch.Size([4, 8]),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
@@ -178,11 +181,11 @@ class TestPrintComparison:
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"
"Raw [shape] [1, 1, 4, 8] vs [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"
"Unify shape: [1, 1, 4, 8] -> [4, 8] "
"(to match [4, 8])\n"
"After unify [shape] [4, 8] vs [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"
@@ -193,7 +196,7 @@ class TestPrintComparison:
"[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 "
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005\n"
)
@@ -202,7 +205,7 @@ class TestPrintComparison:
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]),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
@@ -210,9 +213,9 @@ class TestPrintComparison:
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"After unify [shape] [4, 8] vs [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"
@@ -223,7 +226,7 @@ class TestPrintComparison:
"[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 "
"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"
@@ -236,7 +239,7 @@ class TestPrintComparison:
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]),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
@@ -244,16 +247,16 @@ class TestPrintComparison:
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t"
"After unify [shape] [4, 8] vs [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 "
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005\n"
)

15
test/registered/debug_utils/test_dumper.py
View File

@@ -2002,7 +2002,7 @@ class TestDumperE2E:
assert len(dump_files) > 0, f"No dump files in {dump_dir}"
filenames = {f.name for f in dump_files}
for field in ("input_ids", "positions"):
for field in ("input_ids", "positions", "rids"):
assert any(f"name={field}" in f for f in filenames), (
f"Missing {field} dump from non-intrusive hooks, "
f"got: {sorted(filenames)[:10]}"
@@ -2022,6 +2022,19 @@ class TestDumperE2E:
assert "name" in loaded["meta"]
assert "rank" in loaded["meta"]
assert "step" in loaded["meta"]
rids_files = [f for f in dump_files if "name=rids" in f.name]
rids_loaded = torch.load(
rids_files[0], map_location="cpu", weights_only=False
)
rids_value = rids_loaded["value"]
assert isinstance(
rids_value, list
), f"rids should be a list, got {type(rids_value)}"
assert len(rids_value) > 0, "rids should be non-empty"
assert all(
isinstance(r, str) for r in rids_value
), f"each rid should be a str, got {[type(r) for r in rids_value]}"
finally:
kill_process_tree(proc.pid)