[JazyJiang]

1. Related issues

Related to #6488: the items left open after #6505 (see the status table in #6488 (comment)). Follows #6505.

2. Background

#6505 added whole-model, data-parallel and resume checks for model_zoo entries. Its review against #6488 left open: an optimizer coverage check and tied-parameter declarations (D2), a device audit next to the dtype audit (D3), steps and micro-batches without valid tokens (D4), a direct comparison of the optimizer and scheduler state on resume, a stated RNG / data-progress boundary, and not holding several models on the device at once (D5), and gradient checkpointing in the support matrix (D0).

3. Changes

All under tests/kit/numerics and tests/test_numerics; nothing in colossalai/ or CI changes.

• Parameter contract, checked in every run and failed by judge.require_gradients, with the logical parameter name:
  • the optimizer must hold every trainable parameter, by identity (optim/param/<name>). An optimizer_fn that builds parameter groups and leaves one out does so on both sides, so the numbers alone would pass;
  • NumericsSpec(tied_params=(("embed_tokens.weight", "lm_head.weight"),)): each declared group must be one Parameter in each run's own model, also after the plugin wrapped it (FSDP excepted); undeclared shared parameters fail too (tied/<names>). A tied parameter is compared once, under its first name.
• Device audit: every audit record carries the device type (float32@cpu), and the audit also records what the optimizer updates (optim/params), so master weights or Adam moments offloaded on one side only raise RecipeMismatch before numbers are judged.
• Zero valid tokens (token_weighted=True): a micro-batch without valid tokens records its summed loss, 0, instead of 0 / 0, and still takes part in the gradient reduction; a step without any valid token raises ValueError on every rank before the first forward, so no rank waits in a collective. The README states how DDP's rank averaging and the token weighting relate.
• Resume state: after the last step, run_steps also records, in the plugin's own layout, the tensors the optimizer updates (optim/<group>.<index>/param, the float32 master weights where the plugin keeps them), Adam's state, the scheduler's state, the fp16 loss scale and the RNG states (rng/cpu, rng/cuda); resume.check compares them bit for bit. The docstring and README state what the scenario carries itself and a booster checkpoint does not: data progress and RNG state. A model that draws random numbers while training fails on rng/*.
• One model on the device at a time: prepare keeps the model and reference on the CPU; each run copies the one it needs to the device, and the float64 truth is cast on the way there. For a 111M-parameter Llama this removes 0.44 GiB (one float32 copy) from every run's peak (truth 5.49 → 5.05 GiB, ref 3.27 → 2.83 GiB, LowLevelZero bf16 3.48 → 3.03 GiB).
• Gradient checkpointing: enabled non-reentrant in prepare_fn, verified to recompute every layer and to align through TorchDDP, LowLevelZero (bf16, offload), Gemini, FSDP and HybridParallel (bf16, tp_size=1).
• Two recipe fixes, found by requiring the plugin self-tests to reproduce the reference exactly on one device:
  • FSDP MixedPrecision keeps the model's own float32 weights as its sharded parameters, so the reference's master weights now start from the unrounded float32 weights (Recipe.master_from_float32), not from their cast. Mixtral fsdp_bf16: max ratio 1.42 → 1.00; a model with tied embeddings had reached 237 on one loss.
  • With master weights, the fp16 loss scale is divided out in float32, as GeminiOptimizer does in its step (div_scale). Mixtral gemini_fp16: 1.68 → 1.07. The rest is HybridAdam against torch's AdamW, which differ in the last float32 bit of a few elements; Gemini fp16 is the one configuration exempt from the exactness assertion.
• Reference cache: the key now covers the declared unused and tied parameters, and reads the optimizer's parameter groups from the reference's own parameters (an optimizer_fn that filters parameters by shape no longer fails on a dummy one).

4. User and compatibility impact

• Python API of colossalai, checkpoint format, training behavior, performance: none.
• NumericsSpec gains an optional tied_params (default ()). An entry whose model shares parameters without declaring them now fails tied/...; no model_zoo entry with a NumericsSpec is affected (Mixtral has no tied parameters).
• Recipe gains master_from_float32 (default False). Cached reference runs are invalidated by the changed capture code, as designed.
• Reports: new entry kinds optim/param/..., tied/..., and on resume optim/..., scheduler/..., loss_scale, rng/....

5. Test environment

• OS: Ubuntu 22.04.5 LTS (CI image torch2.13-cu130-tf5.17.0)
• Python: 3.10.12
• GPU: 2× NVIDIA H20
• NVIDIA driver: 595.71.05
• PyTorch: 2.13.0+cu130
• Transformers: 5.17.0
• ColossalAI: this branch, bdab6a9, on main 69c9ff5
• Also the Layer 1 CPU job's environment, rebuilt as ci_layer_1_pr.yml installs it: torch 2.5.1+cpu, requirements/requirements.txt, Python 3.10.12

6. Test commands and results

CUDA_VISIBLE_DEVICES=<two GPUs> python -m pytest -q -rs tests/test_numerics

• 78 passed, 0 failed, 0 skipped, 640 s.

CUDA_VISIBLE_DEVICES= python -m pytest -q -rs tests/test_numerics

• torch 2.13 (CI image): 54 passed, 0 failed, 25 skipped, 254 s. Every skip needs a device.
• torch 2.5.1 (Layer 1 environment): 48 passed, 0 failed, 31 skipped, 192 s: the 25 above and 6 resume tests, which need torch 2.10 or later on CPU (unsharded optimizer checkpoints, see [test] numerical alignment: whole models through booster plugins, data-parallel and resume equivalence, acceptance report #6505).

No xfail, --ignore or manual reruns; FAST_TEST is not used.

Acceptance of the Mixtral entry under all 14 configurations (python -m tests.kit.numerics.acceptance --entry transformers_mixtral, 2 GPUs): 29 passed, 12 failed, 2 unsupported, 14 not run, the same verdicts as on main. Every failure is one of #6490, #6492, #6493, #6496, #6497, #6498. All whole-model alignments have max ratio 1.00 except gemini_fp16 (1.07).

7. Known issues

• Not covered: the full float64 truth of a model too large for one device; such a model is aligned with a reduced configuration, which the README asks to state wherever the result is reported.
• Gemini fp16 is not bit-exact against torch's AdamW (see above).
• The framework issues listed in [test] numerical alignment: whole models through booster plugins, data-parallel and resume equivalence, acceptance report #6505 are unchanged.

8. Risks and rollback

• Risk: the stricter contract checks fail entries that pass today if they share parameters without declaring them; the failure names the parameters and the fix (declare them).
• Rollback: revert this PR; nothing outside tests/ depends on it.

9. Documentation and release impact

• tests/kit/numerics/README.md: parameter contract, device audit, large models, zero valid tokens, resume state and boundary, gradient checkpointing, recipe table, calibration note.
• Installation, API and Docker documentation, changelog, version, PyPI, Docker images: none.