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

Update token layout and cleanup printer in dump comparator (#19457)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:09:15 +08:00
committed by GitHub
parent 8ac64e1487
commit eb0e905fc3
15 changed files with 56 additions and 393 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py
View File

@@ -24,7 +24,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
normalize_parallel_info,
)
from sglang.srt.debug_utils.comparator.dims import ParallelAxis
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows
@@ -61,7 +61,7 @@ def load_and_normalize_aux(
if step_data:
steps_data[step] = step_data
layout: str = plugin.detect_layout(steps_data)
layout: TokenLayout = plugin.detect_layout(steps_data)
step_auxs: dict[int, TokenAlignerStepAux] = {
step: plugin.compute_step_aux(step_data, layout=layout, step=step)

19
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py
View File

@@ -10,6 +10,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SGLangSeqId,
TokenAlignerStepAux,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
@@ -45,11 +46,11 @@ class _AuxFrameworkPlugin(ABC):
return frozenset()
@abstractmethod
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: ...
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: ...
@abstractmethod
def compute_step_aux(
self, step_data: dict[str, object], *, layout: str, step: int
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux: ...
@abstractmethod
@@ -89,11 +90,11 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
def has_required_names(self, names: set[str]) -> bool:
return "input_ids" in names and "seq_lens" in names
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
return "thd"
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
return TokenLayout.T
def compute_step_aux(
self, step_data: dict[str, object], *, layout: str, step: int
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux:
input_ids = step_data["input_ids"]
positions = step_data["positions"]
@@ -154,13 +155,13 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
def has_required_names(self, names: set[str]) -> bool:
return "input_ids" in names and "cu_seqlens_q" in names
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
for step_data in raw.values():
if (qkv_format := step_data.get("qkv_format")) is not None:
fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format)
if "bshd" in fmt.lower():
raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
return "thd"
return TokenLayout.T
input_ids = step_data.get("input_ids")
if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2:
@@ -175,10 +176,10 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
),
)
)
return "thd"
return TokenLayout.T
def compute_step_aux(
self, step_data: dict[str, object], *, layout: str, step: int
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux:
input_ids: torch.Tensor = step_data["input_ids"]

5
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py
View File

@@ -5,6 +5,7 @@ from typing import NamedTuple, Union
from pydantic import model_validator
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import (
Pair,
_check_equal_lengths,
@@ -50,7 +51,7 @@ class TokenAlignerGlobalAux:
step_auxs: dict[int, TokenAlignerStepAux]
framework: str # "sglang" | "megatron"
layout: str # "thd"
layout: TokenLayout
class TokenLocator(_FrozenBase):
@@ -106,7 +107,7 @@ class TokenAlignerSeqsInfo(_FrozenBase):
"""All sequences for one side across all steps."""
sequences: dict[SeqId, TokenAlignerSeqInfo]
layout: str
layout: TokenLayout
class TokenAlignerPlan(_FrozenBase):

5
python/sglang/srt/debug_utils/comparator/dims.py
View File

@@ -8,6 +8,11 @@ BATCH_DIM_NAME: str = "b"
SEQ_DIM_NAME: str = "s"
class TokenLayout(Enum):
T = "t" # single flat token dim
BS = "bs" # separate batch + seq dims, need collapse
class ParallelAxis(Enum):
TP = "tp"
CP = "cp"

1
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
View File

@@ -125,6 +125,7 @@ def _compute_diff(
max_diff_coord=list(max_diff_coord),
baseline_at_max=x_baseline[max_diff_coord].item(),
target_at_max=x_target[max_diff_coord].item(),
diff_threshold=diff_threshold,
passed=(
rel_diff <= diff_threshold
and max_abs_diff <= diff_threshold

5
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
View File

@@ -69,13 +69,10 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list
def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
marker = "✅" if diff.passed else "❌"
return [
prefix_text
+ marker
+ " "
+ "\t".join(
f"{name}={value}"
f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}"
for name, value in [
("rel_diff", diff.rel_diff),
("max_abs_diff", diff.max_abs_diff),

86
python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py
View File

@@ -1,86 +0,0 @@
from sglang.srt.debug_utils.comparator.tensor_comparator.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 = list(TensorStats.model_fields.keys())
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}"
)

1
python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py
View File

@@ -28,6 +28,7 @@ class DiffInfo(_StrictBase):
max_diff_coord: list[int]
baseline_at_max: float
target_at_max: float
diff_threshold: float
passed: bool

15
test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py
View File

@@ -13,6 +13,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SGLangSeqId,
TokenAlignerStepAux,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="default", nightly=True)
@@ -34,7 +35,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.input_ids == [10, 20, 30]
@@ -51,7 +52,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=3
step_data, layout=TokenLayout.T, step=3
)
assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
@@ -65,7 +66,7 @@ class TestNormalizeSGLang:
}
result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
@@ -81,7 +82,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.seq_lens == [3, 2]
@@ -94,7 +95,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.positions == [0, 1, 2, 0, 1]
@@ -108,7 +109,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=0
step_data, layout=TokenLayout.T, step=0
)
assert result.positions == [5, 6, 7, 8, 9]
@@ -121,7 +122,7 @@ class TestNormalizeMegatron:
}
result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
step_data, layout="thd", step=5
step_data, layout=TokenLayout.T, step=5
)
assert result.seq_ids == [
PositionalSeqId(step=5, seq_index=0),

3
test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py
View File

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerStepAux,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
@@ -52,7 +53,7 @@ class TestExecuteAlignment:
side_aux = TokenAlignerGlobalAux(
step_auxs={0: aux, 1: aux_step1},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)

23
test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py
View File

@@ -19,6 +19,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerStepAux,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
@@ -40,7 +41,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -75,7 +76,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -108,7 +109,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -132,7 +133,7 @@ class TestBuildTokenIndexSGLangThd:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -163,7 +164,7 @@ class TestBuildTokenIndexMegatronThd:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -204,7 +205,7 @@ class TestBuildTokenIndexMegatronThd:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index = build_seqs_info(side_aux)
@@ -389,7 +390,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
@@ -401,7 +402,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
index_a = build_seqs_info(side_aux_a)
@@ -428,7 +429,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="sglang",
layout="thd",
layout=TokenLayout.T,
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
@@ -443,7 +444,7 @@ class TestComputeAlignmentPlanCrossLayout:
),
},
framework="megatron",
layout="thd",
layout=TokenLayout.T,
)
index_a = build_seqs_info(side_aux_a)
@@ -467,7 +468,7 @@ def _int_to_seq_id(k: int) -> SeqId:
def _make_index(
*,
sequences: dict[int, tuple[int, ...]],
layout: str = "thd",
layout: TokenLayout = TokenLayout.T,
) -> TokenAlignerSeqsInfo:
"""Create a TokenAlignerSeqsInfo from simplified input_ids-only specification."""
records: dict[SeqId, TokenAlignerSeqInfo] = {}

14
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
View File

@@ -35,6 +35,7 @@ def _make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
diff_threshold: float = 1e-3,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
@@ -44,6 +45,7 @@ def _make_diff(
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=diff_threshold,
passed=passed,
)
@@ -94,7 +96,7 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)
@@ -153,11 +155,11 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"❌ rel_diff=0.002\tmax_abs_diff=0.005\tmean_abs_diff=0.001\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 "
"baseline=1.0 target=1.0005\n"
"When downcast to torch.bfloat16: "
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)
@@ -187,7 +189,7 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)
@@ -215,7 +217,7 @@ class TestFormatComparison:
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005\n"
"x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
@@ -243,7 +245,7 @@ class TestFormatComparison:
"[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\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\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 "
"baseline=1.0 target=1.0005"
)

265
test/registered/debug_utils/comparator/tensor_comparator/test_printer.py
View File

@@ -1,265 +0,0 @@
import sys
import pytest
from sglang.srt.debug_utils.comparator.tensor_comparator.printer import (
print_comparison,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorInfo,
TensorStats,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _make_stats(
mean: float = 0.0,
std: float = 1.0,
min: float = -2.0,
max: float = 2.0,
p1: float | None = -1.8,
p5: float | None = -1.5,
p95: float | None = 1.5,
p99: float | None = 1.8,
) -> TensorStats:
return TensorStats(
mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
)
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],
baseline_at_max=1.0,
target_at_max=1.0005,
passed=passed,
)
def _make_tensor_info(
shape: list[int] | None = None,
dtype: str = "torch.float32",
stats: TensorStats | None = None,
sample: str | None = None,
) -> TensorInfo:
return TensorInfo(
shape=shape if shape is not None else [4, 8],
dtype=dtype,
stats=stats if stats is not None else _make_stats(),
sample=sample,
)
# Snapshot strings below are intentionally spelled out in full per test.
# The shared skeleton (stats block, diff block) looks duplicated, but keeping
# each test self-contained makes failures immediately readable without chasing
# helper functions. Do not extract common fragments.
class TestPrintComparison:
def test_normal(self, capsys):
info = TensorComparisonInfo(
name="test",
baseline=_make_tensor_info(
stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0),
),
target=_make_tensor_info(
stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001),
),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"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"
"[min] -2.0000 vs -2.0001 (diff: -0.0001)\n"
"[max] 2.0000 vs 2.0001 (diff: 0.0001)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"baseline=1.0 target=1.0005\n"
)
def test_shape_mismatch(self, capsys):
info = TensorComparisonInfo(
name="mismatch",
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] [3, 4] vs [5, 6]\t"
"[dtype] torch.float32 vs torch.float32\n"
"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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"⚠️ Shape mismatch\n"
)
def test_with_downcast(self, capsys):
info = TensorComparisonInfo(
name="downcast",
baseline=_make_tensor_info(),
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, passed=False
),
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] [4, 8] vs [4, 8]\t"
"[🟠dtype] torch.float32 vs torch.bfloat16\n"
"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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"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 "
"baseline=1.0 target=1.0005\n"
)
def test_with_shape_unification(self, capsys):
info = TensorComparisonInfo(
name="unify",
baseline=_make_tensor_info(shape=[1, 1, 4, 8]),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [1, 1, 4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"baseline=1.0 target=1.0005\n"
)
def test_with_samples(self, capsys):
info = TensorComparisonInfo(
name="samples",
baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"),
target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\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"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[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 "
"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"
)
def test_none_quantiles(self, capsys):
stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None)
info = TensorComparisonInfo(
name="no_quantiles",
baseline=_make_tensor_info(stats=stats_no_quantiles),
target=_make_tensor_info(stats=stats_no_quantiles),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
print_comparison(info=info, diff_threshold=1e-3)
assert capsys.readouterr().out == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\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"
"[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 "
"baseline=1.0 target=1.0005\n"
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

2
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
View File

@@ -46,6 +46,7 @@ def _make_diff(**overrides) -> DiffInfo:
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=1e-3,
passed=True,
)
defaults.update(overrides)
@@ -76,6 +77,7 @@ class TestStrictBase:
max_diff_coord=[0],
baseline_at_max=0.0,
target_at_max=0.0,
diff_threshold=1e-3,
passed=True,
extra_field=123,
)

1
test/registered/debug_utils/comparator/test_model_validation.py
View File

@@ -185,6 +185,7 @@ def _make_diff_info(*, passed: bool) -> DiffInfo:
max_diff_coord=[0, 0],
baseline_at_max=1.0,
target_at_max=1.01,
diff_threshold=1e-3,
passed=passed,
)