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
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")