From 407f684095d18b5f9eb0e78e006e49970f4e8768 Mon Sep 17 00:00:00 2001 From: JazyJiang Date: Tue, 27 Oct 2026 17:22:43 +0800 Subject: [PATCH 1/2] [test] add numerical alignment core with baseline-error rule Add tests/kit/numerics, a shared tool for judging whether an implementation matches a reference numerically. The rule follows PyTorch TorchInductor accuracy check: the baseline error is how far the reference implementation deviates from a float64 truth at the tested precision, and the tested output may deviate at most 3x (fp16/bf16) or 2x (fp32) as much, with a larger allowance for small tensors. - compare.py: flatten nested outputs, float64 metrics (rmse, max abs, rel l2, cosine) - judge.py: judge / assert_aligned with a per-tensor report - capture.py: run_fwd_bwd with a fixed upstream gradient for op/module checks - README.md: the rule, runnable examples, how to read the report - trace.py: Trace, one run as {name: tensor} plus metadata, saved as a single safetensors file so runs from different machines can be judged offline The unit tests check that a reordered summation passes while a dropped column fails, that saved traces round-trip bit-exactly, and run on CPU in the CPU unit test job. --- .github/scripts/ci/run_cpu_unit_tests.sh | 1 + tests/kit/numerics/README.md | 171 ++++++++++++++++++++++ tests/kit/numerics/__init__.py | 19 +++ tests/kit/numerics/capture.py | 71 +++++++++ tests/kit/numerics/compare.py | 66 +++++++++ tests/kit/numerics/judge.py | 156 ++++++++++++++++++++ tests/kit/numerics/trace.py | 79 ++++++++++ tests/test_numerics/test_numerics_core.py | 152 +++++++++++++++++++ 8 files changed, 715 insertions(+) create mode 100644 tests/kit/numerics/README.md create mode 100644 tests/kit/numerics/__init__.py create mode 100644 tests/kit/numerics/capture.py create mode 100644 tests/kit/numerics/compare.py create mode 100644 tests/kit/numerics/judge.py create mode 100644 tests/kit/numerics/trace.py create mode 100644 tests/test_numerics/test_numerics_core.py diff --git a/.github/scripts/ci/run_cpu_unit_tests.sh b/.github/scripts/ci/run_cpu_unit_tests.sh index ae22911eb0ec..777743d83702 100644 --- a/.github/scripts/ci/run_cpu_unit_tests.sh +++ b/.github/scripts/ci/run_cpu_unit_tests.sh @@ -14,6 +14,7 @@ TESTS=( tests/test_config/test_load_config.py tests/test_infer/test_async_engine/test_request_tracer.py tests/test_lazy/test_chatglm_cpu.py + tests/test_numerics/test_numerics_core.py tests/test_optimizer/test_lr_scheduler.py tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py diff --git a/tests/kit/numerics/README.md b/tests/kit/numerics/README.md new file mode 100644 index 000000000000..ad5bad5b09ae --- /dev/null +++ b/tests/kit/numerics/README.md @@ -0,0 +1,171 @@ +# Numerical alignment + +A shared tool for checking that an implementation (a fused kernel, a sharded module, a plugin) computes the same thing +as a reference implementation, allowing only the error that the precision itself explains. + +## The rule: baseline error + +Run the same inputs and weights three times: + +| Run | What runs | Role | +| --- | --- | --- | +| `truth` | reference implementation in float64 | ground truth | +| `ref` | reference implementation at the tested precision | measures the **baseline error**: how far correct code drifts from the truth at this precision | +| `test` | implementation under test, same precision recipe as `ref` | what is being checked | + +For every tensor (outputs, input gradients, parameter gradients): + +``` +baseline error = RMSE(ref, truth) +test error = RMSE(test, truth) +pass if test error <= multiplier * baseline error + tol / 10 +``` + +`multiplier` is 3 for fp16 / bf16 and 2 for fp32, and at least 3 for tensors with fewer than 1000 elements. A `test` +that is already `allclose` to `ref` (`atol = rtol = tol`, default `1e-4`) passes directly. This is the rule PyTorch +uses to check `torch.compile` across its model suites (`same()` in `torch/_dynamo/utils.py`). + +Why not a fixed tolerance: the error a correct bf16 implementation makes differs by orders of magnitude between models +and tensors. A tolerance tight enough for one model rejects another; a loose one hides real bugs. The baseline error +is measured per tensor, so each tensor gets its own allowance. + +## Quick start: an operator + +Check that `scaled_dot_product_attention` matches a plain attention written from the formula: + +```python +import torch +import torch.nn.functional as F + +from tests.kit.numerics import assert_aligned, run_fwd_bwd + + +def eager_attention(q, k, v): + scores = q @ k.transpose(-1, -2) / q.shape[-1] ** 0.5 + return torch.softmax(scores, dim=-1) @ v + + +def sdpa_attention(q, k, v): + return F.scaled_dot_product_attention(q, k, v) + + +g = torch.Generator().manual_seed(0) +inputs = {name: torch.randn(2, 4, 128, 64, generator=g, dtype=torch.float64) for name in ("q", "k", "v")} + +truth = run_fwd_bwd(eager_attention, inputs, torch.float64) +ref = run_fwd_bwd(eager_attention, inputs, torch.bfloat16) +test = run_fwd_bwd(sdpa_attention, inputs, torch.bfloat16) +assert_aligned(test, ref, truth) +``` + +`run_fwd_bwd(fn, inputs, dtype)` casts the floating-point inputs to `dtype`, runs `fn(**inputs)`, backpropagates a +fixed random upstream gradient (the same in all three runs) and returns a `Trace` with `out/...` for the outputs and +`grad/` for the input gradients. TF32 is disabled during the run. + +## A module + +Pass an `nn.Module` instead of a function. It is deep-copied and cast for each run, so the caller's module is not +modified and every run starts from the same weights. Parameter gradients are recorded as `grad/param/`, so the +module under test must use the same parameter names; load the reference weights into it: + +```python +class SplitKLinear(torch.nn.Linear): + """Same math as nn.Linear, summed in two halves like a sharded kernel would.""" + + def forward(self, x): + h = self.in_features // 2 + return x[..., :h] @ self.weight[:, :h].t() + x[..., h:] @ self.weight[:, h:].t() + self.bias + + +torch.manual_seed(0) +ref_module = torch.nn.Linear(256, 128) +test_module = SplitKLinear(256, 128) +test_module.load_state_dict(ref_module.state_dict()) + +inputs = {"input": torch.randn(64, 256, dtype=torch.float64)} +truth = run_fwd_bwd(ref_module, inputs, torch.float64) +ref = run_fwd_bwd(ref_module, inputs, torch.bfloat16) +test = run_fwd_bwd(test_module, {"x": inputs["input"]}, torch.bfloat16) +``` + +Here the two modules name their input differently (`input` vs `x`), so the input gradients would not match by name. +Rename or drop entries before judging; `Trace(tensors, meta)` builds a new trace from a dict: + +```python +from tests.kit.numerics import Trace, judge + +test = Trace({name.replace("grad/x", "grad/input"): t for name, t in test.items()}, test.meta) +print(judge(test, ref, truth)) +``` + +## Reading the report + +`judge(test, ref, truth)` returns an `AlignmentReport`; `assert_aligned` raises `AssertionError` with the same table. +The two reports below come from the linear-layer cases in `tests/test_numerics`. A correct reordering of the +summation: + +``` +3/3 tensors aligned +tensor baseline test_err ratio verdict +out 4.74e-02 5.58e-02 1.18 PASS +grad/w 2.34e-02 2.34e-02 1.00 PASS +grad/x 3.27e-02 3.27e-02 1.00 PASS +``` + +A sharding bug that drops one column of the reduction dimension: + +``` +0/3 tensors aligned +tensor baseline test_err ratio verdict +grad/w 2.34e-02 5.02e-01 21.49 FAIL (> 3x baseline error) +out 4.74e-02 9.81e-01 20.70 FAIL (> 3x baseline error) +grad/x 3.27e-02 6.72e-01 20.52 FAIL (> 3x baseline error) +``` + +`baseline` and `test_err` are RMSE against the float64 truth, `ratio` is their quotient. Failures are listed first, +largest ratio first. Other verdicts: `missing in test` / `missing in truth` (names differ), `shape ...`, +`non-finite test output`, and `non-float tensors differ` for integer tensors, which must match exactly. Each entry +also carries `metrics` with max abs error, relative L2 error and cosine similarity. + +## Comparing runs from different machines + +A `Trace` is a read-only mapping from tensor name to tensor plus a `meta` dict (dtype, function, torch and +transformers versions, anything passed as keyword arguments). It saves to a single `.safetensors` file with the +metadata as JSON in the header, so a run on one machine can be judged against runs from another: + +```python +truth.save("truth.safetensors") # e.g. on the GPU machine +ref.save("ref.safetensors") +test.save("test.safetensors") # e.g. on the NPU machine + +report = judge(Trace.load("test.safetensors"), Trace.load("ref.safetensors"), Trace.load("truth.safetensors")) +``` + +`Trace.capture(outputs, **meta)` turns any tensor, nested dict / list / tuple or HF `ModelOutput` into a trace. + +## Things that make the comparison meaningless + +- **Different precision recipes.** `ref` and `test` must reduce precision in the same places. Autocast (fp32 weights, + bf16 matmuls, fp32 softmax and norms) has a smaller baseline error than casting the whole model to bf16, so judging + a whole-cast `test` against an autocast `ref` reports failures in correct code. +- **Sharing code between `ref` and `test`.** A bug in shared code appears in both and is never caught. The reference + must be an independent implementation (e.g. HF eager modeling code), not a call into the code under test. +- **A reference that cannot run in float64.** Fused kernels such as flash attention usually have no fp64 path; use a + plain implementation as the reference and test the fused one. +- **Randomness inside the function.** Generate noise, timesteps, dropout masks and the like once, outside, and pass + them in `inputs` so all three runs see the same values. + +## API + +| Name | Purpose | +| --- | --- | +| `assert_aligned(test, ref, truth, tol=1e-4)` | judge and raise with the report if anything is not aligned | +| `judge(test, ref, truth, tol=1e-4, small_tensor_multiplier=False)` | return an `AlignmentReport` | +| `run_fwd_bwd(fn, inputs, dtype, grad_seed=0, **meta)` | run a function or module forward and backward, return a `Trace` | +| `Trace`, `Trace.capture`, `Trace.save`, `Trace.load` | one run as named tensors plus metadata | +| `error_multiplier(dtype, numel)` | the multiplier used for a tensor | +| `tensor_metrics(actual, expected)` | RMSE, max abs, relative L2 and cosine, in float64 | +| `cast_floating`, `full_precision_matmul`, `flatten_tensors` | helpers used by the above | + +Whole-model alignment that builds models from `tests/kit/model_zoo` and runs them through booster plugins is not part +of this module yet. diff --git a/tests/kit/numerics/__init__.py b/tests/kit/numerics/__init__.py new file mode 100644 index 000000000000..a72b95f72c24 --- /dev/null +++ b/tests/kit/numerics/__init__.py @@ -0,0 +1,19 @@ +from .capture import cast_floating, full_precision_matmul, run_fwd_bwd +from .compare import TensorMetrics, flatten_tensors, tensor_metrics +from .judge import AlignmentReport, Entry, assert_aligned, error_multiplier, judge +from .trace import Trace + +__all__ = [ + "AlignmentReport", + "Entry", + "TensorMetrics", + "Trace", + "assert_aligned", + "cast_floating", + "error_multiplier", + "flatten_tensors", + "full_precision_matmul", + "judge", + "run_fwd_bwd", + "tensor_metrics", +] diff --git a/tests/kit/numerics/capture.py b/tests/kit/numerics/capture.py new file mode 100644 index 000000000000..760307d00272 --- /dev/null +++ b/tests/kit/numerics/capture.py @@ -0,0 +1,71 @@ +"""Helpers to produce the three runs (float64 truth, reference, test) for operator- and module-level checks.""" + +import copy +from contextlib import contextmanager +from typing import Any, Callable, Dict + +import torch +from torch import Tensor, nn +from torch.utils._pytree import tree_map + +from .compare import flatten_tensors +from .trace import Trace + + +def cast_floating(obj: Any, dtype: torch.dtype) -> Any: + """Cast every floating-point tensor nested in ``obj`` to ``dtype``; other leaves are returned unchanged.""" + return tree_map(lambda x: x.to(dtype) if isinstance(x, Tensor) and x.is_floating_point() else x, obj) + + +@contextmanager +def full_precision_matmul(): + """Disable TF32 so that float32 runs are really float32 and the float64 truth is not silently degraded.""" + matmul_precision = torch.get_float32_matmul_precision() + cudnn_tf32 = torch.backends.cudnn.allow_tf32 + torch.set_float32_matmul_precision("highest") + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.set_float32_matmul_precision(matmul_precision) + torch.backends.cudnn.allow_tf32 = cudnn_tf32 + + +def run_fwd_bwd(fn: Callable, inputs: Dict[str, Any], dtype: torch.dtype, grad_seed: int = 0, **meta) -> Trace: + """Run ``fn(**inputs)`` in ``dtype`` and backpropagate a fixed random upstream gradient. + + If ``fn`` is an ``nn.Module`` it is deep-copied before casting, so every run starts from the same weights. + The upstream gradient is drawn in float64 from ``grad_seed`` and cast per output, so all runs see the same one. + + Returns a ``Trace`` with ``{"out/...": output, "grad/": input grad, "grad/param/": param + grad}``; ``meta`` (e.g. ``role="truth"``) is recorded alongside dtype, function name and seed. + """ + fn_name = type(fn).__name__ if isinstance(fn, nn.Module) else getattr(fn, "__qualname__", repr(fn)) + if isinstance(fn, nn.Module): + fn = copy.deepcopy(fn).to(dtype) + leaves = {} + for name, value in cast_floating(inputs, dtype).items(): + if isinstance(value, Tensor) and value.is_floating_point(): + value = value.detach().clone().requires_grad_(True) + leaves[name] = value + + with full_precision_matmul(): + outputs = flatten_tensors(fn(**leaves), "out") + generator = torch.Generator().manual_seed(grad_seed) + loss = 0 + for out in outputs.values(): + if out.requires_grad: + upstream = torch.randn(out.shape, generator=generator, dtype=torch.float64).to(out) + loss = loss + (out * upstream).sum() + if isinstance(loss, Tensor): + loss.backward() + + results = {name: out.detach() for name, out in outputs.items()} + for name, value in leaves.items(): + if isinstance(value, Tensor) and value.grad is not None: + results[f"grad/{name}"] = value.grad + if isinstance(fn, nn.Module): + for name, param in fn.named_parameters(): + if param.grad is not None: + results[f"grad/param/{name}"] = param.grad + return Trace.capture(results, fn=fn_name, dtype=str(dtype), grad_seed=grad_seed, **meta) diff --git a/tests/kit/numerics/compare.py b/tests/kit/numerics/compare.py new file mode 100644 index 000000000000..b8be5e1dc186 --- /dev/null +++ b/tests/kit/numerics/compare.py @@ -0,0 +1,66 @@ +"""Tensor-level comparison for numerical alignment tests. + +All metrics are computed in float64 on CPU so that the measurement itself adds no rounding error. +""" + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Dict + +import torch +from torch import Tensor + + +@dataclass +class TensorMetrics: + rmse: float + max_abs: float + rel_l2: float + cosine: float + + +def flatten_tensors(obj: Any, prefix: str = "") -> Dict[str, Tensor]: + """Flatten tensors nested in mappings (dicts, HF ``ModelOutput``, ``Trace``), lists and tuples into ``{"a/b/0": tensor}``. + + Non-tensor leaves such as ``None`` are dropped. + """ + if isinstance(obj, Tensor): + return {prefix or "output": obj} + if isinstance(obj, Mapping): + items = obj.items() + elif isinstance(obj, (list, tuple)): + items = enumerate(obj) + else: + return {} + flat = {} + for key, value in items: + flat.update(flatten_tensors(value, f"{prefix}/{key}" if prefix else str(key))) + return flat + + +def as_float64(tensor: Tensor) -> Tensor: + return tensor.detach().to(device="cpu", dtype=torch.float64) + + +def tensor_metrics(actual: Tensor, expected: Tensor) -> TensorMetrics: + """Distance of ``actual`` from ``expected``. NaN or inf in either input propagates into the metrics.""" + a, e = as_float64(actual).flatten(), as_float64(expected).flatten() + if a.numel() == 0: + return TensorMetrics(0.0, 0.0, 0.0, 1.0) + diff = a - e + a_norm, e_norm, diff_norm = a.norm().item(), e.norm().item(), diff.norm().item() + if e_norm > 0: + rel_l2 = diff_norm / e_norm + else: + rel_l2 = 0.0 if diff_norm == 0 else math.inf + if a_norm > 0 and e_norm > 0: + cosine = (a @ e).item() / (a_norm * e_norm) + else: + cosine = 1.0 if a_norm == e_norm else 0.0 + return TensorMetrics( + rmse=diff.pow(2).mean().sqrt().item(), + max_abs=diff.abs().max().item(), + rel_l2=rel_l2, + cosine=cosine, + ) diff --git a/tests/kit/numerics/judge.py b/tests/kit/numerics/judge.py new file mode 100644 index 000000000000..c09090df370c --- /dev/null +++ b/tests/kit/numerics/judge.py @@ -0,0 +1,156 @@ +"""Pass/fail rules for numerical alignment. + +The baseline-error rule follows PyTorch's TorchInductor accuracy check (``same()`` in ``torch/_dynamo/utils.py``): +the baseline error is how far the reference implementation itself deviates from a float64 ground truth at the +precision being tested, and the implementation under test may deviate at most ``multiplier`` times as much. +""" + +import math +from dataclasses import dataclass, field +from typing import Any, List, Optional + +import torch + +from .compare import TensorMetrics, as_float64, flatten_tensors, tensor_metrics + +LOW_PRECISION_DTYPES = (torch.float16, torch.bfloat16) + + +def error_multiplier(dtype: torch.dtype, numel: int, small_tensor_multiplier: bool = False) -> float: + """Allowed ratio between the test error and the baseline error, as chosen by PyTorch. + + 3 for fp16/bf16 (2 gave false alarms under AMP), 2 otherwise; small tensors are noisier and get at least 3, + or 8-10 with ``small_tensor_multiplier``. + """ + multiplier = 3.0 if dtype in LOW_PRECISION_DTYPES else 2.0 + if small_tensor_multiplier and numel <= 10: + return 10.0 + if small_tensor_multiplier and numel < 1000: + return 8.0 + if numel < 1000: + return max(multiplier, 3.0) + return multiplier + + +@dataclass +class Entry: + name: str + passed: bool + reason: str + baseline_error: float = math.nan + test_error: float = math.nan + allowed_error: float = math.nan + metrics: Optional[TensorMetrics] = None + + @property + def ratio(self) -> float: + if self.baseline_error == 0: + return 0.0 if self.test_error == 0 else math.inf + return self.test_error / self.baseline_error + + +@dataclass +class AlignmentReport: + entries: List[Entry] = field(default_factory=list) + + @property + def failures(self) -> List[Entry]: + return [e for e in self.entries if not e.passed] + + @property + def passed(self) -> bool: + return not self.failures + + def format(self, limit: int = 20) -> str: + rows = sorted(self.entries, key=lambda e: (e.passed, -e.ratio if not math.isnan(e.ratio) else 0.0)) + lines = [ + f"{len(self.entries) - len(self.failures)}/{len(self.entries)} tensors aligned", + f"{'tensor':<56} {'baseline':>9} {'test_err':>9} {'ratio':>8} verdict", + ] + for e in rows[:limit]: + verdict = "PASS" if e.passed else f"FAIL ({e.reason})" + lines.append(f"{e.name:<56} {e.baseline_error:9.2e} {e.test_error:9.2e} {e.ratio:8.2f} {verdict}") + if len(rows) > limit: + lines.append(f"... {len(rows) - limit} more") + return "\n".join(lines) + + def __str__(self) -> str: + return self.format() + + def assert_pass(self) -> None: + assert self.passed, "numerical alignment failed\n" + self.format() + + +def judge( + test: Any, + ref: Any, + truth: Any, + tol: float = 1e-4, + small_tensor_multiplier: bool = False, +) -> AlignmentReport: + """Judge ``test`` against a float64 ``truth``, allowing ``multiplier`` times the baseline error. + + The baseline error is the error of ``ref`` against ``truth``: what correct code naturally gets wrong at this + precision. + + Args: + test: Outputs of the implementation under test, at the precision being tested. + ref: Outputs of the reference implementation run with the same precision recipe as ``test``. + truth: Outputs of the reference implementation in float64. + tol: Fast-path tolerance: ``test`` within ``allclose(ref, atol=tol, rtol=tol)`` passes outright; + ``tol / 10`` is also added to the allowed error, as in PyTorch. + + Each argument may be a tensor or any nesting of dicts / lists / tuples / ``ModelOutput``; tensors are matched + by their flattened names (see ``flatten_tensors``). + """ + test_flat, ref_flat, truth_flat = flatten_tensors(test), flatten_tensors(ref), flatten_tensors(truth) + report = AlignmentReport() + for name in sorted(set(truth_flat) - set(test_flat)): + report.entries.append(Entry(name, False, "missing in test")) + for name in sorted(set(test_flat) - set(truth_flat)): + report.entries.append(Entry(name, False, "missing in truth")) + for name in sorted(set(truth_flat) & set(test_flat)): + report.entries.append( + _judge_one(name, test_flat[name], ref_flat.get(name), truth_flat[name], tol, small_tensor_multiplier) + ) + return report + + +def _judge_one(name, test, ref, truth, tol, small_tensor_multiplier) -> Entry: + if test.shape != truth.shape: + return Entry(name, False, f"shape {tuple(test.shape)} vs {tuple(truth.shape)}") + if not test.is_floating_point(): + equal = bool(torch.equal(test.cpu(), truth.cpu())) + return Entry(name, equal, "exact" if equal else "non-float tensors differ", 0.0, 0.0, 0.0) + if ref is None: + return Entry(name, False, "missing in ref") + if ref.shape != truth.shape: + return Entry(name, False, f"ref shape {tuple(ref.shape)} vs {tuple(truth.shape)}") + + test64, ref64, truth64 = as_float64(test), as_float64(ref), as_float64(truth) + metrics = tensor_metrics(test64, truth64) + baseline_error, test_error = tensor_metrics(ref64, truth64).rmse, metrics.rmse + multiplier = error_multiplier(test.dtype, test.numel(), small_tensor_multiplier) + allowed = multiplier * baseline_error + tol / 10.0 + entry = Entry(name, False, "", baseline_error, test_error, allowed, metrics) + + if not torch.isfinite(truth64).all(): + entry.reason = "non-finite truth, rerun the reference in higher precision" + elif not torch.isfinite(test64).all(): + entry.reason = "non-finite test output" + elif torch.allclose(test64, ref64, atol=tol, rtol=tol): + entry.passed, entry.reason = True, "allclose to ref" + elif math.isnan(baseline_error): + entry.reason = "non-finite ref output" + elif test_error <= allowed: + entry.passed, entry.reason = True, "within baseline" + else: + entry.reason = f"> {multiplier:g}x baseline error" + return entry + + +def assert_aligned(test: Any, ref: Any, truth: Any, **kwargs) -> AlignmentReport: + """``judge`` and raise ``AssertionError`` with a per-tensor table if anything is not aligned.""" + report = judge(test, ref, truth, **kwargs) + report.assert_pass() + return report diff --git a/tests/kit/numerics/trace.py b/tests/kit/numerics/trace.py new file mode 100644 index 000000000000..b766316c593c --- /dev/null +++ b/tests/kit/numerics/trace.py @@ -0,0 +1,79 @@ +"""Trace: the record of one run, ``{name: tensor}`` plus metadata, that the compare and judge layers consume. + +A trace is saved as a single ``.safetensors`` file with its metadata as JSON in the file header, so runs made on +different machines (e.g. GPU and NPU) can be compared offline without unpickling anything. +""" + +import json +import os +import platform +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Any, Dict, Iterator, Optional, Union + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from torch import Tensor + +from .compare import flatten_tensors + +META_KEY = "numerics_trace" + + +def environment_meta() -> Dict[str, Any]: + meta = { + "created": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "host": platform.node(), + "python": platform.python_version(), + "torch": torch.__version__, + } + try: + import transformers + + meta["transformers"] = transformers.__version__ + except ImportError: + pass + return meta + + +class Trace(Mapping): + """Read-only mapping of tensor name to tensor, with a ``meta`` dict describing how the run was made. + + Tensors are detached and kept on CPU in their original dtype. Anything accepted by ``judge``, including another + trace, can be turned into a trace with ``Trace.capture``. + """ + + def __init__(self, tensors: Dict[str, Tensor], meta: Optional[Dict[str, Any]] = None): + self._tensors = {name: t.detach().cpu() for name, t in tensors.items()} + self.meta = dict(meta or {}) + + @classmethod + def capture(cls, outputs: Any, **meta) -> "Trace": + """Flatten ``outputs`` (tensor, nested dicts / lists, ``ModelOutput``) into a trace; ``meta`` is recorded.""" + return cls(flatten_tensors(outputs), {**environment_meta(), **meta}) + + def save(self, path: Union[str, os.PathLike]) -> None: + tensors = {name: t.contiguous().clone() for name, t in self._tensors.items()} + save_file(tensors, path, metadata={META_KEY: json.dumps(self.meta, default=str)}) + + @classmethod + def load(cls, path: Union[str, os.PathLike]) -> "Trace": + with safe_open(path, framework="pt") as f: + header = f.metadata() or {} + if META_KEY not in header: + raise ValueError(f"{path} is a safetensors file but not a numerics trace") + tensors = {name: f.get_tensor(name) for name in f.keys()} + return cls(tensors, json.loads(header[META_KEY])) + + def __getitem__(self, name: str) -> Tensor: + return self._tensors[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._tensors) + + def __len__(self) -> int: + return len(self._tensors) + + def __repr__(self) -> str: + return f"Trace({len(self)} tensors, meta={self.meta})" diff --git a/tests/test_numerics/test_numerics_core.py b/tests/test_numerics/test_numerics_core.py new file mode 100644 index 000000000000..43268349aceb --- /dev/null +++ b/tests/test_numerics/test_numerics_core.py @@ -0,0 +1,152 @@ +import math + +import pytest +import torch +from safetensors.torch import save_file + +from tests.kit.numerics import Trace, assert_aligned, error_multiplier, flatten_tensors, judge, run_fwd_bwd + +HIDDEN = 256 + + +def linear(x, w): + return x @ w.t() + + +def linear_split_k(x, w): + # same math, different summation order: what a fused or sharded kernel legitimately does + half = HIDDEN // 2 + return x[:, :half] @ w[:, :half].t() + x[:, half:] @ w[:, half:].t() + + +def linear_off_by_one(x, w): + # a typical sharding bug: one column of the reduction dimension is dropped + half = HIDDEN // 2 + return x[:, :half] @ w[:, :half].t() + x[:, half + 1 :] @ w[:, half + 1 :].t() + + +def _inputs(): + g = torch.Generator().manual_seed(0) + return { + "x": torch.randn(64, HIDDEN, generator=g, dtype=torch.float64), + "w": torch.randn(128, HIDDEN, generator=g, dtype=torch.float64), + } + + +def _three_runs(test_fn, dtype=torch.bfloat16): + inputs = _inputs() + truth = run_fwd_bwd(linear, inputs, torch.float64) + ref = run_fwd_bwd(linear, inputs, dtype) + test = run_fwd_bwd(test_fn, inputs, dtype) + return test, ref, truth + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_reordered_summation_is_within_baseline(dtype): + test, ref, truth = _three_runs(linear_split_k, dtype) + report = assert_aligned(test, ref, truth) + assert {e.name for e in report.entries} == {"out", "grad/x", "grad/w"} + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_dropped_column_is_caught(dtype): + test, ref, truth = _three_runs(linear_off_by_one, dtype) + report = judge(test, ref, truth) + failed = {e.name for e in report.failures} + assert {"out", "grad/x", "grad/w"} <= failed + assert all(e.ratio > 3 for e in report.failures) + with pytest.raises(AssertionError, match="grad/w"): + report.assert_pass() + + +def test_baseline_must_come_from_the_same_precision_recipe(): + # a float32 "reference" has a much smaller baseline error than bf16, so a correct bf16 kernel looks broken against it + inputs = _inputs() + truth = run_fwd_bwd(linear, inputs, torch.float64) + mismatched_ref = run_fwd_bwd(linear, inputs, torch.float32) + test = run_fwd_bwd(linear_split_k, inputs, torch.bfloat16) + assert not judge(test, mismatched_ref, truth).passed + + +def test_module_params_get_gradients_and_same_weights(): + torch.manual_seed(0) + mlp = torch.nn.Sequential(torch.nn.Linear(32, 64), torch.nn.GELU(), torch.nn.Linear(64, 16)) + inputs = {"input": torch.randn(8, 32, dtype=torch.float64)} + truth = run_fwd_bwd(mlp, inputs, torch.float64) + ref = run_fwd_bwd(mlp, inputs, torch.bfloat16) + assert "grad/param/0.weight" in truth and "grad/input" in truth + assert_aligned(ref, ref, truth) + assert next(mlp.parameters()).dtype == torch.float32 # the caller's module is not modified + + +def test_error_multiplier_follows_pytorch(): + assert error_multiplier(torch.bfloat16, 10**6) == 3.0 + assert error_multiplier(torch.float16, 10**6) == 3.0 + assert error_multiplier(torch.float32, 10**6) == 2.0 + assert error_multiplier(torch.float32, 100) == 3.0 + assert error_multiplier(torch.float32, 100, small_tensor_multiplier=True) == 8.0 + assert error_multiplier(torch.bfloat16, 4, small_tensor_multiplier=True) == 10.0 + + +def test_structural_mismatches_fail_with_reason(): + t = torch.ones(4) + report = judge({"a": t, "extra": t}, {"a": t}, {"a": t, "b": t}) + reasons = {e.name: e.reason for e in report.failures} + assert reasons == {"b": "missing in test", "extra": "missing in truth"} + + report = judge({"a": torch.ones(2, 2)}, {"a": torch.ones(4)}, {"a": torch.ones(4)}) + assert report.failures[0].reason.startswith("shape") + + +def test_non_finite_and_integer_tensors(): + truth = {"y": torch.ones(8, dtype=torch.float64), "idx": torch.arange(8)} + ref = {"y": torch.ones(8, dtype=torch.bfloat16), "idx": torch.arange(8)} + nan_test = {"y": torch.full((8,), math.nan, dtype=torch.bfloat16), "idx": torch.arange(8)} + report = judge(nan_test, ref, truth) + assert [e.reason for e in report.failures] == ["non-finite test output"] + + wrong_idx = {"y": torch.ones(8, dtype=torch.bfloat16), "idx": torch.arange(8).flip(0)} + assert [e.name for e in judge(wrong_idx, ref, truth).failures] == ["idx"] + + +def test_flatten_tensors_names(): + nested = {"logits": torch.ones(1), "past": [torch.ones(1), None, (torch.ones(1),)]} + assert sorted(flatten_tensors(nested)) == ["logits", "past/0", "past/2/0"] + assert list(flatten_tensors(torch.ones(1))) == ["output"] + + +def test_trace_roundtrip_is_bit_exact(tmp_path): + traces = { + "truth": run_fwd_bwd(linear, _inputs(), torch.float64, role="truth"), + "ref": run_fwd_bwd(linear, _inputs(), torch.bfloat16, role="ref"), + "mixed": Trace.capture({"idx": torch.arange(4), "y": torch.ones(2, dtype=torch.float16)}, role="mixed"), + } + for role, trace in traces.items(): + path = tmp_path / f"{role}.safetensors" + trace.save(path) + loaded = Trace.load(path) + assert loaded.meta == trace.meta + assert set(loaded) == set(trace) + for name in trace: + assert loaded[name].dtype == trace[name].dtype + assert torch.equal(loaded[name], trace[name]) + assert traces["ref"].meta["dtype"] == "torch.bfloat16" + assert {"created", "torch", "transformers", "fn"} <= set(traces["ref"].meta) + + +def test_judging_saved_traces_matches_in_memory(tmp_path): + runs = _three_runs(linear_off_by_one) + loaded = [] + for i, trace in enumerate(runs): + trace.save(tmp_path / f"{i}.safetensors") + loaded.append(Trace.load(tmp_path / f"{i}.safetensors")) + in_memory, offline = judge(*runs), judge(*loaded) + assert [(e.name, e.passed, e.ratio) for e in in_memory.entries] == [ + (e.name, e.passed, e.ratio) for e in offline.entries + ] + + +def test_load_rejects_plain_safetensors(tmp_path): + save_file({"a": torch.ones(1)}, tmp_path / "plain.safetensors") + with pytest.raises(ValueError, match="not a numerics trace"): + Trace.load(tmp_path / "plain.safetensors") From 515e8354a044619af579147b9916a0141086444b Mon Sep 17 00:00:00 2001 From: JazyJiang Date: Tue, 27 Oct 2026 18:15:24 +0800 Subject: [PATCH 2/2] [ci] run the numerics core tests in the Layer 1 cpu batch The Layer 1 CPU job calls run_colossalai_batch.sh cpu, which keeps its own test list separate from run_cpu_unit_tests.sh used by the fork-safe gate. Without this entry the new tests would only run in the fork-safe gate. --- .github/scripts/ci/run_colossalai_batch.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/scripts/ci/run_colossalai_batch.sh b/.github/scripts/ci/run_colossalai_batch.sh index a44385f993f9..f030d4ae76ef 100644 --- a/.github/scripts/ci/run_colossalai_batch.sh +++ b/.github/scripts/ci/run_colossalai_batch.sh @@ -76,6 +76,7 @@ case "${BATCH}" in tests/test_config/test_load_config.py tests/test_infer/test_async_engine/test_request_tracer.py tests/test_lazy/test_chatglm_cpu.py + tests/test_numerics/test_numerics_core.py tests/test_optimizer/test_lr_scheduler.py tests/test_pipeline/test_pipeline_utils/test_t5_pipeline_utils.py tests/test_pipeline/test_pipeline_utils/test_whisper_pipeline_utils.py