[JazyJiang]

No linked issue. This adds shared test tooling; it does not fix a reported bug.

1. Background and problem

Current behaviour. Tests that compare an implementation with a reference use hand-picked atol / rtol per test. Under tests/test_shardformer/test_model alone there are more than ten different pairs, a few loosened to 4e-2 or 9e-2 with no recorded reason.

Problem. The error a correct low-precision implementation makes differs by orders of magnitude between models and between tensors of the same model. A fixed tolerance tight enough for one rejects correct code in another, and a loose one hides real bugs. Once a tolerance has been loosened, nobody can tell whether it absorbs rounding noise or a bug. Each new model or kernel test picks its own number again, and there is no way to compare runs made on different machines (e.g. GPU and NPU).

Expected behaviour. One shared rule whose allowance is measured, not chosen: for every tensor, the implementation under test may deviate from a float64 ground truth by at most a small multiple of how far the reference implementation itself deviates at the same precision.

Why now. New models, kernels and optimizers are being added that all need a reference comparison. Having one tool before they write their tests avoids a new set of fixed tolerances.

2. What changed

New module tests/kit/numerics (test-only, nothing under colossalai/ changes):

• judge.py: judge(test, ref, truth) and assert_aligned(...). Per tensor, baseline error = RMSE(ref, truth), test error = RMSE(test, truth); passes if test error <= multiplier * baseline error + tol / 10, with multiplier 3 for fp16/bf16 and 2 for fp32 (at least 3 below 1000 elements); an allclose(test, ref) fast path passes first. This is the rule PyTorch uses for torch.compile accuracy checks (same(..., fp64_ref=...) in torch/_dynamo/utils.py). Returns a per-tensor report; failures are listed first with their ratio.
• capture.py: run_fwd_bwd(fn, inputs, dtype) runs a function or nn.Module forward and backward at a given dtype with a fixed random upstream gradient, so the float64 truth, the reference and the test see identical inputs, weights and gradients. TF32 is disabled during the run.
• trace.py: Trace, one run as {name: tensor} plus metadata. Saved as a single .safetensors file with the metadata as JSON in the header (no pickle, dtypes preserved), so runs from different machines can be judged offline.
• compare.py: flattening of nested outputs and float64 metrics (RMSE, max abs, relative L2, cosine).
• README.md: the rule, runnable operator and module examples, how to read the report, and what makes a comparison meaningless (different precision recipes, shared code between reference and test).
• tests/test_numerics/test_numerics_core.py: 14 CPU tests. Registered in both CPU test lists: the cpu batch of .github/scripts/ci/run_colossalai_batch.sh (Layer 1) and .github/scripts/ci/run_cpu_unit_tests.sh (Fork-safe PR Gate). The two lists are maintained separately.

Not included in this PR:

• Whole-model alignment that builds models from tests/kit/model_zoo and runs them through booster plugins.
• Data-parallel equivalence and checkpoint-resume checks.
• Migrating any existing test to the new rule. No existing tolerance is changed.

3. User and compatibility impact

Area	Impact
Python API	None (test-only module under tests/kit)
Checkpoint format	None
Default behaviour	None
Model accuracy	None
Performance / memory	None
Python / PyTorch / CUDA versions	None
Custom C++/CUDA operators	None
Docker images	None
Existing training scripts	None

CI: one more file in the CPU unit test job, about 3 s.

4. Test environment

CPU only; no GPU is needed or used by these tests. Two environments, matching the two CI jobs that run CPU tests.

A. Repository CI image (the image used by the GPU batches), on the CI host:

OS	Ubuntu 22.04.5 LTS (container), host Linux 5.15.0, Intel Xeon Platinum 8563C
Python	3.10.12
GPU	none (CUDA_VISIBLE_DEVICES= empty)
PyTorch	2.13.0+cu130 (torch.version.cuda 13.0, not exercised)
Transformers	5.17.0
diffusers	0.35.2
safetensors	0.8.0
pytest	7.4.4
flash-attn	not installed
Image	sha256:ae8ff34bf340a692803edca6e0f38a576fc1a2b206c01b4ef6f7c2a40dfd8c10 (current COLOSSAL_CI_IMAGE)
ColossalAI commit	515e8354 (this branch, on top of cb3b2234)

B. Same packages as the GitHub-hosted CPU jobs: fresh Python 3.10.12 virtualenv with torch==2.5.1 from the PyTorch CPU index, pip install -r requirements/requirements.txt (Transformers 5.17.0, safetensors 0.8.0) and pytest, as in ci_layer_1_pr.yml / fork_pr_checks.yml. Commit 407f6840 (the second commit only adds a list entry).

5. Test commands and results

5.1 Layer 1 cpu batch (environment A)

CUDA_VISIBLE_DEVICES= bash .github/scripts/ci/run_colossalai_batch.sh cpu

• 127 tests: 124 passed, 0 failed, 0 errors, 3 skipped; exit code 0; 17.5 s
• 14 of them are the new tests/test_numerics cases, all passed
• The 3 skips are pre-existing CUDA-guarded cases in tests/test_checkpoint_io/test_optimizer_safetensors_v2.py ("Requires a guarded CUDA device")
• FAST_TEST=1 (set by the batch script); it does not affect these tests. No rerun, no xfail, no ignore
• JUnit XML and log written to $COLOSSAL_RESULTS by the batch script

5.2 CI helper tests (environment A)

python -m pytest -q .github/scripts/ci/

• 66 passed, exit code 0 (checks the batch scripts themselves)

5.3 GitHub CPU job packages (environment B)

COLOSSAL_PYTHON=<venv>/bin/python bash .github/scripts/ci/run_cpu_unit_tests.sh
<venv>/bin/python -m pytest -q tests/test_numerics

• run_cpu_unit_tests.sh: exit code 0; 2 skipped (the same CUDA-guarded cases); the total count was not recorded in this run
• tests/test_numerics: 14 passed in 1.34 s

5.4 README examples

Every Python block in tests/kit/numerics/README.md was extracted and executed in order in environment A: the attention example passes assert_aligned, the module example reports 4/4 tensors aligned (max ratio 1.33), and the saved traces reload and judge identically.

6. Known issues

None introduced by this PR.

• In an older experimental CI image (tf5.17.0-exp-20261019, not the current COLOSSAL_CI_IMAGE), tests/test_infer/test_async_engine/test_request_tracer.py fails at collection because its diffusers imports FLAX_WEIGHTS_NAME, which Transformers 5.17 removed. Reproduced on main (cb3b2234) without this change, and absent in the current image (diffusers 0.35.2). Not addressed here.
• GPU paths are not exercised: the module only runs what the caller passes in, and its own tests are CPU-only.

7. Risks and rollback

Risk. A multiplier that is too tight for some operator would make a correct kernel fail. The self-tests leave a wide margin: a legitimate reordering of the summation reaches a ratio of at most 1.18 (limit 3), while dropping one column of the reduction dimension gives at least 20 in bf16 and over 10^5 in fp32.

Rollback. Revert the commit. Nothing outside tests/ and one line of the CPU test list depends on it.

8. Docs and release impact

• README: added for the new module (tests/kit/numerics/README.md).
• Installation, API and Docker docs, changelog, version.txt, PyPI, DockerHub: none.