From af4d9c2d5bf547eaeff550e1aec6dcb9e196b7ae Mon Sep 17 00:00:00 2001 From: maxiaosong1124 Date: Wed, 12 Aug 2026 00:23:39 +0800 Subject: [PATCH 1/4] feat(ws1): land C1 four-judgment numerical contract (#267) Freeze the WS1 numerical SSOT for issue #267: four-judgment tolerance rows, dtype/TF32/FP8 policy, comparison roles, chain logprob aggregates, shared resolver, and op_checks wiring so forward and gradient accuracy no longer share one threshold path. Add schema tests, usage docs, and a migration checklist for remaining private-atol call sites (C3/C4/C8). Closes #267 --- docs/contributing/gtest-usage.md | 389 +++++++ docs/contributing/testing.md | 21 + docs/design/ws1-gtest-migration-checklist.md | 248 +++++ docs/design/ws1-numerical-contract.md | 187 ++++ rl_engine/kernels/gtest/__init__.py | 12 + rl_engine/kernels/gtest/op_checks.py | 121 ++- rl_engine/kernels/gtest/tolerance.py | 967 +++++++++++++++++- .../kernels/gtest/tolerance_contract.json | 239 ++++- tests/test_op_checks.py | 99 +- tests/test_tolerance_contract.py | 446 +++++++- 10 files changed, 2716 insertions(+), 13 deletions(-) create mode 100644 docs/contributing/gtest-usage.md create mode 100644 docs/design/ws1-gtest-migration-checklist.md create mode 100644 docs/design/ws1-numerical-contract.md diff --git a/docs/contributing/gtest-usage.md b/docs/contributing/gtest-usage.md new file mode 100644 index 00000000..036c1b65 --- /dev/null +++ b/docs/contributing/gtest-usage.md @@ -0,0 +1,389 @@ +# gtest usage guide (operator candidate vs gold) + +> **Audience:** contributors implementing train–inference / batch-invariant operators +> **Entry point:** `scripts/check_operator.py` + `rl_engine/kernels/gtest/*` +> **Numerical SSOT:** [#267](https://github.com/RL-Align/RL-Kernel/issues/267) four-judgment contract +> **Related:** [WS1 numerical contract](../design/ws1-numerical-contract.md) · [migration checklist](../design/ws1-gtest-migration-checklist.md) + +This is the official how-to for the gtest harness: register an op, build inputs, run the CLI for forward/backward checks, and obtain tolerances from the shared contract (not private `atol`/`rtol`). + +--- + +## 1. What gtest is for + +gtest validates a **single operator**: + +| Capability | Meaning | +|------------|---------| +| Gold | Usually a PyTorch / `forward_fp32` reference path | +| Candidate | CUDA / Triton / arch-specific implementation | +| Forward check | Outputs within contract tolerance (`forward_accuracy`) | +| Backward check | Selected input gradients within contract tolerance (`gradient_accuracy`, **independent of forward**) | + +It is **not**: + +- The full Qwen3-8B model-level gate (#266 C9/C10) +- The final cross-config invariance harness (C3/C4 build on the same contract) +- Real vLLM vs Megatron engine alignment + +The CLI primarily covers **accuracy** (candidate vs gold). +**Invariance** (bitwise across configs) and **train/infer aggregates** use the contract APIs / later harnesses—do not invent private gate thresholds in tests. + +--- + +## 2. End-to-end flow + +```text +1) (Optional) register the op in the runtime registry + ↓ +2) gtest/operator_specs.py → OP_SPECS: gold + candidates + ↓ +3) gtest/operator_inputs.py → build input shapes / values + ↓ +4) scripts/check_operator.py → run suite, load tolerance_contract.json + ↓ +5) report max_abs / tol / passed +``` + +### 2.1 Key files + +| Path | Role | +|------|------| +| `rl_engine/kernels/gtest/operator_specs.py` | `OP_SPECS`: name, `op_class`, gold, candidates, grad inputs | +| `rl_engine/kernels/gtest/operator_inputs.py` | Default Qwen3-8B dims + `make_operator_inputs` | +| `rl_engine/kernels/gtest/op_checks.py` | Suite execution and comparison | +| `rl_engine/kernels/gtest/tolerance_contract.json` | Numerical contract SSOT | +| `rl_engine/kernels/gtest/tolerance.py` | `load_contract` / `resolve_tolerance` / chain aggregates | +| `scripts/check_operator.py` | **CLI entry** | + +--- + +## 3. Step 1: register the op in `OP_SPECS` + +Edit `rl_engine/kernels/gtest/operator_specs.py` and add an entry to `OP_SPECS`. Example shape (logp / linear_logp): + +```python +"logp": OperatorSpec( + name="logp", + op_class="logprob", # selects the contract op_class row + gold_path="rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp", + gold_method="forward_fp32", # method invoked on the gold instance + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp", + "cuda": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op", + }, + grad_input_names=("logits",), # inputs compared under --check-grad +), +``` + +### 3.1 `OperatorSpec` fields + +| Field | Meaning | +|-------|---------| +| `name` | Value for CLI `--op` | +| `op_class` | Contract class: `elementwise` / `reduction` / `logprob` / `attention` | +| `gold_path` | Gold class path `module.Class` | +| `gold_method` | Method name, e.g. `forward_fp32`, `apply`, `__call__` | +| `candidate_paths` | Map `candidate name → implementation class`; CLI `--candidate cuda` looks up this map | +| `grad_input_names` | With `--check-grad`, enable grads and compare these inputs; missing config errors | + +**Only ops registered in `OP_SPECS` can be invoked via `check_operator.py`.** + +Currently registered (source of truth is the code): + +```text +rms_norm, attention, logp, linear_logp, embedding, lm_head, +det_gemm, rope, silu, swiglu, batch_invariant_logp +``` + +--- + +## 4. Step 2: build inputs + +File: `rl_engine/kernels/gtest/operator_inputs.py`. + +### 4.1 Default model dims (Qwen3-8B Dense semantics) + +Macros at the top of the file (local experiments may change them; WS1 full-model EXIT uses the official config fingerprint): + +```text +DEFAULT_HIDDEN = 4096 +DEFAULT_N_HEADS = 32 +DEFAULT_N_KV_HEADS = 8 +DEFAULT_HEAD_DIM = 128 +DEFAULT_INTERMEDIATE = 12288 +DEFAULT_VOCAB = 151936 +DEFAULT_ROPE_THETA = 1.0e6 +DEFAULT_RMS_EPS = 1.0e-6 +``` + +### 4.2 Shape names and input builders + +- `operator_shape_name(op_name, args)` — human-readable case name (e.g. `2x16x257`) +- `_make_*_inputs` / `make_operator_inputs` — build the input dict from `--op` and CLI args + - `random`: reproducible randomness (`--seed` plus per-tensor offsets) + - `constant`: fixed values for debugging (`--constant-value` / `--token-value`) + +When adding an op: extend the shape map and implement the matching `_make_xxx_inputs`. + +### 4.3 Suggested GRPO-oriented shapes (local sweeps) + +For GRPO, `B = P × G`. With `G=8`, batch is often a multiple of 8. +`B=1` is fine for smoke; fuller sweeps may use: + +```text +B ∈ {1, 8, 16, 32, 64} +S ∈ {1, 31, 33, 127, 129, 255, 256, 257, 512, 1024, 4096, 8192} +``` + +Prefer short `S` when VRAM is tight; full-model gates are owned by #266 / C2. + +--- + +## 5. Step 3: run the CLI + +```bash +# From the repo root; prefer an editable install: pip install -e . +python scripts/check_operator.py --op logp --candidate pytorch --device cpu --dtype fp32 --batch 1 --seq 2 --vocab 17 +``` + +### 5.1 Common examples + +**Smoke (CPU / PyTorch self-check)** + +```bash +python scripts/check_operator.py \ + --op logp --candidate pytorch --device cpu --dtype fp32 \ + --batch 1 --seq 2 --vocab 17 +``` + +**Triton `linear_logp` + backward (BF16)** + +```bash +python scripts/check_operator.py \ + --op linear_logp --candidate triton --device cuda --dtype bf16 \ + --batch 1 --seq 2 --vocab 1024 --normalized-dim 4096 \ + --check-grad +``` + +**CUDA deterministic attention + gradients** + +```bash +python scripts/check_operator.py \ + --op attention --candidate cuda --device cuda --dtype bf16 \ + --batch 2 --seq 64 --check-grad --grad-mode random +``` + +**Full JSON report** + +```bash +python scripts/check_operator.py --op rms_norm --candidate cuda --dtype bf16 --device cuda --json +``` + +### 5.2 CLI flags + +| Flag | Meaning | +|------|---------| +| `--op` | Operator name from `OP_SPECS` | +| `--candidate` | Backend: `pytorch` / `cuda` / `cuda-generic` / `cuda-sm90` / `triton` / … (see that op’s `candidate_paths`) | +| `--dtype` | `fp32` / `bf16` / `fp16`; selects input dtype and contract row | +| `--device` | `auto` / `cpu` / `cuda` | +| `--batch` / `--seq` | Batch size and sequence length for inputs | +| `--vocab` | Vocab size; logp logits `[B,S,V]`; linear_logp weight `[V,H]` | +| `--input-mode` | `random` (default) or `constant` | +| `--constant-value` | Float fill in constant mode | +| `--token-value` | Token id in constant mode | +| `--normalized-dim` | Hidden dim for rms_norm / linear_logp, etc. | +| `--k-dim` / `--n-dim` | Matmul / det_gemm dims | +| `--theta` | RoPE theta | +| `--eps` | RMSNorm epsilon | +| `--seed` | Input RNG seed (per-tensor offsets still apply) | +| `--arch-key` | Arch override key, e.g. `sm90` (contract `arch_overrides`) | +| `--check-grad` | Also compare gradients (requires `grad_input_names`) | +| `--grad-mode` | `random` (default, stricter) / `ones` (≈ `output.sum().backward()`) | +| `--grad-seed` | Seed for random upstream gradients | +| `--json` | Print the full structured report | + +--- + +## 6. Where tolerances come from (after #267) + +### 6.1 Before vs after C1 + +| Before | After (C1 / #267) | +|--------|-------------------| +| Mostly `accuracy[op_class][dtype]` | **Four judgments**: forward/gradient × accuracy/invariance | +| Forward and grad often shared one tol | **Grad uses `gradient_accuracy` only** (no silent forward inheritance) | +| Flat threshold table | Plus dtype policy, comparison roles, chain logprob aggregates | + +### 6.2 Which judgments the CLI / `op_checks` use + +`run_operator_suite` / `check_operator.py`: + +| Comparison | Judgment | +|------------|----------| +| Output vs gold | `forward_accuracy` | +| Gradient vs gold | `gradient_accuracy` | + +Batch/chunk **bitwise invariance** and train/infer **three aggregates** are not separate CLI switches. Use: + +```python +from rl_engine.kernels.gtest.tolerance import ( + load_contract, + resolve_tolerance, + compute_logprob_aggregates, + judge_logprob_aggregates, + default_clip_interval, +) + +contract = load_contract() +# Cross-config invariance (gate path) +inv = resolve_tolerance( + contract, + judgment="forward_invariance", # or gradient_invariance + op_class="attention", + dtype="bfloat16", + backend_profile="cuda_bf16", +) +# inv.mode == "bitwise", inv.atol == inv.rtol == 0 + +# Train vs infer selected-logprob +agg = compute_logprob_aggregates( + train_logp, + rollout_logp, + active_mask, + contract=contract, + report_kind="train_infer_logprob_parity", + clip_interval=default_clip_interval(contract), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", +) +verdict = judge_logprob_aggregates(agg, contract, execution_dtype="bfloat16") +``` + +### 6.3 Policy locks (WS1) + +| Item | Value | +|------|--------| +| Execution | BF16 mandatory for EXIT (CLI may still exercise fp32/fp16) | +| Reference / accumulation | FP32 | +| FP8 | Out of scope (resolve hard-fails) | +| TF32 | Disabled | + +WS1 evidence must attach checked provenance to its candidate report: + +```python +from rl_engine.kernels.gtest import BackendProvenance, CandidateSpec + +provenance = BackendProvenance( + backend_profile="cuda_bf16", # use triton_cuda_bf16 + triton for Triton + requested_backend="cuda", + actual_backend="cuda", + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, +) +candidate = CandidateSpec( + name="cuda-candidate", + backend="cuda", + fn=op, + provenance=provenance, +) +``` + +The suite rejects backend fallback, dtype drift, TF32 enablement, and observed output +dtypes that disagree with this provenance before producing a passing report. +| Profiles | `cuda_bf16` and `triton_cuda_bf16` share **the same** thresholds | + +**Do not** use private `atol=1e-5` (etc.) as WS1 gate evidence. Inventory of legacy private thresholds: [migration checklist](../design/ws1-gtest-migration-checklist.md). + +### 6.4 Report `tol=(atol=..., rtol=...)` + +The CLI summary line: + +```text +tol=(atol=..., rtol=...) +``` + +comes from the shared resolver—not hard-coded constants inside `check_operator.py`. + +--- + +## 7. Recommended local test order + +```text +1. --candidate pytorch --device cpu --dtype fp32 + → registration / inputs / plumbing smoke + +2. Same shape with --dtype bf16 --device cuda --candidate triton|cuda + → real candidate forward + +3. Add --check-grad --grad-mode random + → gradients (random upstream grads catch more bugs than ones) + +4. --arch-key sm90 only when you need arch-specific contract overrides + +5. Cross batch/layout: not CLI-only; use invariance judgments + dedicated tests +``` + +--- + +## 8. Common failures + +| Symptom | Likely cause | +|---------|----------------| +| Unsupported / missing `--op` choice | Not registered in `OP_SPECS` | +| `--check-grad` missing grad inputs | Empty/wrong `grad_input_names` vs input keys | +| Candidate import error | Bad `candidate_paths` or extension not built | +| BF16 over tolerance | Confirm gold is `forward_fp32`; check contract row; do not loosen private atol | +| Missing SM90 symbols | Build without SM90 / non-sm90 GPU; pick another candidate or rebuild | +| Want FP8 | Hard-fail under WS1 contract; out of scope | + +--- + +## 9. Relationship to pytest + +| Path | Use | +|------|-----| +| `python scripts/check_operator.py ...` | Fast single-op shape/debug loops | +| `pytest tests/test_*.py` | Regression, invariance, integration | +| `pytest tests/test_tolerance_contract.py` | Contract schema / resolver | + +Both paths should take thresholds from `tolerance_contract.json`. +New pytest code should call `resolve_tolerance` instead of copying magic numbers. + +--- + +## 10. Minimal checklist for a new operator + +- [ ] Implementation under `rl_engine/kernels/ops/{pytorch,cuda,triton}/...` +- [ ] (Optional) runtime `registry` registration +- [ ] `OP_SPECS` entry: gold + candidates + `op_class` + `grad_input_names` +- [ ] `operator_inputs` shape name + input builder +- [ ] `check_operator.py` smoke + bf16 + `--check-grad` green +- [ ] Contract already has the `op_class` row (extend schema + `test_tolerance_contract` if not) +- [ ] No new private `atol`/`rtol` as gate evidence +- [ ] Operator docs point at the contract for thresholds (do not restate ad-hoc numbers) + +--- + +## 11. Further reading + +| Doc | Content | +|-----|---------| +| [ws1-numerical-contract.md](../design/ws1-numerical-contract.md) | Four judgments, roles, aggregate formulas | +| [ws1-gtest-migration-checklist.md](../design/ws1-gtest-migration-checklist.md) | Which tests still use private thresholds and when to migrate | +| [testing.md](testing.md) | Short testing entry points | +| Issues [#266](https://github.com/RL-Align/RL-Kernel/issues/266) / [#267](https://github.com/RL-Align/RL-Kernel/issues/267) | WS1 closeout and C1 contract | + +--- + +## 12. Changelog + +| Date | Notes | +|------|--------| +| 2026-08-11 | Initial English guide aligned with C1; documents CLI, `OP_SPECS`, inputs, and contract usage | diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index a96c0014..b0924ba0 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -2,6 +2,21 @@ RL-Kernel uses focused tests for dispatch behavior and operator accuracy. +## gtest (operator candidate vs gold) + +Primary entry for single-operator forward/backward checks against a PyTorch gold path: + +```bash +python scripts/check_operator.py --op logp --candidate pytorch --device cpu --dtype fp32 +``` + +Full usage (register `OP_SPECS`, build inputs, CLI flags, and the WS1 four-judgment +tolerance contract after #267): + +- **[gtest usage guide](gtest-usage.md)** (operator CLI + `OP_SPECS` + contract; English) +- [WS1 numerical contract](../design/ws1-numerical-contract.md) +- [gtest private-threshold migration checklist](../design/ws1-gtest-migration-checklist.md) + ## Dispatch Tests ```bash @@ -14,6 +29,12 @@ python -m pytest rl_engine/tests/test_dispatch.py -v python tests/test_op_accuracy.py ``` +Contract schema / resolver: + +```bash +python -m pytest tests/test_tolerance_contract.py tests/test_op_checks.py -q +``` + ## Documentation Build ```bash diff --git a/docs/design/ws1-gtest-migration-checklist.md b/docs/design/ws1-gtest-migration-checklist.md new file mode 100644 index 00000000..3b0a9346 --- /dev/null +++ b/docs/design/ws1-gtest-migration-checklist.md @@ -0,0 +1,248 @@ +# WS1 gtest 阈值迁移清单 + +> **关联：** [#266](https://github.com/RL-Align/RL-Kernel/issues/266) 父收尾 · [#267](https://github.com/RL-Align/RL-Kernel/issues/267) C1 契约 · [数值契约说明](ws1-numerical-contract.md) +> **目的：** 盘点「哪些测试仍用私有 `atol`/`rtol`、哪些已走 SSOT、何时必须迁到 `resolve_tolerance`」。 +> **快照：** 基于 `feat/ws1-c1-tolerance-contract-267` 落地 C1 后的仓库状态；文件增减时请更新本表。 + +--- + +## 0. 迁移总原则 + +### 0.1 SSOT 入口（改后唯一推荐） + +```python +from rl_engine.kernels.gtest.tolerance import ( + load_contract, + resolve_tolerance, + compute_logprob_aggregates, + judge_logprob_aggregates, + default_clip_interval, +) + +contract = load_contract() +spec = resolve_tolerance( + contract, + judgment="forward_accuracy", # 或 forward_invariance / gradient_* + op_class="logprob", # elementwise | reduction | logprob | attention + dtype="bfloat16", + backend_profile="cuda_bf16", # 与 triton_cuda_bf16 同阈值 +) +# assert_close(..., atol=spec.atol, rtol=spec.rtol) +# 不变性：spec.mode == "bitwise" 且 atol=rtol=0 → 优先 torch.equal +``` + +| Judgment | 用于 | +|----------|------| +| `forward_accuracy` | BF16 candidate vs FP32 reference | +| `forward_invariance` | 同逻辑 workload 跨 batch/chunk/layout（**bitwise**） | +| `gradient_accuracy` | 梯度 vs FP32 参考（**不得**读 forward 行） | +| `gradient_invariance` | 梯度跨 config（**bitwise**） | +| 三聚合 API | 链级 / 训推 selected-logprob（`max_abs_dlogp` / `approx_kl0` / `clipfrac0`） | + +### 0.2 什么叫「私有阈值」（禁止作为 WS1 gate 证据） + +- 测试文件内字面量：`atol=1e-5`、`atol=5e-2`、模块常量 `_DECODE_ATOL` 等 +- 文档里写死但未从 `tolerance_contract.json` resolve 的数 +- 从 `contract["accuracy"]...` 手抄数值后本地再改（应用 resolve，不要复制常量） +- 用非零 `atol` 充当 Batch/Chunk **invariance** 通过条件 + +### 0.3 什么可以保留（不必硬迁） + +| 场景 | 处理 | +|------|------| +| **bitwise 身份断言**（`torch.equal`） | 合法；对应 invariance judgment 的 `mode=bitwise` | +| **非数值语义**（mask 形状、版本单调、manifest 字段） | 不迁 | +| **框架/集成单测**（bridge、vLLM mock、DeepSpeed worker 编排） | 非 WS1 op gate；可保留宽松 `allclose`，但**不能**当作 #266 EXIT 证据 | +| **生产 FA / SDPA 对齐**（`test_attention_correctness`） | 非 BI 候选路径；阈值可独立，**不得**写进 WS1 EXIT claim | +| **legacy `accuracy` 键** | 仅兼容；新代码禁止新增依赖，应改 `resolve_tolerance` | + +### 0.4 建议迁移句式 + +```python +# BAD — 私有阈值 +torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + +# GOOD — accuracy +spec = resolve_tolerance(contract, judgment="forward_accuracy", op_class="reduction", dtype=dtype) +torch.testing.assert_close(out, ref, atol=spec.atol, rtol=spec.rtol) + +# GOOD — invariance +ispec = resolve_tolerance(contract, judgment="forward_invariance", op_class="attention", dtype=dtype) +assert ispec.mode == "bitwise" and ispec.atol == 0.0 +assert torch.equal(a, b) # 或 assert_close(..., atol=0, rtol=0) + +# GOOD — gradient accuracy（独立 judgment） +gspec = resolve_tolerance(contract, judgment="gradient_accuracy", op_class="logprob", dtype=dtype) +``` + +--- + +## 1. 状态总表（`tests/`） + +图例： + +| 标记 | 含义 | +|------|------| +| **A** | 已走 resolver / gtest suite（目标态） | +| **B** | 走 `load_contract` 旧键或 gtest 间接路径（过渡） | +| **C** | WS1 相关但 **私有 atol**（应迁） | +| **D** | 多为 `torch.equal` / 结构断言（ok 或仅需声明 judgment） | +| **E** | 非 WS1 门禁（框架/产品路径，低优先级） | + +### 1.1 已对齐或接近 SSOT + +| 文件 | 状态 | 说明 | 下一步 | +|------|------|------|--------| +| `test_tolerance_contract.py` | **A** | C1 schema + resolve + 聚合 | 保持；契约变更必跑 | +| `test_op_checks.py` | **A/B** | suite 已按 judgment 解析；部分用例注入最小 contract | 新 fixture 尽量带 `judgments` | +| `test_operator_inputs.py` | **B** | 输入/规格，无数值阈值主责 | 无需迁阈值 | +| `test_swiglu.py` | **B/D** | issue-108 harness + 大量 `torch.equal` | accuracy 路径确认走 suite；字面量 atol 清零 | +| `test_det_gemm.py` | **B** | `load_contract()["accuracy"]...` | **优先迁**：改为 `resolve_tolerance(..., forward/gradient_accuracy)` | +| `test_deterministic_attention_cuda.py` | **B/C** | 部分用 suite；仍见 `5e-2/2e-2` 字面量 | 字面量改为 resolve；invariance 保持 equal | + +### 1.2 WS1 算子测试 — 私有阈值（应迁，按优先级） + +| 优先级 | 文件 | op_class 建议 | 现状摘要 | 何时必须迁 | +|--------|------|---------------|----------|------------| +| **P0** | `test_batch_invariant_logp.py` | `logprob` | 大量 `1e-6`…`1e-2` 私有；含 bwd | 接 C3/C4/C8 证据前 | +| **P0** | `test_linear_logp.py` | `logprob` | `1e-5`…`1.5e-1` 混用；bf16 松阈值 | 同上；链级改用三聚合 API | +| **P0** | `test_logp.py` / `test_deterministic_logp.py` | `logprob` | 私有 atol | 关 #148 residual / C8 前 | +| **P0** | `test_rms_norm.py` | `reduction` | `1e-5`…`8e-2`；bwd 混用 | C8 RMSNorm 证据前 | +| **P0** | `test_triton_batch_invariant_attention.py` | `attention` | 混 `1e-5` 与 `5e-2/2e-2` | C8 Attention 证据前 | +| **P0** | `test_attention.py` | `attention` | native GT；`1e-4`/`2e-6` 等 | 与 contract `attention` 行对齐 | +| **P1** | `test_kv_cache_attention.py` | `attention` | 含 `2e-6` 等；#152 相关 | **C6/C7 前必须**消私有 decode 阈值 | +| **P1** | `test_issue151_embedding_lm_head_invariance.py` | emb + lm_head + logp | bf16 `5e-2` 手写 | C8 emb/lm_head 证据前 | +| **P1** | `test_lm_head.py` | `reduction` | 多 equal；grad `1e-5` 私有 | 迁 grad → `gradient_accuracy` | +| **P1** | `test_embedding.py` | `elementwise` | 多为 equal | 若有 tolerance 路径再 resolve | +| **P1** | `test_rope.py` | `elementwise` | `1e-3`…`2e-2` | C5 RoPE 证据前 | +| **P1** | `test_matmul.py` | `reduction` | 私有 `1e-4/1e-5` | 与 det_gemm 统一 | +| **P2** | `test_pack.py` | `elementwise` | 几乎 equal；gradcheck `1e-6` | packing 纳入 #150 时 | +| **P2** | `test_grpo_loss.py` / `test_ratio_kl.py` | （loss，契约暂无独立 class） | `1e-4` 等 | 若进 chain 则扩展 op_class 或显式 N/A | +| **P3** | `test_attention_correctness.py` | 非 BI EXIT | FA/SDPA 私有表 | **不迁入 WS1 SSOT**；文档标明 out of WS1 claim | +| **P3** | `test_op_accuracy.py` | 杂项 harness | `1e-3` | 废弃或改走 `check_operator` + contract | + +### 1.3 非 WS1 门禁（低优先级 / 不阻塞 #267） + +| 文件 | 状态 | 说明 | +|------|------|------| +| `test_deepspeed_training_worker.py` | **E** | 训练 worker；`atol=1e-5` 编排级 | +| `test_stateless_training_contract.py` | **E** | 契约字段/数值 smoke | +| `test_rl_kernel_loss_step.py` | **E** | 端到端 loss 步 | +| `test_sampler_temperature.py` | **E** | 采样 | +| `test_weight_sync_bridge.py` 等 | **D/E** | bridge / IPC | +| `test_vllm_rollout_sampler.py` | **D/E** | vLLM mock | +| `test_alignment_model_wrappers.py` | **D/E** | wrapper 行为 | +| `test_rl_batch_fixture.py` | **D** | fixture 身份 | +| `test_stateless_executor.py` / `*_hf_integration*` | **D/E** | 执行器集成 | + +这些**不**作为 #266 Full WS1 EXIT 的数值证据来源；C10/C11 不得引用其私有阈值刷绿。 + +--- + +## 2. 按 #266 子 issue 的「何时必须迁」 + +| 子 issue | 阻塞迁移范围 | 完成信号 | +|----------|--------------|----------| +| **C1 #267** | 契约 + resolver + schema/报告测试 | **实现完成；待 CI / issue evidence** | +| **C3 #269** forward harness | 所有 **forward_accuracy / forward_invariance** 的 op 单测证据路径 | 无 private forward atol 作为 gate | +| **C4 #270** grad harness | 所有 **gradient_*** 证据路径 | 无「grad 抄 forward 字面量」 | +| **C5 #271** RoPE/elementwise | `test_rope.py`、activation/swiglu residual | audit 报告阈值均来自 resolve | +| **C6/C7 #272/#273** KV | `test_kv_cache_attention.py` 及后续 kv harness | **禁止** `_DECODE_ATOL` 类私有常量 | +| **C8 #274** closed-op 矩阵 | rmsnorm / gemm / attn / logp / emb / lm_head 测试 | 每格 `requested/actual backend` + resolve 阈值 | +| **C10 #276** 全模型 gate | 仅用 resolver + 三聚合；禁止任何测试内字面量阈值 | gate 报告无 private tol 字段 | +| **C11 #277** CI | CI 只跑 resolve 路径 | fail-closed | + +**规则：** 某 op 的 PR 若声称「满足 #266/C8」，则该 PR 触达的 assert **必须**来自 `resolve_tolerance` / 聚合 API，而不是文件顶部的魔法数。 + +--- + +## 3. 文件级迁移清单（可勾选） + +### 3.1 P0 — 直接挡 C3/C4/C8 + +- [ ] `tests/test_batch_invariant_logp.py` — fwd/bwd accuracy + invariance 拆 judgment +- [ ] `tests/test_linear_logp.py` — 同上；训推/链级改用三聚合 +- [ ] `tests/test_logp.py` +- [ ] `tests/test_deterministic_logp.py` +- [ ] `tests/test_rms_norm.py` +- [ ] `tests/test_triton_batch_invariant_attention.py` +- [ ] `tests/test_attention.py` +- [ ] `tests/test_det_gemm.py` — 去掉 `contract["accuracy"]` 直读 + +### 3.2 P1 — C5/C6/C7/C8 residual + +- [ ] `tests/test_kv_cache_attention.py` +- [ ] `tests/test_issue151_embedding_lm_head_invariance.py` +- [ ] `tests/test_lm_head.py` +- [ ] `tests/test_embedding.py`（若有 non-bitwise 路径） +- [ ] `tests/test_rope.py` +- [ ] `tests/test_matmul.py` +- [ ] `tests/test_deterministic_attention_cuda.py` 中剩余字面量 +- [ ] `tests/test_swiglu.py` 中任何 residual 字面量 + +### 3.3 P2 — 进 chain 时 + +- [ ] `tests/test_pack.py` +- [ ] `tests/test_grpo_loss.py` / `tests/test_ratio_kl.py`（先扩 contract op_class 或标 N/A） +- [ ] `scripts/check_operator.py` 报告字段确认只回传 resolve 结果（已间接） + +### 3.4 明确不迁入 WS1 SSOT + +- [x] `tests/test_attention_correctness.py` — 生产 FA；文档标注非 EXIT +- [x] bridge / vLLM / DeepSpeed / sampler 类 **E** 组 + +--- + +## 4. 推荐落地动作（每个测试文件） + +1. **分类每条 assert** + - identity / batch-invariance → `forward_invariance` 或 `gradient_invariance` + `torch.equal` + - vs fp32 gold → `*_accuracy` + - train vs infer logp → 三聚合，不用单点 atol 冒充 +2. **删除模块级 `_ATOL` / `_RTOL`** +3. **dtype 参数化** 时用 `resolve_tolerance(..., dtype=dtype)`，禁止 bf16 写死 `5e-2` +4. **报告**（若有）写上 `comparison_lhs_role` / `comparison_rhs_role`（从 spec 取） +5. **禁止** 为让 invariance 通过而调大 atol + +### 4.1 与契约行不一致时怎么办 + +| 情况 | 动作 | +|------|------| +| 测试私有更松，契约更紧 → 测试红 | **修 kernel** 或开 Blocker；**禁止**在测试放宽 | +| 测试私有更紧，契约更松 | 迁到契约后可能变绿；可保留额外严格 assert 但须标注 *non-gate* | +| 需要新 op_class（如 `grpo_loss`） | 先改 `tolerance_contract.json` + schema 测试，再迁测试 | +| decode vs prefill 无法 bitwise | 用 contract 已声明的 semantic 行 / 三聚合；**不要**私设 `_DECODE_ATOL` | + +--- + +## 5. 工具与 CI 建议（后续，非 C1 范围） + +| 建议 | 作用 | +|------|------| +| 简单 lint：`tests/**/*.py` 禁止 `atol=\d`（allowlist 契约测试与 FA 测试） | 防回流 | +| `pytest` marker：`ws1_gate` 仅收集 resolve 路径 | C11 门禁清晰 | +| 在 `check_operator.py` 输出中强制打印 `judgment` + roles | 证据可检索 | + +C1 **不**强制上 lint；C8/C10 前建议至少做 allowlist 扫描。 + +--- + +## 6. 现状一句话 + +| 层 | 状态 | +|----|------| +| **契约 + resolver（gtest 核心）** | 已就绪；待 CI / issue evidence（#267） | +| **op_checks 接入** | 已按 judgment 分叉，并持久化 roles / provenance | +| **存量 op 单测** | **多数仍私有 atol**（上表 P0/P1） | +| **#266 EXIT** | 依赖后续把 P0/P1 迁完，而不是只合 C1 | + +**C1 的价值是「唯一入口已存在」；清单的价值是「知道还欠哪些文件」。** +未完成 P0/P1 迁移前，**不得**声称「全仓测试已统一走 WS1 数值契约」。 + +--- + +## 7. 修订记录 + +| 日期 | 说明 | +|------|------| +| 2026-08-11 | 初版：C1 落地后基于 `tests/` 扫描的迁移清单与优先级 | diff --git a/docs/design/ws1-numerical-contract.md b/docs/design/ws1-numerical-contract.md new file mode 100644 index 00000000..76410fcb --- /dev/null +++ b/docs/design/ws1-numerical-contract.md @@ -0,0 +1,187 @@ +# WS1 Numerical Contract (C1 / #267) + +> **Parent:** [#266](https://github.com/RL-Align/RL-Kernel/issues/266) +> **Issue:** [#267](https://github.com/RL-Align/RL-Kernel/issues/267) +> **SSOT files:** `rl_engine/kernels/gtest/tolerance_contract.json`, `rl_engine/kernels/gtest/tolerance.py` + +This document freezes the **sole** numerical judgment source for WS1 ablations and +gates. New gates must obtain thresholds only through the shared resolver APIs; private +`atol` / `rtol` constants are forbidden. + +## Scope boundary + +**Allowed WS1 claim after full exit (#266):** single-GPU model-level train–inference +consistency for full Qwen3-8B Dense under required CUDA BF16 and Triton-on-CUDA BF16 +profiles (in-repo BI stack). + +**Not claimed here:** multi-GPU (WS2), real vLLM vs Megatron / vime product alignment +(WS3), or kernel bug fixes (open Blockers). + +## Dtype policy + +| Field | Lock | +| --- | --- | +| `execution_dtype` | **BF16** (mandatory) | +| `accumulation_dtype` | **FP32** | +| `reference_dtype` | **FP32** | +| `output_dtype.default` | follows execution | +| logprob aggregates compute dtype | **FP32** | +| FP8 | **out of scope** (request → hard fail) | +| FP16 | optional; rows complete when declared | +| TF32 (reference + candidate) | **disabled** (repo-wide single policy) | +| Backend profiles | `cuda_bf16`, `triton_cuda_bf16` (same thresholds) | +| Backend-private tolerance relaxation | **forbidden** | + +Execution, accumulation, output, and reference dtypes resolve **independently** via +`resolve_dtype_policy()`. + +## Four judgments + +| Judgment | What it compares | Default mode | +| --- | --- | --- | +| `forward_accuracy` | BF16 candidate vs FP32 reference outputs | tolerance | +| `forward_invariance` | transformed vs canonical config, same backend/dtype/logical workload | **bitwise** (`atol=0`, `rtol=0`) | +| `gradient_accuracy` | candidate gradient/VJP vs FP32 reference gradient/VJP | tolerance (independent of forward) | +| `gradient_invariance` | transformed vs canonical gradients, same logical workload | **bitwise** (`atol=0`, `rtol=0`) | + +Every declared-applicable `(judgment, op_class, dtype)` tuple must resolve. Missing +applicable cells hard-fail. Explicit `not_applicable` / `out_of_scope` is allowed only +when present in the schema. Use `resolve_tolerance_support()` to persist the explicit +support status; requesting thresholds for an N/A or out-of-scope cell still hard-fails. + +**Batch/Chunk invariance** (issue #150 / C10 matrix) **must** use the invariance +judgments in bitwise mode. Nonzero tolerance cannot satisfy that gate. + +Op classes: `elementwise`, `reduction`, `logprob`, `attention`. + +## Comparison roles + +Reports must record `comparison_lhs_role` / `comparison_rhs_role`. A bare `baseline` +field is forbidden. C2 `singleton_aggregate` is an **execution/aggregation mode**, not +a comparison role. + +| Report kind | `comparison_lhs_role` | `comparison_rhs_role` | +| --- | --- | --- | +| `forward_accuracy` | `bf16_candidate` | `fp32_reference` | +| `forward_invariance` | `transformed_config` | `canonical_config` | +| `train_infer_logprob_parity` | `training_style_teacher_forcing` | `inference_style_rollout_decode` | +| `gradient_accuracy` | `bf16_candidate` | `fp32_reference` | +| `gradient_invariance` | `transformed_config` | `canonical_config` | + +Direction is locked so train/infer preserves: + +```text +dlogp = train_logp - rollout_logp +ratio0 = exp(dlogp) +``` + +Swapping lhs/rhs without a different declared contract row hard-fails. + +API: `resolve_comparison_roles()`, `assert_comparison_roles()`. + +Aggregate callers must also provide `contract`, `report_kind`, and both roles; the +compute API validates the direction before calculating any metric. gtest reports +persist these roles on every output verdict. Backend reports must include +`BackendProvenance` (requested/actual backend, all four dtypes, and TF32 state), +which `validate_backend_provenance()` checks against the selected profile. + +## Chain-level logprob aggregates + +These three metrics are the **only** chain-level logprob / ablation aggregates for WS1 +pass/fail. Gradients use independent `gradient_*` tensor verdicts and **do not** use +these aggregates. + +Computed in FP32 on **active selected tokens only**: + +```text +dlogp = comparison_lhs_logp - comparison_rhs_logp +max_abs_dlogp = max(abs(dlogp)) +approx_kl0 = mean(exp(dlogp) - 1 - dlogp) +clipfrac0 = mean(1[exp(dlogp) outside clip_interval]) +``` + +Rules: + +- **All three** must pass (`require_all=true`). +- Empty active-token set → hard fail. +- NaN / Inf in `dlogp`, `ratio0`, or aggregates → hard fail. +- Clip interval is pinned by the workload manifest (C2); the contract stores a default + and the field name `clip_interval`. + +API: `compute_logprob_aggregates()`, `judge_logprob_aggregates()`, +`resolve_chain_aggregate_thresholds()`. + +## Resolver usage + +```python +from rl_engine.kernels.gtest.tolerance import ( + load_contract, + resolve_tolerance, + resolve_dtype_policy, + compute_logprob_aggregates, + judge_logprob_aggregates, + default_clip_interval, +) + +contract = load_contract() +policy = resolve_dtype_policy(contract) + +fwd = resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="logprob", + dtype="bfloat16", + backend_profile="cuda_bf16", +) +bwd = resolve_tolerance( + contract, + judgment="gradient_accuracy", + op_class="logprob", + dtype="bfloat16", + backend_profile="triton_cuda_bf16", # same thresholds as cuda_bf16 +) +inv = resolve_tolerance( + contract, + judgment="forward_invariance", + op_class="attention", + dtype="bfloat16", +) +# inv.mode == "bitwise", inv.atol == 0.0, inv.rtol == 0.0 +``` + +`op_checks.run_operator_suite` resolves **forward_accuracy** for outputs and +**gradient_accuracy** for gradients. + +## Compatibility keys + +For older tests that still dig into: + +- `contract["accuracy"]["default"][op_class][dtype]` — mirror of `forward_accuracy` +- `contract["batch_invariance"]` — `{atol: 0, rtol: 0}` + +New code should call the resolvers above. Schema validation fails if the compatibility +mirror drifts from `forward_accuracy` or if invariance rows leave bitwise mode. + +## Related issues + +| ID | Role | +| --- | --- | +| #266 | WS1 closeout parent | +| #267 | This contract (C1) | +| #268 | Full-model workload / clip interval pin in manifest | +| #269 / #270 | Forward / gradient invariance harnesses | +| #276 | Full-model train/infer gate consuming this contract | +| #154 / #108 | Historical contract owners (superseded remaining work → C1) | + +## Migration of existing tests + +Most operator tests still use **private** `atol` / `rtol` literals. That is expected +after C1: the SSOT exists, but call sites have not all moved. + +See the full inventory, priority, and “when it must migrate” map: + +- [WS1 gtest 阈值迁移清单](ws1-gtest-migration-checklist.md) + +How to register ops and run the CLI (post-#267): + +- [gtest usage guide](../contributing/gtest-usage.md) diff --git a/rl_engine/kernels/gtest/__init__.py b/rl_engine/kernels/gtest/__init__.py index c3fc3665..61b43218 100644 --- a/rl_engine/kernels/gtest/__init__.py +++ b/rl_engine/kernels/gtest/__init__.py @@ -2,9 +2,21 @@ # Copyright (c) 2026 RL-Kernel Contributors from .op_checks import CandidateSpec, OperatorCase, run_operator_suite +from .tolerance import ( + BackendProvenance, + ContractResolveError, + ContractSchemaError, + resolve_tolerance_support, + validate_backend_provenance, +) __all__ = [ "CandidateSpec", "OperatorCase", "run_operator_suite", + "BackendProvenance", + "ContractResolveError", + "ContractSchemaError", + "resolve_tolerance_support", + "validate_backend_provenance", ] diff --git a/rl_engine/kernels/gtest/op_checks.py b/rl_engine/kernels/gtest/op_checks.py index efea31a3..085b5f89 100644 --- a/rl_engine/kernels/gtest/op_checks.py +++ b/rl_engine/kernels/gtest/op_checks.py @@ -9,7 +9,13 @@ import torch -from rl_engine.kernels.gtest.tolerance import load_contract +from rl_engine.kernels.gtest.tolerance import ( + BackendProvenance, + ContractResolveError, + load_contract, + resolve_tolerance, + validate_backend_provenance, +) @dataclass(frozen=True) @@ -32,6 +38,7 @@ class CandidateSpec: fn: Callable[..., Any] | Any backend: str = "unknown" arch_key: str | None = None + provenance: BackendProvenance | None = None @dataclass(frozen=True) @@ -48,6 +55,9 @@ class OutputCheck: mean_abs_error: float max_rel_error: float passed: bool + judgment: str + comparison_lhs_role: str + comparison_rhs_role: str message: str = "" @@ -73,6 +83,7 @@ class CandidateReport: pass_rate: float passed: bool cases: list[CaseCheck] + backend_provenance: BackendProvenance | None = None @dataclass(frozen=True) @@ -140,6 +151,19 @@ def _run_candidate( grad_mode: str, grad_seed: int, ) -> CandidateReport: + if candidate.provenance is not None: + validate_backend_provenance(contract, candidate.provenance) + if candidate.backend != candidate.provenance.actual_backend: + raise ContractResolveError( + f"candidate backend {candidate.backend!r} disagrees with reported actual backend " + f"{candidate.provenance.actual_backend!r}" + ) + for case in cases: + if _dtype_name(case.dtype) != candidate.provenance.execution_dtype: + raise ContractResolveError( + f"case {case.name!r} dtype {case.dtype} does not match " + f"provenance execution_dtype {candidate.provenance.execution_dtype!r}" + ) if check_grad: case_checks = [ _run_case_backward( @@ -164,6 +188,7 @@ def _run_candidate( pass_rate=pass_rate, passed=passed_outputs == total_outputs, cases=case_checks, + backend_provenance=candidate.provenance, ) @@ -216,14 +241,29 @@ def _run_case_backward( candidate_outputs, gold_outputs, ).outputs - # Reuse the same tolerance class for gradients as for values. This is a - # first conservative default; operator-specific gradient tolerances can be - # split out later if a real backend shows different numerical behavior. + # Gradient thresholds come from the independent gradient_accuracy judgment + # (#267); they must not silently inherit forward_accuracy rows. atol, rtol = _resolve_tolerance( contract, op_class=case.op_class, dtype=case.dtype, arch_key=candidate.arch_key, + backend_profile=(candidate.provenance.backend_profile if candidate.provenance else None), + judgment="gradient_accuracy", + ) + gradient_spec = ( + resolve_tolerance( + contract, + judgment="gradient_accuracy", + op_class=case.op_class, + dtype=case.dtype, + arch_key=candidate.arch_key, + backend_profile=( + candidate.provenance.backend_profile if candidate.provenance else None + ), + ) + if "judgments" in contract + else None ) grad_checks = [ _compare_output( @@ -232,6 +272,13 @@ def _run_case_backward( output_index=len(output_checks) + index, atol=atol, rtol=rtol, + judgment="gradient_accuracy", + comparison_lhs_role=( + gradient_spec.comparison_lhs_role if gradient_spec is not None else "bf16_candidate" + ), + comparison_rhs_role=( + gradient_spec.comparison_rhs_role if gradient_spec is not None else "fp32_reference" + ), message=f"gradient:{name}", ) for index, (name, candidate_grad, gold_grad) in enumerate( @@ -265,7 +312,37 @@ def _compare_case_outputs( op_class=case.op_class, dtype=case.dtype, arch_key=candidate.arch_key, + backend_profile=(candidate.provenance.backend_profile if candidate.provenance else None), + judgment="forward_accuracy", ) + forward_spec = ( + resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class=case.op_class, + dtype=case.dtype, + arch_key=candidate.arch_key, + backend_profile=( + candidate.provenance.backend_profile if candidate.provenance else None + ), + ) + if "judgments" in contract + else None + ) + if candidate.provenance is not None: + for candidate_output, gold_output in zip(candidate_outputs, gold_outputs, strict=True): + candidate_dtype = _dtype_name(candidate_output.dtype) + gold_dtype = _dtype_name(gold_output.dtype) + if candidate_dtype != candidate.provenance.output_dtype: + raise ContractResolveError( + f"candidate output dtype {candidate_dtype!r} disagrees with provenance " + f"output_dtype {candidate.provenance.output_dtype!r}" + ) + if gold_dtype != candidate.provenance.reference_dtype: + raise ContractResolveError( + f"gold output dtype {gold_dtype!r} disagrees with provenance " + f"reference_dtype {candidate.provenance.reference_dtype!r}" + ) output_checks = [ _compare_output( candidate_output, @@ -273,6 +350,13 @@ def _compare_case_outputs( output_index=index, atol=atol, rtol=rtol, + judgment="forward_accuracy", + comparison_lhs_role=( + forward_spec.comparison_lhs_role if forward_spec is not None else "bf16_candidate" + ), + comparison_rhs_role=( + forward_spec.comparison_rhs_role if forward_spec is not None else "fp32_reference" + ), ) for index, (candidate_output, gold_output) in enumerate( zip(candidate_outputs, gold_outputs, strict=True) @@ -407,7 +491,27 @@ def _resolve_tolerance( op_class: str, dtype: torch.dtype, arch_key: str | None = None, + backend_profile: str | None = None, + judgment: str = "forward_accuracy", ) -> tuple[float, float]: + """Resolve thresholds via the shared four-judgment contract (#267). + + Falls back to the legacy ``accuracy`` mirror only when the four-judgment + block is absent (older fixture contracts in unit tests). + """ + + if "judgments" in contract: + spec = resolve_tolerance( + contract, + judgment=judgment, + op_class=op_class, + dtype=dtype, + arch_key=arch_key, + backend_profile=backend_profile, + ) + return float(spec.atol), float(spec.rtol) + + # Legacy fixtures used by some unit tests that inject a minimal contract. dtype_name = _dtype_name(dtype) if arch_key is not None: arch_values = ( @@ -441,6 +545,9 @@ def _compare_output( output_index: int, atol: float, rtol: float, + judgment: str = "forward_accuracy", + comparison_lhs_role: str = "bf16_candidate", + comparison_rhs_role: str = "fp32_reference", message: str = "", ) -> OutputCheck: if candidate.shape != gold.shape: @@ -455,6 +562,9 @@ def _compare_output( mean_abs_error=float("inf"), max_rel_error=float("inf"), passed=False, + judgment=judgment, + comparison_lhs_role=comparison_lhs_role, + comparison_rhs_role=comparison_rhs_role, message=f"shape mismatch: candidate={tuple(candidate.shape)} gold={tuple(gold.shape)}", ) @@ -482,6 +592,9 @@ def _compare_output( mean_abs_error=mean_abs_error, max_rel_error=max_rel_error, passed=bool(torch.allclose(candidate_fp32, gold_fp32, atol=atol, rtol=rtol)), + judgment=judgment, + comparison_lhs_role=comparison_lhs_role, + comparison_rhs_role=comparison_rhs_role, message=message, ) diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index d0481e83..230dd95d 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -1,20 +1,977 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +"""WS1 numerical contract loader and resolver (#267 / C1 of #266). + +This module is the sole authority for: +- dtype policy (BF16 execution, FP32 reference/accumulation, FP8 out) +- four-judgment tolerances +- comparison roles +- chain-level logprob aggregates (max_abs_dlogp / approx_kl0 / clipfrac0) + +Gates must obtain thresholds only through the resolvers defined here. +""" + from __future__ import annotations import json +import math +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any +from typing import Any, Mapping, Sequence _CONTRACT_PATH = Path(__file__).with_name("tolerance_contract.json") +JUDGMENTS = ( + "forward_accuracy", + "forward_invariance", + "gradient_accuracy", + "gradient_invariance", +) +OP_CLASSES = ("elementwise", "reduction", "logprob", "attention") +MANDATORY_DTYPES = ("float32", "bfloat16") +OPTIONAL_DTYPES = ("float16",) +OUT_OF_SCOPE_DTYPES = ("float8",) +ALL_DTYPES = MANDATORY_DTYPES + OPTIONAL_DTYPES + OUT_OF_SCOPE_DTYPES +CHAIN_AGGREGATE_METRICS = ("max_abs_dlogp", "approx_kl0", "clipfrac0") +INVARIANCE_JUDGMENTS = ("forward_invariance", "gradient_invariance") +REPORT_KINDS = ( + "forward_accuracy", + "forward_invariance", + "train_infer_logprob_parity", + "gradient_accuracy", + "gradient_invariance", +) + + +class ContractError(ValueError): + """Base error for contract load / resolve failures.""" + + +class ContractSchemaError(ContractError): + """Contract JSON failed schema validation.""" + + +class ContractResolveError(ContractError): + """A resolve request cannot be satisfied under the contract.""" + + +@dataclass(frozen=True) +class DtypePolicy: + """Resolved WS1 dtype / TF32 / FP8 policy.""" + + execution_dtype: str + accumulation_dtype: str + reference_dtype: str + output_dtype_default: str + logprob_aggregates_dtype: str + fp8: str + fp16_status: str + tf32_reference: str + tf32_candidate_execution: str + backend_profiles: tuple[str, ...] + backend_private_tolerance_relaxation: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class BackendProvenance: + """Actual backend and dtype facts persisted by a WS1 report.""" + + backend_profile: str + requested_backend: str + actual_backend: str + execution_dtype: str + accumulation_dtype: str + output_dtype: str + reference_dtype: str + candidate_tf32_enabled: bool + reference_tf32_enabled: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ToleranceSupport: + """Schema-level support result, including explicit N/A and out-of-scope cells.""" + + judgment: str + op_class: str + dtype_name: str + status: str + reason: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ComparisonRoles: + """lhs/rhs roles for a report kind.""" + + report_kind: str + comparison_lhs_role: str + comparison_rhs_role: str + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ToleranceSpec: + """Resolved tolerance for one (judgment, op_class, dtype) request.""" + + judgment: str + op_class: str + dtype_name: str + status: str + mode: str + atol: float + rtol: float + comparison_lhs_role: str + comparison_rhs_role: str + backend_profile: str | None = None + arch_key: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) -def load_contract(path: str | Path = _CONTRACT_PATH) -> dict[str, Any]: - """Load the dtype/operator-class tolerance contract.""" + +@dataclass(frozen=True) +class LogprobAggregates: + """Three chain-level logprob aggregates (FP32).""" + + max_abs_dlogp: float + approx_kl0: float + clipfrac0: float + active_token_count: int + clip_interval: tuple[float, float] + report_kind: str + comparison_lhs_role: str + comparison_rhs_role: str + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["clip_interval"] = list(self.clip_interval) + return data + + +@dataclass(frozen=True) +class AggregateMetricVerdict: + metric: str + value: float + threshold: float + passed: bool + + +@dataclass(frozen=True) +class LogprobAggregateVerdict: + aggregates: LogprobAggregates + metrics: tuple[AggregateMetricVerdict, ...] + passed: bool + report_kind: str + comparison_lhs_role: str + comparison_rhs_role: str + + def to_dict(self) -> dict[str, Any]: + return { + "aggregates": self.aggregates.to_dict(), + "metrics": [asdict(m) for m in self.metrics], + "passed": self.passed, + "report_kind": self.report_kind, + "comparison_lhs_role": self.comparison_lhs_role, + "comparison_rhs_role": self.comparison_rhs_role, + } + + +def load_contract( + path: str | Path = _CONTRACT_PATH, + *, + validate: bool = True, +) -> dict[str, Any]: + """Load the WS1 dtype/operator-class tolerance contract.""" with Path(path).open("r", encoding="utf-8") as handle: - return json.load(handle) + contract = json.load(handle) + if validate: + validate_contract_schema(contract) + return contract + + +def validate_contract_schema(contract: Mapping[str, Any]) -> None: + """Validate four-judgment schema, dtype policy, roles, and aggregates.""" + + if not isinstance(contract, Mapping): + raise ContractSchemaError("contract must be a mapping") + + for key in ( + "version", + "policy", + "comparison_roles", + "judgments", + "chain_logprob_aggregates", + ): + if key not in contract: + raise ContractSchemaError(f"contract missing required key {key!r}") + + _validate_policy(contract["policy"]) + _validate_comparison_roles(contract["comparison_roles"]) + _validate_judgments(contract["judgments"]) + _validate_chain_aggregates(contract["chain_logprob_aggregates"]) + _validate_compat_views(contract) + + +def resolve_dtype_policy(contract: Mapping[str, Any]) -> DtypePolicy: + """Resolve independent execution / accumulation / output / reference dtypes.""" + + policy = contract["policy"] + output = policy["output_dtype"] + tf32 = policy["tf32"] + fp16 = policy["fp16"] + return DtypePolicy( + execution_dtype=str(policy["execution_dtype"]), + accumulation_dtype=str(policy["accumulation_dtype"]), + reference_dtype=str(policy["reference_dtype"]), + output_dtype_default=( + str(policy["execution_dtype"]) + if output["default"] == "execution" + else str(output["default"]) + ), + logprob_aggregates_dtype=str(output["logprob_aggregates"]), + fp8=str(policy["fp8"]), + fp16_status=str(fp16["status"]), + tf32_reference=str(tf32["reference"]), + tf32_candidate_execution=str(tf32["candidate_execution"]), + backend_profiles=tuple(str(p) for p in policy["backend_profiles"]), + backend_private_tolerance_relaxation=bool(policy["backend_private_tolerance_relaxation"]), + ) + + +def validate_backend_provenance( + contract: Mapping[str, Any], + provenance: BackendProvenance, +) -> BackendProvenance: + """Fail closed when reported backend or dtype facts violate the WS1 profile.""" + + policy = resolve_dtype_policy(contract) + if provenance.backend_profile not in policy.backend_profiles: + raise ContractResolveError(f"unknown backend_profile {provenance.backend_profile!r}") + profile_contract = contract["policy"]["backend_profile_contracts"][provenance.backend_profile] + expected_backend = str(profile_contract["backend_family"]) + for field_name, actual in ( + ("requested_backend", provenance.requested_backend), + ("actual_backend", provenance.actual_backend), + ): + if actual != expected_backend: + raise ContractResolveError( + f"backend provenance mismatch for {field_name}: expected " + f"{expected_backend!r}, got {actual!r}" + ) + + expected_dtypes = { + "execution_dtype": policy.execution_dtype, + "accumulation_dtype": policy.accumulation_dtype, + "output_dtype": policy.output_dtype_default, + "reference_dtype": policy.reference_dtype, + } + for field_name, expected in expected_dtypes.items(): + actual = _dtype_name(getattr(provenance, field_name)) + if actual != expected: + raise ContractResolveError( + f"backend provenance mismatch for {field_name}: expected " + f"{expected!r}, got {actual!r}" + ) + for field_name in ("candidate_tf32_enabled", "reference_tf32_enabled"): + if getattr(provenance, field_name): + raise ContractResolveError( + f"backend provenance reports {field_name}=true; WS1 requires disabled" + ) + return provenance + + +def resolve_comparison_roles( + contract: Mapping[str, Any], + report_kind: str, +) -> ComparisonRoles: + """Return lhs/rhs roles for a report kind.""" + + roles_root = contract["comparison_roles"] + forbidden = set(roles_root.get("forbidden", ())) + by_kind = roles_root["by_report_kind"] + if report_kind not in by_kind: + raise ContractResolveError(f"unknown report_kind {report_kind!r}") + entry = by_kind[report_kind] + lhs = str(entry["comparison_lhs_role"]) + rhs = str(entry["comparison_rhs_role"]) + for role in (lhs, rhs): + if role in forbidden: + raise ContractResolveError( + f"forbidden comparison role {role!r} for report_kind {report_kind!r}" + ) + if role not in roles_root["allowed"]: + raise ContractResolveError( + f"unknown comparison role {role!r} for report_kind {report_kind!r}" + ) + return ComparisonRoles( + report_kind=report_kind, + comparison_lhs_role=lhs, + comparison_rhs_role=rhs, + ) + + +def assert_comparison_roles( + contract: Mapping[str, Any], + report_kind: str, + comparison_lhs_role: str, + comparison_rhs_role: str, +) -> ComparisonRoles: + """Hard-fail if report roles are reversed, unknown, or forbidden.""" + + expected = resolve_comparison_roles(contract, report_kind) + if comparison_lhs_role in contract["comparison_roles"].get("forbidden", ()): + raise ContractResolveError(f"forbidden comparison_lhs_role {comparison_lhs_role!r}") + if comparison_rhs_role in contract["comparison_roles"].get("forbidden", ()): + raise ContractResolveError(f"forbidden comparison_rhs_role {comparison_rhs_role!r}") + if ( + comparison_lhs_role != expected.comparison_lhs_role + or comparison_rhs_role != expected.comparison_rhs_role + ): + raise ContractResolveError( + f"role mismatch for {report_kind!r}: expected " + f"lhs={expected.comparison_lhs_role!r}, rhs={expected.comparison_rhs_role!r}; " + f"got lhs={comparison_lhs_role!r}, rhs={comparison_rhs_role!r}" + ) + return expected + + +def resolve_tolerance( + contract: Mapping[str, Any], + *, + judgment: str, + op_class: str, + dtype: str | Any, + arch_key: str | None = None, + backend_profile: str | None = None, +) -> ToleranceSpec: + """Resolve one four-judgment tolerance cell. + + ``cuda_bf16`` and ``triton_cuda_bf16`` share the same rows. Backend-private + threshold relaxation is forbidden. + """ + + if judgment not in JUDGMENTS: + raise ContractResolveError(f"unknown judgment {judgment!r}") + if op_class not in OP_CLASSES: + raise ContractResolveError(f"unknown op_class {op_class!r}") + + dtype_name = _dtype_name(dtype) + policy = resolve_dtype_policy(contract) + + if backend_profile is not None: + if backend_profile not in policy.backend_profiles: + raise ContractResolveError( + f"unknown backend_profile {backend_profile!r}; " + f"allowed={list(policy.backend_profiles)}" + ) + if policy.backend_private_tolerance_relaxation: + raise ContractResolveError( + "backend_private_tolerance_relaxation must remain false under WS1 C1" + ) + + if dtype_name in OUT_OF_SCOPE_DTYPES or policy.fp8 == "out_of_scope" and dtype_name == "float8": + raise ContractResolveError( + f"dtype {dtype_name!r} is out of scope for WS1 (FP8 requests hard-fail)" + ) + + support = resolve_tolerance_support( + contract, + judgment=judgment, + op_class=op_class, + dtype=dtype_name, + arch_key=arch_key, + ) + judgment_root = contract["judgments"][judgment] + cell = _lookup_cell(judgment_root, op_class=op_class, dtype_name=dtype_name, arch_key=arch_key) + if cell is None: + raise ContractResolveError( + f"missing declared cell for judgment={judgment!r}, " + f"op_class={op_class!r}, dtype={dtype_name!r}" + ) + + status = support.status + if status == "out_of_scope": + raise ContractResolveError( + f"cell out_of_scope for judgment={judgment!r}, " + f"op_class={op_class!r}, dtype={dtype_name!r}" + ) + if status == "not_applicable": + raise ContractResolveError( + f"cell not_applicable for judgment={judgment!r}, " + f"op_class={op_class!r}, dtype={dtype_name!r}; " + "callers must not request non-applicable judgments without an explicit N/A path" + ) + if status not in {"applicable", "optional"}: + raise ContractResolveError( + f"invalid status {status!r} for judgment={judgment!r}, " + f"op_class={op_class!r}, dtype={dtype_name!r}" + ) + + mode = str(cell.get("mode", judgment_root.get("default_mode", "tolerance"))) + if "atol" not in cell or "rtol" not in cell: + raise ContractResolveError( + f"cell missing atol/rtol for judgment={judgment!r}, " + f"op_class={op_class!r}, dtype={dtype_name!r}" + ) + atol = float(cell["atol"]) + rtol = float(cell["rtol"]) + + if judgment in INVARIANCE_JUDGMENTS and status == "applicable": + if mode != "bitwise" or atol != 0.0 or rtol != 0.0: + raise ContractResolveError( + f"Batch/Chunk invariance requires bitwise atol=0 rtol=0; got " + f"mode={mode!r}, atol={atol}, rtol={rtol} for {judgment}/{op_class}/{dtype_name}" + ) + + roles = resolve_comparison_roles(contract, judgment) + return ToleranceSpec( + judgment=judgment, + op_class=op_class, + dtype_name=dtype_name, + status=status, + mode=mode, + atol=atol, + rtol=rtol, + comparison_lhs_role=roles.comparison_lhs_role, + comparison_rhs_role=roles.comparison_rhs_role, + backend_profile=backend_profile, + arch_key=arch_key, + ) + + +def resolve_tolerance_support( + contract: Mapping[str, Any], + *, + judgment: str, + op_class: str, + dtype: str | Any, + arch_key: str | None = None, +) -> ToleranceSupport: + """Resolve schema support without pretending N/A cells have thresholds.""" + + if judgment not in JUDGMENTS: + raise ContractResolveError(f"unknown judgment {judgment!r}") + if op_class not in OP_CLASSES: + raise ContractResolveError(f"unknown op_class {op_class!r}") + dtype_name = _dtype_name(dtype) + cell = _lookup_cell( + contract["judgments"][judgment], + op_class=op_class, + dtype_name=dtype_name, + arch_key=arch_key, + ) + if cell is None: + raise ContractResolveError( + f"missing declared cell for judgment={judgment!r}, " + f"op_class={op_class!r}, dtype={dtype_name!r}" + ) + status = str(cell.get("status", "")) + if status not in {"applicable", "optional", "not_applicable", "out_of_scope"}: + raise ContractResolveError( + f"invalid support status {status!r} for {judgment}/{op_class}/{dtype_name}" + ) + reason = cell.get("reason") + return ToleranceSupport( + judgment=judgment, + op_class=op_class, + dtype_name=dtype_name, + status=status, + reason=str(reason) if reason is not None else None, + ) + + +def resolve_chain_aggregate_thresholds( + contract: Mapping[str, Any], + metric_name: str, + execution_dtype: str | Any, +) -> float: + """Named resolve for max_abs_dlogp / approx_kl0 / clipfrac0 thresholds.""" + + if metric_name not in CHAIN_AGGREGATE_METRICS: + raise ContractResolveError( + f"unknown chain aggregate metric {metric_name!r}; " + f"only {list(CHAIN_AGGREGATE_METRICS)} are allowed" + ) + dtype_name = _dtype_name(execution_dtype) + metrics = contract["chain_logprob_aggregates"]["metrics"] + by_dtype = metrics[metric_name]["by_execution_dtype"] + if dtype_name not in by_dtype: + raise ContractResolveError( + f"missing chain aggregate threshold for metric={metric_name!r}, " + f"execution_dtype={dtype_name!r}" + ) + return float(by_dtype[dtype_name]["threshold"]) + + +def compute_logprob_aggregates( + lhs_logp: Any, + rhs_logp: Any, + active_mask: Any, + *, + contract: Mapping[str, Any], + report_kind: str, + clip_interval: Sequence[float] | tuple[float, float], + comparison_lhs_role: str, + comparison_rhs_role: str, +) -> LogprobAggregates: + """Compute the three chain-level logprob aggregates in FP32. + + ``dlogp = lhs_logp - rhs_logp`` on active selected tokens only. + Empty active set / NaN / Inf → hard fail. + """ + + assert_comparison_roles(contract, report_kind, comparison_lhs_role, comparison_rhs_role) + + try: + import torch + except ImportError as exc: # pragma: no cover + raise ContractResolveError("torch is required for aggregate computation") from exc + + if len(clip_interval) != 2: + raise ContractResolveError("clip_interval must be a length-2 [lo, hi] pair") + lo, hi = float(clip_interval[0]), float(clip_interval[1]) + if not (lo < hi): + raise ContractResolveError(f"clip_interval requires lo < hi, got [{lo}, {hi}]") + + lhs = torch.as_tensor(lhs_logp).detach().float().reshape(-1) + rhs = torch.as_tensor(rhs_logp).detach().float().reshape(-1) + mask = torch.as_tensor(active_mask).detach().reshape(-1).bool() + if lhs.shape != rhs.shape or lhs.shape != mask.shape: + raise ContractResolveError( + f"lhs/rhs/mask shape mismatch: {tuple(lhs.shape)} vs " + f"{tuple(rhs.shape)} vs {tuple(mask.shape)}" + ) + active = int(mask.sum().item()) + if active == 0: + raise ContractResolveError("empty active-token set is a hard fail for logprob aggregates") + + dlogp = lhs[mask] - rhs[mask] + if not torch.isfinite(dlogp).all(): + raise ContractResolveError("NaN/Inf in dlogp is a hard fail for logprob aggregates") + + ratio0 = torch.exp(dlogp) + if not torch.isfinite(ratio0).all(): + raise ContractResolveError("NaN/Inf in ratio0 is a hard fail for logprob aggregates") + + max_abs_dlogp = float(dlogp.abs().max().item()) + approx_kl0 = float((ratio0 - 1.0 - dlogp).mean().item()) + outside = (ratio0 < lo) | (ratio0 > hi) + clipfrac0 = float(outside.float().mean().item()) + + for name, value in ( + ("max_abs_dlogp", max_abs_dlogp), + ("approx_kl0", approx_kl0), + ("clipfrac0", clipfrac0), + ): + if not math.isfinite(value): + raise ContractResolveError(f"NaN/Inf in aggregate {name} is a hard fail") + + return LogprobAggregates( + max_abs_dlogp=max_abs_dlogp, + approx_kl0=approx_kl0, + clipfrac0=clipfrac0, + active_token_count=active, + clip_interval=(lo, hi), + report_kind=report_kind, + comparison_lhs_role=comparison_lhs_role, + comparison_rhs_role=comparison_rhs_role, + ) + + +def judge_logprob_aggregates( + aggregates: LogprobAggregates, + contract: Mapping[str, Any], + *, + execution_dtype: str | Any, + clip_interval: Sequence[float] | tuple[float, float] | None = None, +) -> LogprobAggregateVerdict: + """Judge all three chain logprob aggregates; all must pass.""" + + assert_comparison_roles( + contract, + aggregates.report_kind, + aggregates.comparison_lhs_role, + aggregates.comparison_rhs_role, + ) + + if clip_interval is not None: + lo, hi = float(clip_interval[0]), float(clip_interval[1]) + if (lo, hi) != aggregates.clip_interval: + raise ContractResolveError( + "clip_interval mismatch between compute and judge " + f"(computed={aggregates.clip_interval}, judge=({lo}, {hi}))" + ) + + metrics: list[AggregateMetricVerdict] = [] + for name in CHAIN_AGGREGATE_METRICS: + threshold = resolve_chain_aggregate_thresholds(contract, name, execution_dtype) + value = float(getattr(aggregates, name)) + if not math.isfinite(value): + raise ContractResolveError(f"NaN/Inf in aggregate {name} is a hard fail") + metrics.append( + AggregateMetricVerdict( + metric=name, + value=value, + threshold=threshold, + passed=value <= threshold, + ) + ) + require_all = bool(contract["chain_logprob_aggregates"].get("require_all", True)) + passed = all(m.passed for m in metrics) if require_all else any(m.passed for m in metrics) + return LogprobAggregateVerdict( + aggregates=aggregates, + metrics=tuple(metrics), + passed=passed, + report_kind=aggregates.report_kind, + comparison_lhs_role=aggregates.comparison_lhs_role, + comparison_rhs_role=aggregates.comparison_rhs_role, + ) + + +def default_clip_interval(contract: Mapping[str, Any]) -> tuple[float, float]: + """Return the contract default clip interval for clipfrac0.""" + + interval = contract["chain_logprob_aggregates"]["default_clip_interval"] + return float(interval[0]), float(interval[1]) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _validate_policy(policy: Mapping[str, Any]) -> None: + required = ( + "execution_dtype", + "accumulation_dtype", + "reference_dtype", + "output_dtype", + "fp8", + "fp16", + "tf32", + "backend_profiles", + "backend_profile_contracts", + "backend_private_tolerance_relaxation", + ) + for key in required: + if key not in policy: + raise ContractSchemaError(f"policy missing {key!r}") + if policy["execution_dtype"] != "bfloat16": + raise ContractSchemaError("policy.execution_dtype must be bfloat16 for WS1") + if policy["accumulation_dtype"] != "float32": + raise ContractSchemaError("policy.accumulation_dtype must be float32 for WS1") + if policy["reference_dtype"] != "float32": + raise ContractSchemaError("policy.reference_dtype must be float32 for WS1") + if policy["fp8"] != "out_of_scope": + raise ContractSchemaError("policy.fp8 must be out_of_scope for WS1") + output = policy["output_dtype"] + for key in ("default", "logprob_aggregates"): + if key not in output: + raise ContractSchemaError(f"policy.output_dtype missing {key!r}") + if output["logprob_aggregates"] != "float32": + raise ContractSchemaError("logprob aggregates must be computed in float32") + if output["default"] != "execution": + raise ContractSchemaError("policy.output_dtype.default must follow execution") + if policy["fp16"].get("status") != "optional": + raise ContractSchemaError("policy.fp16.status must be optional for WS1") + tf32 = policy["tf32"] + for key in ("reference", "candidate_execution"): + if key not in tf32: + raise ContractSchemaError(f"policy.tf32 missing {key!r}") + if tf32[key] != "disabled": + raise ContractSchemaError( + f"policy.tf32.{key} must be 'disabled' under the WS1 single policy" + ) + profiles = list(policy["backend_profiles"]) + profile_contracts = policy["backend_profile_contracts"] + required_profile_families = { + "cuda_bf16": "cuda", + "triton_cuda_bf16": "triton", + } + for required_profile, expected_family in required_profile_families.items(): + if required_profile not in profiles: + raise ContractSchemaError(f"policy.backend_profiles must include {required_profile!r}") + if required_profile not in profile_contracts: + raise ContractSchemaError( + f"policy.backend_profile_contracts missing {required_profile!r}" + ) + actual_family = profile_contracts[required_profile].get("backend_family") + if actual_family != expected_family: + raise ContractSchemaError( + f"profile {required_profile!r} requires backend_family " + f"{expected_family!r}, got {actual_family!r}" + ) + if policy["backend_private_tolerance_relaxation"] is not False: + raise ContractSchemaError("backend_private_tolerance_relaxation must be false") + + +def _validate_comparison_roles(roles_root: Mapping[str, Any]) -> None: + for key in ("allowed", "forbidden", "by_report_kind"): + if key not in roles_root: + raise ContractSchemaError(f"comparison_roles missing {key!r}") + forbidden = set(roles_root["forbidden"]) + for name in ("baseline", "singleton_aggregate"): + if name not in forbidden: + raise ContractSchemaError(f"comparison_roles.forbidden must include {name!r}") + by_kind = roles_root["by_report_kind"] + for kind in REPORT_KINDS: + if kind not in by_kind: + raise ContractSchemaError(f"comparison_roles.by_report_kind missing {kind!r}") + entry = by_kind[kind] + for role_key in ("comparison_lhs_role", "comparison_rhs_role"): + if role_key not in entry: + raise ContractSchemaError( + f"comparison_roles.by_report_kind[{kind!r}] missing {role_key!r}" + ) + role = entry[role_key] + if role in forbidden: + raise ContractSchemaError(f"report_kind {kind!r} uses forbidden role {role!r}") + if role not in roles_root["allowed"]: + raise ContractSchemaError(f"report_kind {kind!r} uses unknown role {role!r}") + + +def _validate_judgments(judgments: Mapping[str, Any]) -> None: + for judgment in JUDGMENTS: + if judgment not in judgments: + raise ContractSchemaError(f"judgments missing {judgment!r}") + root = judgments[judgment] + if "by_op_class" not in root: + raise ContractSchemaError(f"judgments[{judgment!r}] missing by_op_class") + by_op = root["by_op_class"] + for op_class in OP_CLASSES: + if op_class not in by_op: + raise ContractSchemaError( + f"judgments[{judgment!r}].by_op_class missing {op_class!r}" + ) + dtype_map = by_op[op_class] + for dtype_name in ALL_DTYPES: + if dtype_name not in dtype_map: + raise ContractSchemaError( + f"missing cell judgments[{judgment!r}][{op_class!r}][{dtype_name!r}]" + ) + cell = dtype_map[dtype_name] + status = cell.get("status") + if status is None: + raise ContractSchemaError( + f"cell missing status: {judgment}/{op_class}/{dtype_name}" + ) + if dtype_name in OUT_OF_SCOPE_DTYPES: + if status != "out_of_scope": + raise ContractSchemaError( + f"FP8 cell must be out_of_scope: {judgment}/{op_class}/{dtype_name}" + ) + continue + if status == "not_applicable" and not cell.get("reason"): + raise ContractSchemaError( + f"not_applicable cell requires reason: {judgment}/{op_class}/{dtype_name}" + ) + if dtype_name in MANDATORY_DTYPES and status not in { + "applicable", + "not_applicable", + }: + # BF16/FP32 must be explicitly applicable (or explicit N/A). + if status != "applicable": + raise ContractSchemaError( + f"mandatory dtype cell must be applicable: " + f"{judgment}/{op_class}/{dtype_name} status={status!r}" + ) + if status in {"applicable", "optional"}: + for thr in ("atol", "rtol", "mode"): + if thr not in cell: + raise ContractSchemaError( + f"cell missing {thr}: {judgment}/{op_class}/{dtype_name}" + ) + if judgment in INVARIANCE_JUDGMENTS and status == "applicable": + mode = cell.get("mode") + atol = float(cell.get("atol", 1.0)) + rtol = float(cell.get("rtol", 1.0)) + if mode != "bitwise" or atol != 0.0 or rtol != 0.0: + raise ContractSchemaError( + f"invariance applicable cells must be bitwise 0/0: " + f"{judgment}/{op_class}/{dtype_name}" + ) + + +def _validate_chain_aggregates(root: Mapping[str, Any]) -> None: + for key in ( + "compute_dtype", + "require_all", + "nan_inf_policy", + "empty_active_token_set", + "dlogp_definition", + "default_clip_interval", + "sole_chain_level_logprob_metrics", + "metrics", + ): + if key not in root: + raise ContractSchemaError(f"chain_logprob_aggregates missing {key!r}") + if root["compute_dtype"] != "float32": + raise ContractSchemaError("chain aggregates must use compute_dtype=float32") + if root["nan_inf_policy"] != "hard_fail": + raise ContractSchemaError("nan_inf_policy must be hard_fail") + if root["empty_active_token_set"] != "hard_fail": + raise ContractSchemaError("empty_active_token_set must be hard_fail") + if not root["require_all"]: + raise ContractSchemaError("require_all must be true for chain logprob aggregates") + sole = list(root["sole_chain_level_logprob_metrics"]) + if set(sole) != set(CHAIN_AGGREGATE_METRICS) or len(sole) != 3: + raise ContractSchemaError( + "sole_chain_level_logprob_metrics must be exactly " f"{list(CHAIN_AGGREGATE_METRICS)}" + ) + interval = root["default_clip_interval"] + if len(interval) != 2 or float(interval[0]) >= float(interval[1]): + raise ContractSchemaError("default_clip_interval must be [lo, hi] with lo < hi") + metrics = root["metrics"] + for name in CHAIN_AGGREGATE_METRICS: + if name not in metrics: + raise ContractSchemaError(f"chain metrics missing {name!r}") + by_dtype = metrics[name].get("by_execution_dtype") + if not isinstance(by_dtype, Mapping): + raise ContractSchemaError(f"metric {name!r} missing by_execution_dtype") + for dtype_name in ("bfloat16", "float32"): + if dtype_name not in by_dtype or "threshold" not in by_dtype[dtype_name]: + raise ContractSchemaError(f"metric {name!r} missing threshold for {dtype_name}") + + +def _validate_compat_views(contract: Mapping[str, Any]) -> None: + """Legacy accuracy / batch_invariance must mirror the four-judgment SSOT.""" + + if "batch_invariance" not in contract: + raise ContractSchemaError("compat key batch_invariance is required") + bi = contract["batch_invariance"] + if float(bi.get("atol", 1.0)) != 0.0 or float(bi.get("rtol", 1.0)) != 0.0: + raise ContractSchemaError("batch_invariance must remain bitwise 0/0") + + if "accuracy" not in contract: + raise ContractSchemaError("compat key accuracy is required") + accuracy = contract["accuracy"]["default"] + fwd = contract["judgments"]["forward_accuracy"]["by_op_class"] + for op_class in OP_CLASSES: + if op_class not in accuracy: + raise ContractSchemaError(f"compat accuracy missing op_class {op_class!r}") + for dtype_name in MANDATORY_DTYPES + OPTIONAL_DTYPES: + if dtype_name not in accuracy[op_class]: + raise ContractSchemaError(f"compat accuracy missing {op_class}/{dtype_name}") + cell = fwd[op_class][dtype_name] + if cell.get("status") not in {"applicable", "optional"}: + continue + acc = accuracy[op_class][dtype_name] + if float(acc["atol"]) != float(cell["atol"]) or float(acc["rtol"]) != float( + cell["rtol"] + ): + raise ContractSchemaError( + f"compat accuracy mismatch vs forward_accuracy for " f"{op_class}/{dtype_name}" + ) + + +def _lookup_cell( + judgment_root: Mapping[str, Any], + *, + op_class: str, + dtype_name: str, + arch_key: str | None, +) -> Mapping[str, Any] | None: + if arch_key is not None: + arch_cell = ( + judgment_root.get("arch_overrides", {}) + .get(arch_key, {}) + .get(op_class, {}) + .get(dtype_name) + ) + if arch_cell is not None: + return arch_cell + return judgment_root.get("by_op_class", {}).get(op_class, {}).get(dtype_name) + + +def _dtype_name(dtype: str | Any) -> str: + if isinstance(dtype, str): + name = dtype + # Accept torch-style aliases. + aliases = { + "torch.float32": "float32", + "torch.bfloat16": "bfloat16", + "torch.float16": "float16", + "torch.float8": "float8", + "fp32": "float32", + "bf16": "bfloat16", + "fp16": "float16", + "fp8": "float8", + } + name = aliases.get(name, name) + if name not in ALL_DTYPES: + raise ContractResolveError(f"unsupported dtype name {dtype!r}") + return name + + # torch.dtype without importing torch at module import time for non-torch tests. + module = getattr(type(dtype), "__module__", "") + qual = getattr(dtype, "name", None) or str(dtype) + if module.startswith("torch") or "torch" in str(type(dtype)): + mapping = { + "torch.float32": "float32", + "torch.bfloat16": "bfloat16", + "torch.float16": "float16", + "float32": "float32", + "bfloat16": "bfloat16", + "float16": "float16", + } + # torch.dtype str is like "torch.float32" + as_str = str(dtype) + if as_str in mapping: + return mapping[as_str] + if qual in mapping: + return mapping[qual] + try: + import torch + + if dtype is torch.float32: + return "float32" + if dtype is torch.bfloat16: + return "bfloat16" + if dtype is torch.float16: + return "float16" + except ImportError: # pragma: no cover + pass + raise ContractResolveError(f"unsupported dtype: {dtype!r}") -__all__ = ["load_contract"] +__all__ = [ + "ALL_DTYPES", + "CHAIN_AGGREGATE_METRICS", + "JUDGMENTS", + "OP_CLASSES", + "AggregateMetricVerdict", + "BackendProvenance", + "ComparisonRoles", + "ContractError", + "ContractResolveError", + "ContractSchemaError", + "DtypePolicy", + "LogprobAggregateVerdict", + "LogprobAggregates", + "ToleranceSpec", + "ToleranceSupport", + "assert_comparison_roles", + "compute_logprob_aggregates", + "default_clip_interval", + "judge_logprob_aggregates", + "load_contract", + "resolve_chain_aggregate_thresholds", + "resolve_comparison_roles", + "resolve_dtype_policy", + "resolve_tolerance", + "resolve_tolerance_support", + "validate_backend_provenance", + "validate_contract_schema", +] diff --git a/rl_engine/kernels/gtest/tolerance_contract.json b/rl_engine/kernels/gtest/tolerance_contract.json index 975ae450..e7645b75 100644 --- a/rl_engine/kernels/gtest/tolerance_contract.json +++ b/rl_engine/kernels/gtest/tolerance_contract.json @@ -1,5 +1,239 @@ { - "batch_invariance": {"atol": 0.0, "rtol": 0.0}, + "version": "ws1-c1-v1", + "policy": { + "execution_dtype": "bfloat16", + "accumulation_dtype": "float32", + "reference_dtype": "float32", + "output_dtype": { + "default": "execution", + "logprob_aggregates": "float32" + }, + "fp8": "out_of_scope", + "fp16": { + "status": "optional", + "note": "FP16 rows are complete when declared; not mandatory for WS1 EXIT." + }, + "tf32": { + "reference": "disabled", + "candidate_execution": "disabled", + "policy": "Repo-wide single policy: TF32 is disabled for FP32 reference and for candidate execution under this contract." + }, + "backend_profiles": ["cuda_bf16", "triton_cuda_bf16"], + "backend_profile_contracts": { + "cuda_bf16": {"backend_family": "cuda"}, + "triton_cuda_bf16": {"backend_family": "triton"} + }, + "backend_private_tolerance_relaxation": false + }, + "comparison_roles": { + "allowed": [ + "bf16_candidate", + "fp32_reference", + "canonical_config", + "transformed_config", + "training_style_teacher_forcing", + "inference_style_rollout_decode" + ], + "forbidden": ["baseline", "singleton_aggregate"], + "by_report_kind": { + "forward_accuracy": { + "comparison_lhs_role": "bf16_candidate", + "comparison_rhs_role": "fp32_reference" + }, + "forward_invariance": { + "comparison_lhs_role": "transformed_config", + "comparison_rhs_role": "canonical_config" + }, + "train_infer_logprob_parity": { + "comparison_lhs_role": "training_style_teacher_forcing", + "comparison_rhs_role": "inference_style_rollout_decode" + }, + "gradient_accuracy": { + "comparison_lhs_role": "bf16_candidate", + "comparison_rhs_role": "fp32_reference" + }, + "gradient_invariance": { + "comparison_lhs_role": "transformed_config", + "comparison_rhs_role": "canonical_config" + } + } + }, + "judgments": { + "forward_accuracy": { + "default_mode": "tolerance", + "by_op_class": { + "elementwise": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-5, "rtol": 1.0e-5}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 2.0e-2, "rtol": 1.6e-2}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 1.0e-3, "rtol": 1.0e-3}, + "float8": {"status": "out_of_scope"} + }, + "reduction": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-4, "rtol": 1.0e-4}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 5.0e-2, "rtol": 2.0e-2}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 1.0e-3, "rtol": 1.0e-3}, + "float8": {"status": "out_of_scope"} + }, + "logprob": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-5, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 5.0e-2, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 5.0e-3, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "attention": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-4, "rtol": 1.0e-4}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 5.0e-2, "rtol": 2.0e-2}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 1.0e-3, "rtol": 1.0e-3}, + "float8": {"status": "out_of_scope"} + } + }, + "arch_overrides": { + "sm90": {} + } + }, + "forward_invariance": { + "default_mode": "bitwise", + "scope": "batch_chunk_padding_layout", + "by_op_class": { + "elementwise": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "reduction": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "logprob": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "attention": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + } + }, + "arch_overrides": { + "sm90": {} + } + }, + "gradient_accuracy": { + "default_mode": "tolerance", + "by_op_class": { + "elementwise": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-5, "rtol": 1.0e-5}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 2.0e-2, "rtol": 1.6e-2}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 1.0e-3, "rtol": 1.0e-3}, + "float8": {"status": "out_of_scope"} + }, + "reduction": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-4, "rtol": 1.0e-4}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 5.0e-2, "rtol": 2.0e-2}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 1.0e-3, "rtol": 1.0e-3}, + "float8": {"status": "out_of_scope"} + }, + "logprob": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-5, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 5.0e-2, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 5.0e-3, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "attention": { + "float32": {"status": "applicable", "mode": "tolerance", "atol": 1.0e-4, "rtol": 1.0e-4}, + "bfloat16": {"status": "applicable", "mode": "tolerance", "atol": 5.0e-2, "rtol": 2.0e-2}, + "float16": {"status": "optional", "mode": "tolerance", "atol": 1.0e-3, "rtol": 1.0e-3}, + "float8": {"status": "out_of_scope"} + } + }, + "arch_overrides": { + "sm90": {} + } + }, + "gradient_invariance": { + "default_mode": "bitwise", + "scope": "batch_chunk_padding_layout", + "by_op_class": { + "elementwise": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "reduction": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "logprob": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + }, + "attention": { + "float32": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "bfloat16": {"status": "applicable", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float16": {"status": "optional", "mode": "bitwise", "atol": 0.0, "rtol": 0.0}, + "float8": {"status": "out_of_scope"} + } + }, + "arch_overrides": { + "sm90": {} + } + } + }, + "chain_logprob_aggregates": { + "compute_dtype": "float32", + "require_all": true, + "nan_inf_policy": "hard_fail", + "empty_active_token_set": "hard_fail", + "active_token_policy": "active selected tokens only", + "dlogp_definition": "comparison_lhs_logp - comparison_rhs_logp", + "clip_interval_field": "clip_interval", + "default_clip_interval": [0.8, 1.2], + "sole_chain_level_logprob_metrics": [ + "max_abs_dlogp", + "approx_kl0", + "clipfrac0" + ], + "metrics": { + "max_abs_dlogp": { + "formula": "max(abs(dlogp))", + "pass_rule": "value <= threshold", + "by_execution_dtype": { + "bfloat16": {"threshold": 5.0e-2}, + "float32": {"threshold": 1.0e-5}, + "float16": {"threshold": 5.0e-3} + } + }, + "approx_kl0": { + "formula": "mean(exp(dlogp) - 1 - dlogp)", + "pass_rule": "value <= threshold", + "by_execution_dtype": { + "bfloat16": {"threshold": 5.0e-2}, + "float32": {"threshold": 1.0e-5}, + "float16": {"threshold": 5.0e-3} + } + }, + "clipfrac0": { + "formula": "mean(1[exp(dlogp) outside clip_interval])", + "pass_rule": "value <= threshold", + "by_execution_dtype": { + "bfloat16": {"threshold": 0.0}, + "float32": {"threshold": 0.0}, + "float16": {"threshold": 0.0} + } + } + } + }, "accuracy": { "default": { "elementwise": { @@ -26,5 +260,6 @@ "arch_overrides": { "sm90": {} } - } + }, + "batch_invariance": {"atol": 0.0, "rtol": 0.0} } diff --git a/tests/test_op_checks.py b/tests/test_op_checks.py index e076e106..35de05a5 100644 --- a/tests/test_op_checks.py +++ b/tests/test_op_checks.py @@ -5,6 +5,7 @@ import argparse +import pytest import torch from rl_engine.kernels.gtest.op_checks import CandidateSpec, OperatorCase, run_operator_suite @@ -13,6 +14,7 @@ make_operator_case, operator_names, ) +from rl_engine.kernels.gtest.tolerance import BackendProvenance, ContractResolveError from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp from rl_engine.kernels.ops.pytorch.linear.lm_head import NativeLMHeadOp from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp @@ -119,7 +121,11 @@ def test_embedding_native_candidate_suite_passes_issue_108_helper(): ) assert report.passed - assert report.candidates[0].cases[0].outputs[1].message == "gradient:weight" + gradient = report.candidates[0].cases[0].outputs[1] + assert gradient.message == "gradient:weight" + assert gradient.judgment == "gradient_accuracy" + assert gradient.comparison_lhs_role == "bf16_candidate" + assert gradient.comparison_rhs_role == "fp32_reference" def test_lm_head_native_candidate_suite_passes_issue_108_helper(): @@ -182,6 +188,97 @@ def test_suite_report_to_dict_contains_error_metrics(): assert "passed" in output +def test_ws1_report_persists_roles_and_backend_provenance(): + provenance = BackendProvenance( + backend_profile="cuda_bf16", + requested_backend="cuda", + actual_backend="cuda", + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + ) + report = run_operator_suite( + "logp", + candidates=[ + CandidateSpec( + name="cuda-logp", + backend="cuda", + fn=NativeLogpOp(), + provenance=provenance, + ) + ], + cases=[_logp_case("bf16", torch.bfloat16, seed=12)], + ) + output = report.candidates[0].cases[0].outputs[0] + assert output.judgment == "forward_accuracy" + assert output.comparison_lhs_role == "bf16_candidate" + assert output.comparison_rhs_role == "fp32_reference" + data = report.to_dict()["candidates"][0] + assert data["backend_provenance"]["actual_backend"] == "cuda" + assert "baseline" not in data["cases"][0]["outputs"][0] + + +def test_ws1_report_rejects_backend_provenance_mismatch(): + provenance = BackendProvenance( + backend_profile="cuda_bf16", + requested_backend="cuda", + actual_backend="triton", + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + ) + with pytest.raises(ContractResolveError, match="actual_backend"): + run_operator_suite( + "logp", + candidates=[ + CandidateSpec( + name="bad", + backend="triton", + fn=NativeLogpOp(), + provenance=provenance, + ) + ], + cases=[_logp_case("bf16", torch.bfloat16, seed=13)], + ) + + +def test_ws1_report_checks_observed_output_dtype_against_provenance(): + provenance = BackendProvenance( + backend_profile="cuda_bf16", + requested_backend="cuda", + actual_backend="cuda", + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + ) + + def wrong_output_dtype(logits, token_ids): + return NativeLogpOp().forward(logits, token_ids).float() + + with pytest.raises(ContractResolveError, match="candidate output dtype"): + run_operator_suite( + "logp", + candidates=[ + CandidateSpec( + name="wrong-output", + backend="cuda", + fn=wrong_output_dtype, + provenance=provenance, + ) + ], + cases=[_logp_case("bf16", torch.bfloat16, seed=14)], + ) + + def test_candidate_arch_key_uses_tolerance_override(): def slightly_shifted_logp(logits, token_ids): return NativeLogpOp().forward_fp32(logits, token_ids) + 0.02 diff --git a/tests/test_tolerance_contract.py b/tests/test_tolerance_contract.py index 5eb75cbd..5f711061 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -1,9 +1,36 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +"""Schema and resolver tests for WS1 C1 four-judgment contract (#267).""" + from __future__ import annotations -from rl_engine.kernels.gtest.tolerance import load_contract +import copy +import math + +import pytest +import torch + +from rl_engine.kernels.gtest.tolerance import ( + CHAIN_AGGREGATE_METRICS, + JUDGMENTS, + OP_CLASSES, + BackendProvenance, + ContractResolveError, + ContractSchemaError, + assert_comparison_roles, + compute_logprob_aggregates, + default_clip_interval, + judge_logprob_aggregates, + load_contract, + resolve_chain_aggregate_thresholds, + resolve_comparison_roles, + resolve_dtype_policy, + resolve_tolerance, + resolve_tolerance_support, + validate_backend_provenance, + validate_contract_schema, +) def test_load_contract_contains_expected_operator_classes(): @@ -34,3 +61,420 @@ def test_attention_bfloat16_tolerance_matches_contract(): tolerance = contract["accuracy"]["default"]["attention"]["bfloat16"] assert tolerance["atol"] >= 5.0e-2 assert tolerance["rtol"] >= 2.0e-2 + + +def test_contract_schema_validates_on_load(): + contract = load_contract(validate=True) + validate_contract_schema(contract) + + +def test_dtype_policy_locks_bf16_fp32_fp8_tf32(): + policy = resolve_dtype_policy(load_contract()) + assert policy.execution_dtype == "bfloat16" + assert policy.accumulation_dtype == "float32" + assert policy.reference_dtype == "float32" + assert policy.output_dtype_default == "bfloat16" + assert policy.logprob_aggregates_dtype == "float32" + assert policy.fp8 == "out_of_scope" + assert policy.fp16_status == "optional" + assert policy.tf32_reference == "disabled" + assert policy.tf32_candidate_execution == "disabled" + assert "cuda_bf16" in policy.backend_profiles + assert "triton_cuda_bf16" in policy.backend_profiles + assert policy.backend_private_tolerance_relaxation is False + + +def test_four_judgments_present_and_complete(): + contract = load_contract() + assert set(contract["judgments"]) == set(JUDGMENTS) + for judgment in JUDGMENTS: + by_op = contract["judgments"][judgment]["by_op_class"] + assert set(by_op) == set(OP_CLASSES) + for op_class in OP_CLASSES: + for dtype_name in ("float32", "bfloat16", "float16", "float8"): + assert dtype_name in by_op[op_class] + + +def test_invariance_rows_are_bitwise_zero(): + contract = load_contract() + for judgment in ("forward_invariance", "gradient_invariance"): + for op_class in OP_CLASSES: + for dtype_name in ("float32", "bfloat16"): + spec = resolve_tolerance( + contract, + judgment=judgment, + op_class=op_class, + dtype=dtype_name, + ) + assert spec.mode == "bitwise" + assert spec.atol == 0.0 + assert spec.rtol == 0.0 + + +def test_cuda_and_triton_profiles_share_thresholds(): + contract = load_contract() + for profile in ("cuda_bf16", "triton_cuda_bf16"): + a = resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="logprob", + dtype="bfloat16", + backend_profile=profile, + ) + b = resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="logprob", + dtype="bfloat16", + backend_profile="cuda_bf16" if profile != "cuda_bf16" else "triton_cuda_bf16", + ) + assert a.atol == b.atol and a.rtol == b.rtol and a.mode == b.mode + + +def test_unknown_backend_profile_hard_fails(): + contract = load_contract() + with pytest.raises(ContractResolveError, match="backend_profile"): + resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="reduction", + dtype="bfloat16", + backend_profile="private_backend", + ) + + +def test_backend_provenance_checks_profile_backend_and_all_dtypes(): + contract = load_contract() + provenance = BackendProvenance( + backend_profile="cuda_bf16", + requested_backend="cuda", + actual_backend="cuda", + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + ) + assert validate_backend_provenance(contract, provenance) == provenance + with pytest.raises(ContractResolveError, match="actual_backend"): + validate_backend_provenance( + contract, + BackendProvenance(**{**provenance.to_dict(), "actual_backend": "triton"}), + ) + with pytest.raises(ContractResolveError, match="output_dtype"): + validate_backend_provenance( + contract, + BackendProvenance(**{**provenance.to_dict(), "output_dtype": "float32"}), + ) + with pytest.raises(ContractResolveError, match="candidate_tf32_enabled"): + validate_backend_provenance( + contract, + BackendProvenance(**{**provenance.to_dict(), "candidate_tf32_enabled": True}), + ) + + +def test_not_applicable_has_explicit_support_result_but_no_threshold(): + contract = copy.deepcopy(load_contract()) + cell = contract["judgments"]["forward_accuracy"]["by_op_class"]["elementwise"]["float16"] + cell["status"] = "not_applicable" + cell["reason"] = "profile does not declare FP16" + validate_contract_schema(contract) + support = resolve_tolerance_support( + contract, judgment="forward_accuracy", op_class="elementwise", dtype="float16" + ) + assert support.status == "not_applicable" + with pytest.raises(ContractResolveError, match="not_applicable"): + resolve_tolerance( + contract, judgment="forward_accuracy", op_class="elementwise", dtype="float16" + ) + + +def test_fp8_request_hard_fails(): + contract = load_contract() + with pytest.raises(ContractResolveError, match="out of scope"): + resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="reduction", + dtype="float8", + ) + + +def test_missing_applicable_cell_hard_fails(): + contract = copy.deepcopy(load_contract()) + del contract["judgments"]["forward_accuracy"]["by_op_class"]["attention"]["bfloat16"] + with pytest.raises(ContractSchemaError): + validate_contract_schema(contract) + # Resolver path: re-insert schema-invalid by skipping validate, then resolve. + with pytest.raises(ContractResolveError, match="missing declared cell"): + resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="attention", + dtype="bfloat16", + ) + + +def test_gradient_thresholds_do_not_inherit_forward(): + contract = copy.deepcopy(load_contract()) + # Mutate only forward_accuracy BF16 reduction. + contract["judgments"]["forward_accuracy"]["by_op_class"]["reduction"]["bfloat16"]["atol"] = 9.9 + # Keep compat mirror in sync is not required for this unit test of independence. + fwd = resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="reduction", + dtype="bfloat16", + ) + grad = resolve_tolerance( + contract, + judgment="gradient_accuracy", + op_class="reduction", + dtype="bfloat16", + ) + assert fwd.atol == 9.9 + assert grad.atol == 5.0e-2 + assert grad.atol != fwd.atol + + +def test_comparison_roles_by_report_kind(): + contract = load_contract() + expected = { + "forward_accuracy": ("bf16_candidate", "fp32_reference"), + "forward_invariance": ("transformed_config", "canonical_config"), + "train_infer_logprob_parity": ( + "training_style_teacher_forcing", + "inference_style_rollout_decode", + ), + "gradient_accuracy": ("bf16_candidate", "fp32_reference"), + "gradient_invariance": ("transformed_config", "canonical_config"), + } + for kind, (lhs, rhs) in expected.items(): + roles = resolve_comparison_roles(contract, kind) + assert roles.comparison_lhs_role == lhs + assert roles.comparison_rhs_role == rhs + assert_comparison_roles(contract, kind, lhs, rhs) + + +def test_forbidden_and_reversed_roles_hard_fail(): + contract = load_contract() + with pytest.raises(ContractResolveError, match="role mismatch"): + assert_comparison_roles( + contract, + "train_infer_logprob_parity", + "inference_style_rollout_decode", + "training_style_teacher_forcing", + ) + + +def test_aggregate_requires_declared_roles_and_direction(): + contract = load_contract() + values = torch.zeros(2) + with pytest.raises(ContractResolveError, match="role mismatch"): + compute_logprob_aggregates( + values, + values, + torch.ones(2, dtype=torch.bool), + contract=contract, + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="inference_style_rollout_decode", + comparison_rhs_role="training_style_teacher_forcing", + ) + with pytest.raises(ContractResolveError, match="forbidden"): + assert_comparison_roles( + contract, + "forward_accuracy", + "baseline", + "fp32_reference", + ) + with pytest.raises(ContractResolveError, match="forbidden"): + assert_comparison_roles( + contract, + "forward_invariance", + "singleton_aggregate", + "canonical_config", + ) + + +def test_resolve_tolerance_attaches_roles(): + contract = load_contract() + spec = resolve_tolerance( + contract, + judgment="forward_accuracy", + op_class="logprob", + dtype=torch.bfloat16, + ) + assert spec.comparison_lhs_role == "bf16_candidate" + assert spec.comparison_rhs_role == "fp32_reference" + assert "baseline" not in (spec.comparison_lhs_role, spec.comparison_rhs_role) + + +def test_chain_aggregate_named_resolve(): + contract = load_contract() + for metric in CHAIN_AGGREGATE_METRICS: + thr = resolve_chain_aggregate_thresholds(contract, metric, "bfloat16") + assert thr >= 0.0 + with pytest.raises(ContractResolveError, match="unknown chain aggregate"): + resolve_chain_aggregate_thresholds(contract, "mean_abs_dlogp", "bfloat16") + + +def test_compute_logprob_aggregates_formulas(): + # lhs - rhs = [0.0, 0.1, -0.2] + lhs = torch.tensor([1.0, 2.1, 0.8], dtype=torch.float32) + rhs = torch.tensor([1.0, 2.0, 1.0], dtype=torch.float32) + mask = torch.tensor([True, True, True]) + clip = (0.8, 1.2) + agg = compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=clip, + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + dlogp = torch.tensor([0.0, 0.1, -0.2]) + expected_max = float(dlogp.abs().max()) + expected_kl = float((torch.exp(dlogp) - 1.0 - dlogp).mean()) + ratio = torch.exp(dlogp) + expected_clip = float(((ratio < clip[0]) | (ratio > clip[1])).float().mean()) + assert math.isclose(agg.max_abs_dlogp, expected_max, rel_tol=0.0, abs_tol=1e-6) + assert math.isclose(agg.approx_kl0, expected_kl, rel_tol=0.0, abs_tol=1e-6) + assert math.isclose(agg.clipfrac0, expected_clip, rel_tol=0.0, abs_tol=1e-6) + assert agg.active_token_count == 3 + + +def test_active_mask_filters_tokens(): + lhs = torch.tensor([0.0, 10.0], dtype=torch.float32) + rhs = torch.tensor([0.0, 0.0], dtype=torch.float32) + mask = torch.tensor([True, False]) + agg = compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=default_clip_interval(load_contract()), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + assert agg.max_abs_dlogp == 0.0 + assert agg.active_token_count == 1 + + +def test_empty_active_set_hard_fails(): + lhs = torch.zeros(2) + rhs = torch.zeros(2) + mask = torch.tensor([False, False]) + with pytest.raises(ContractResolveError, match="empty active-token"): + compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + + +def test_nan_inf_hard_fail(): + lhs = torch.tensor([float("nan"), 0.0]) + rhs = torch.tensor([0.0, 0.0]) + mask = torch.tensor([True, True]) + with pytest.raises(ContractResolveError, match="NaN/Inf"): + compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + + lhs = torch.tensor([float("inf"), 0.0]) + with pytest.raises(ContractResolveError, match="NaN/Inf"): + compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + + +def test_judge_requires_all_three_aggregates(): + contract = load_contract() + clip = default_clip_interval(contract) + # Perfect match → all pass. + lhs = torch.zeros(4) + rhs = torch.zeros(4) + mask = torch.ones(4, dtype=torch.bool) + agg = compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=contract, + report_kind="train_infer_logprob_parity", + clip_interval=clip, + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + verdict = judge_logprob_aggregates(agg, contract, execution_dtype="bfloat16") + assert verdict.passed + assert {m.metric for m in verdict.metrics} == set(CHAIN_AGGREGATE_METRICS) + assert all(m.passed for m in verdict.metrics) + + # Large drift fails max_abs_dlogp / approx_kl0 / possibly clipfrac. + lhs = torch.tensor([0.0, 1.0]) + rhs = torch.zeros(2) + mask = torch.ones(2, dtype=torch.bool) + agg = compute_logprob_aggregates( + lhs, + rhs, + mask, + contract=contract, + report_kind="train_infer_logprob_parity", + clip_interval=clip, + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + verdict = judge_logprob_aggregates(agg, contract, execution_dtype="bfloat16") + assert not verdict.passed + assert any(not m.passed for m in verdict.metrics) + + +def test_compat_accuracy_mirrors_forward_accuracy(): + contract = load_contract() + for op_class in OP_CLASSES: + for dtype_name in ("float32", "bfloat16", "float16"): + acc = contract["accuracy"]["default"][op_class][dtype_name] + cell = contract["judgments"]["forward_accuracy"]["by_op_class"][op_class][dtype_name] + assert acc["atol"] == cell["atol"] + assert acc["rtol"] == cell["rtol"] + assert contract["batch_invariance"] == {"atol": 0.0, "rtol": 0.0} + + +def test_schema_rejects_nonzero_invariance_tolerance(): + contract = copy.deepcopy(load_contract()) + contract["judgments"]["forward_invariance"]["by_op_class"]["logprob"]["bfloat16"]["atol"] = 1e-3 + with pytest.raises(ContractSchemaError, match="bitwise"): + validate_contract_schema(contract) + + +def test_schema_rejects_baseline_role(): + contract = copy.deepcopy(load_contract()) + contract["comparison_roles"]["by_report_kind"]["forward_accuracy"][ + "comparison_lhs_role" + ] = "baseline" + with pytest.raises(ContractSchemaError, match="forbidden role"): + validate_contract_schema(contract) From 70286fb9f18247d197a700433a1dade50a0f2503 Mon Sep 17 00:00:00 2001 From: maxiaosong1124 Date: Wed, 12 Aug 2026 00:27:38 +0800 Subject: [PATCH 2/4] docs(ws1): add #267 C1 closeout evidence map Record acceptance-criteria mapping, verification commands, and residual scope so issue #267 can close without implying full #266 exit. --- docs/design/ws1-c1-267-closeout-evidence.md | 72 +++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 docs/design/ws1-c1-267-closeout-evidence.md diff --git a/docs/design/ws1-c1-267-closeout-evidence.md b/docs/design/ws1-c1-267-closeout-evidence.md new file mode 100644 index 00000000..4005b02d --- /dev/null +++ b/docs/design/ws1-c1-267-closeout-evidence.md @@ -0,0 +1,72 @@ +# #267 (C1) closeout evidence + +**Issue:** [#267](https://github.com/RL-Align/RL-Kernel/issues/267) +**Parent:** [#266](https://github.com/RL-Align/RL-Kernel/issues/266) (C1 only; does **not** close #266) +**Branch:** `feat/ws1-c1-tolerance-contract-267` +**Commit:** `af4d9c2` (and follow-ups on the same branch) + +## Acceptance criteria map + +| AC | Status | Where | +| --- | --- | --- | +| BF16 exec, FP32 ref/accum, FP8 out, TF32 policy documented + tested | Pass | `tolerance_contract.json` `policy`; `test_dtype_policy_*`; `docs/design/ws1-numerical-contract.md` | +| Independent execution/accum/output/reference + backend provenance | Pass | `resolve_dtype_policy`, `validate_backend_provenance`; tests | +| Missing applicable four-judgment cell → hard fail; explicit N/A only when declared | Pass | `resolve_tolerance` / `resolve_tolerance_support`; schema + unit tests | +| BF16+FP32 mandatory; FP16 optional complete; FP8 hard-fail | Pass | schema validation + resolve tests | +| Gradient tolerances independent of forward | Pass | separate `gradient_accuracy` rows; `op_checks` uses `gradient_accuracy`; independence test | +| Batch/Chunk inv rows bitwise `atol=0,rtol=0` | Pass | `forward_invariance` / `gradient_invariance`; schema rejects nonzero | +| Aggregate formulas, roles, active mask, clip, empty/NaN rules + boundary tests | Pass | `compute_logprob_aggregates` / `judge_logprob_aggregates`; tests | +| Reports persist `comparison_lhs_role` / `comparison_rhs_role`; reversed roles hard-fail | Pass | `OutputCheck` fields; `assert_comparison_roles` | +| No bare `baseline`; `singleton_aggregate` not a comparison role | Pass | `comparison_roles.forbidden` + schema tests | +| Named resolve for three aggregates; all three in logprob pass/fail | Pass | `resolve_chain_aggregate_thresholds`; `require_all` | +| Docs: three aggregates sole chain logprob metrics; grads independent | Pass | numerical contract + gtest usage guide | +| New gates obtain thresholds only via shared resolver | Pass for gtest path | `op_checks` / `check_operator`; residual private-atol inventory tracked in migration checklist (C3/C4/C8) | +| CUDA + Triton same contract rows; no backend-private relaxation | Pass | shared thresholds; `backend_private_tolerance_relaxation=false` | + +## Docking paths + +- `rl_engine/kernels/gtest/tolerance_contract.json` +- `rl_engine/kernels/gtest/tolerance.py` +- `rl_engine/kernels/gtest/op_checks.py` +- `tests/test_tolerance_contract.py` +- `tests/test_op_checks.py` +- `docs/design/ws1-numerical-contract.md` +- `docs/contributing/gtest-usage.md` +- `docs/design/ws1-gtest-migration-checklist.md` + +## Local verification + +```bash +python -m pytest tests/test_tolerance_contract.py tests/test_op_checks.py -q +# 41 passed +``` + +Sample resolve (logprob BF16): + +```text +forward_accuracy: mode=tolerance atol=0.05 rtol=0.0 lhs=bf16_candidate rhs=fp32_reference +forward_invariance: mode=bitwise atol=0.0 rtol=0.0 lhs=transformed_config rhs=canonical_config +gradient_accuracy: mode=tolerance (independent keys; does not read forward) +gradient_invariance: mode=bitwise atol=0.0 rtol=0.0 +``` + +## Explicitly out of this issue (still open under #266) + +- C2 full-model workload / manifest pin of clip interval (#268) +- C3/C4 shared invariance harnesses (#269/#270) +- Migrating every historical private-atol pytest (checklist; C8 evidence) +- Full-model train/infer gate and CI (#276/#277) +- Closing parent #266 + +## Suggested issue comment when PR is green + +```text +C1 complete on . + +- Contract + resolver + schema tests (41 passed locally) +- op_checks: forward_accuracy vs gradient_accuracy +- Docs: ws1-numerical-contract.md, gtest-usage.md, migration checklist +- Residual private-atol in legacy op tests tracked for C3/C4/C8; not a C1 dock gap + +Closing #267. Parent #266 remains open (C2–C11). +``` From 087156e5486be02e6104707000fad7cdfcc808c2 Mon Sep 17 00:00:00 2001 From: maxiaosong1124 Date: Wed, 12 Aug 2026 00:55:33 +0800 Subject: [PATCH 3/4] docs(ws1): streamline C1 gtest documentation --- docs/contributing/gtest-usage.md | 12 +- docs/contributing/testing.md | 3 +- docs/design/ws1-c1-267-closeout-evidence.md | 72 ------ docs/design/ws1-gtest-migration-checklist.md | 248 ------------------- docs/design/ws1-numerical-contract.md | 187 -------------- 5 files changed, 8 insertions(+), 514 deletions(-) delete mode 100644 docs/design/ws1-c1-267-closeout-evidence.md delete mode 100644 docs/design/ws1-gtest-migration-checklist.md delete mode 100644 docs/design/ws1-numerical-contract.md diff --git a/docs/contributing/gtest-usage.md b/docs/contributing/gtest-usage.md index 036c1b65..40438c77 100644 --- a/docs/contributing/gtest-usage.md +++ b/docs/contributing/gtest-usage.md @@ -3,7 +3,6 @@ > **Audience:** contributors implementing train–inference / batch-invariant operators > **Entry point:** `scripts/check_operator.py` + `rl_engine/kernels/gtest/*` > **Numerical SSOT:** [#267](https://github.com/RL-Align/RL-Kernel/issues/267) four-judgment contract -> **Related:** [WS1 numerical contract](../design/ws1-numerical-contract.md) · [migration checklist](../design/ws1-gtest-migration-checklist.md) This is the official how-to for the gtest harness: register an op, build inputs, run the CLI for forward/backward checks, and obtain tolerances from the shared contract (not private `atol`/`rtol`). @@ -270,6 +269,7 @@ verdict = judge_logprob_aggregates(agg, contract, execution_dtype="bfloat16") | Reference / accumulation | FP32 | | FP8 | Out of scope (resolve hard-fails) | | TF32 | Disabled | +| Profiles | `cuda_bf16` and `triton_cuda_bf16` share **the same** thresholds | WS1 evidence must attach checked provenance to its candidate report: @@ -297,9 +297,13 @@ candidate = CandidateSpec( The suite rejects backend fallback, dtype drift, TF32 enablement, and observed output dtypes that disagree with this provenance before producing a passing report. -| Profiles | `cuda_bf16` and `triton_cuda_bf16` share **the same** thresholds | -**Do not** use private `atol=1e-5` (etc.) as WS1 gate evidence. Inventory of legacy private thresholds: [migration checklist](../design/ws1-gtest-migration-checklist.md). +`check_operator.py` is a local debugging CLI and does not construct provenance on its +own. A WS1 gate must create `CandidateSpec(..., provenance=provenance)` in its harness; +use `--json` to retain the resolved judgment and comparison-role fields in CLI reports. + +**Do not** use private `atol=1e-5` (etc.) as WS1 gate evidence. Migrate a gate to +the shared resolver before using it as WS1 evidence. ### 6.4 Report `tol=(atol=..., rtol=...)` @@ -375,8 +379,6 @@ New pytest code should call `resolve_tolerance` instead of copying magic numbers | Doc | Content | |-----|---------| -| [ws1-numerical-contract.md](../design/ws1-numerical-contract.md) | Four judgments, roles, aggregate formulas | -| [ws1-gtest-migration-checklist.md](../design/ws1-gtest-migration-checklist.md) | Which tests still use private thresholds and when to migrate | | [testing.md](testing.md) | Short testing entry points | | Issues [#266](https://github.com/RL-Align/RL-Kernel/issues/266) / [#267](https://github.com/RL-Align/RL-Kernel/issues/267) | WS1 closeout and C1 contract | diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index b0924ba0..7d8aff90 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -14,8 +14,7 @@ Full usage (register `OP_SPECS`, build inputs, CLI flags, and the WS1 four-judgm tolerance contract after #267): - **[gtest usage guide](gtest-usage.md)** (operator CLI + `OP_SPECS` + contract; English) -- [WS1 numerical contract](../design/ws1-numerical-contract.md) -- [gtest private-threshold migration checklist](../design/ws1-gtest-migration-checklist.md) +- **[gtest 使用指南](gtest-usage.zh-CN.md)**（算子 CLI、`OP_SPECS` 与数值合同；中文） ## Dispatch Tests diff --git a/docs/design/ws1-c1-267-closeout-evidence.md b/docs/design/ws1-c1-267-closeout-evidence.md deleted file mode 100644 index 4005b02d..00000000 --- a/docs/design/ws1-c1-267-closeout-evidence.md +++ /dev/null @@ -1,72 +0,0 @@ -# #267 (C1) closeout evidence - -**Issue:** [#267](https://github.com/RL-Align/RL-Kernel/issues/267) -**Parent:** [#266](https://github.com/RL-Align/RL-Kernel/issues/266) (C1 only; does **not** close #266) -**Branch:** `feat/ws1-c1-tolerance-contract-267` -**Commit:** `af4d9c2` (and follow-ups on the same branch) - -## Acceptance criteria map - -| AC | Status | Where | -| --- | --- | --- | -| BF16 exec, FP32 ref/accum, FP8 out, TF32 policy documented + tested | Pass | `tolerance_contract.json` `policy`; `test_dtype_policy_*`; `docs/design/ws1-numerical-contract.md` | -| Independent execution/accum/output/reference + backend provenance | Pass | `resolve_dtype_policy`, `validate_backend_provenance`; tests | -| Missing applicable four-judgment cell → hard fail; explicit N/A only when declared | Pass | `resolve_tolerance` / `resolve_tolerance_support`; schema + unit tests | -| BF16+FP32 mandatory; FP16 optional complete; FP8 hard-fail | Pass | schema validation + resolve tests | -| Gradient tolerances independent of forward | Pass | separate `gradient_accuracy` rows; `op_checks` uses `gradient_accuracy`; independence test | -| Batch/Chunk inv rows bitwise `atol=0,rtol=0` | Pass | `forward_invariance` / `gradient_invariance`; schema rejects nonzero | -| Aggregate formulas, roles, active mask, clip, empty/NaN rules + boundary tests | Pass | `compute_logprob_aggregates` / `judge_logprob_aggregates`; tests | -| Reports persist `comparison_lhs_role` / `comparison_rhs_role`; reversed roles hard-fail | Pass | `OutputCheck` fields; `assert_comparison_roles` | -| No bare `baseline`; `singleton_aggregate` not a comparison role | Pass | `comparison_roles.forbidden` + schema tests | -| Named resolve for three aggregates; all three in logprob pass/fail | Pass | `resolve_chain_aggregate_thresholds`; `require_all` | -| Docs: three aggregates sole chain logprob metrics; grads independent | Pass | numerical contract + gtest usage guide | -| New gates obtain thresholds only via shared resolver | Pass for gtest path | `op_checks` / `check_operator`; residual private-atol inventory tracked in migration checklist (C3/C4/C8) | -| CUDA + Triton same contract rows; no backend-private relaxation | Pass | shared thresholds; `backend_private_tolerance_relaxation=false` | - -## Docking paths - -- `rl_engine/kernels/gtest/tolerance_contract.json` -- `rl_engine/kernels/gtest/tolerance.py` -- `rl_engine/kernels/gtest/op_checks.py` -- `tests/test_tolerance_contract.py` -- `tests/test_op_checks.py` -- `docs/design/ws1-numerical-contract.md` -- `docs/contributing/gtest-usage.md` -- `docs/design/ws1-gtest-migration-checklist.md` - -## Local verification - -```bash -python -m pytest tests/test_tolerance_contract.py tests/test_op_checks.py -q -# 41 passed -``` - -Sample resolve (logprob BF16): - -```text -forward_accuracy: mode=tolerance atol=0.05 rtol=0.0 lhs=bf16_candidate rhs=fp32_reference -forward_invariance: mode=bitwise atol=0.0 rtol=0.0 lhs=transformed_config rhs=canonical_config -gradient_accuracy: mode=tolerance (independent keys; does not read forward) -gradient_invariance: mode=bitwise atol=0.0 rtol=0.0 -``` - -## Explicitly out of this issue (still open under #266) - -- C2 full-model workload / manifest pin of clip interval (#268) -- C3/C4 shared invariance harnesses (#269/#270) -- Migrating every historical private-atol pytest (checklist; C8 evidence) -- Full-model train/infer gate and CI (#276/#277) -- Closing parent #266 - -## Suggested issue comment when PR is green - -```text -C1 complete on . - -- Contract + resolver + schema tests (41 passed locally) -- op_checks: forward_accuracy vs gradient_accuracy -- Docs: ws1-numerical-contract.md, gtest-usage.md, migration checklist -- Residual private-atol in legacy op tests tracked for C3/C4/C8; not a C1 dock gap - -Closing #267. Parent #266 remains open (C2–C11). -``` diff --git a/docs/design/ws1-gtest-migration-checklist.md b/docs/design/ws1-gtest-migration-checklist.md deleted file mode 100644 index 3b0a9346..00000000 --- a/docs/design/ws1-gtest-migration-checklist.md +++ /dev/null @@ -1,248 +0,0 @@ -# WS1 gtest 阈值迁移清单 - -> **关联：** [#266](https://github.com/RL-Align/RL-Kernel/issues/266) 父收尾 · [#267](https://github.com/RL-Align/RL-Kernel/issues/267) C1 契约 · [数值契约说明](ws1-numerical-contract.md) -> **目的：** 盘点「哪些测试仍用私有 `atol`/`rtol`、哪些已走 SSOT、何时必须迁到 `resolve_tolerance`」。 -> **快照：** 基于 `feat/ws1-c1-tolerance-contract-267` 落地 C1 后的仓库状态；文件增减时请更新本表。 - ---- - -## 0. 迁移总原则 - -### 0.1 SSOT 入口（改后唯一推荐） - -```python -from rl_engine.kernels.gtest.tolerance import ( - load_contract, - resolve_tolerance, - compute_logprob_aggregates, - judge_logprob_aggregates, - default_clip_interval, -) - -contract = load_contract() -spec = resolve_tolerance( - contract, - judgment="forward_accuracy", # 或 forward_invariance / gradient_* - op_class="logprob", # elementwise | reduction | logprob | attention - dtype="bfloat16", - backend_profile="cuda_bf16", # 与 triton_cuda_bf16 同阈值 -) -# assert_close(..., atol=spec.atol, rtol=spec.rtol) -# 不变性：spec.mode == "bitwise" 且 atol=rtol=0 → 优先 torch.equal -``` - -| Judgment | 用于 | -|----------|------| -| `forward_accuracy` | BF16 candidate vs FP32 reference | -| `forward_invariance` | 同逻辑 workload 跨 batch/chunk/layout（**bitwise**） | -| `gradient_accuracy` | 梯度 vs FP32 参考（**不得**读 forward 行） | -| `gradient_invariance` | 梯度跨 config（**bitwise**） | -| 三聚合 API | 链级 / 训推 selected-logprob（`max_abs_dlogp` / `approx_kl0` / `clipfrac0`） | - -### 0.2 什么叫「私有阈值」（禁止作为 WS1 gate 证据） - -- 测试文件内字面量：`atol=1e-5`、`atol=5e-2`、模块常量 `_DECODE_ATOL` 等 -- 文档里写死但未从 `tolerance_contract.json` resolve 的数 -- 从 `contract["accuracy"]...` 手抄数值后本地再改（应用 resolve，不要复制常量） -- 用非零 `atol` 充当 Batch/Chunk **invariance** 通过条件 - -### 0.3 什么可以保留（不必硬迁） - -| 场景 | 处理 | -|------|------| -| **bitwise 身份断言**（`torch.equal`） | 合法；对应 invariance judgment 的 `mode=bitwise` | -| **非数值语义**（mask 形状、版本单调、manifest 字段） | 不迁 | -| **框架/集成单测**（bridge、vLLM mock、DeepSpeed worker 编排） | 非 WS1 op gate；可保留宽松 `allclose`，但**不能**当作 #266 EXIT 证据 | -| **生产 FA / SDPA 对齐**（`test_attention_correctness`） | 非 BI 候选路径；阈值可独立，**不得**写进 WS1 EXIT claim | -| **legacy `accuracy` 键** | 仅兼容；新代码禁止新增依赖，应改 `resolve_tolerance` | - -### 0.4 建议迁移句式 - -```python -# BAD — 私有阈值 -torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) - -# GOOD — accuracy -spec = resolve_tolerance(contract, judgment="forward_accuracy", op_class="reduction", dtype=dtype) -torch.testing.assert_close(out, ref, atol=spec.atol, rtol=spec.rtol) - -# GOOD — invariance -ispec = resolve_tolerance(contract, judgment="forward_invariance", op_class="attention", dtype=dtype) -assert ispec.mode == "bitwise" and ispec.atol == 0.0 -assert torch.equal(a, b) # 或 assert_close(..., atol=0, rtol=0) - -# GOOD — gradient accuracy（独立 judgment） -gspec = resolve_tolerance(contract, judgment="gradient_accuracy", op_class="logprob", dtype=dtype) -``` - ---- - -## 1. 状态总表（`tests/`） - -图例： - -| 标记 | 含义 | -|------|------| -| **A** | 已走 resolver / gtest suite（目标态） | -| **B** | 走 `load_contract` 旧键或 gtest 间接路径（过渡） | -| **C** | WS1 相关但 **私有 atol**（应迁） | -| **D** | 多为 `torch.equal` / 结构断言（ok 或仅需声明 judgment） | -| **E** | 非 WS1 门禁（框架/产品路径，低优先级） | - -### 1.1 已对齐或接近 SSOT - -| 文件 | 状态 | 说明 | 下一步 | -|------|------|------|--------| -| `test_tolerance_contract.py` | **A** | C1 schema + resolve + 聚合 | 保持；契约变更必跑 | -| `test_op_checks.py` | **A/B** | suite 已按 judgment 解析；部分用例注入最小 contract | 新 fixture 尽量带 `judgments` | -| `test_operator_inputs.py` | **B** | 输入/规格，无数值阈值主责 | 无需迁阈值 | -| `test_swiglu.py` | **B/D** | issue-108 harness + 大量 `torch.equal` | accuracy 路径确认走 suite；字面量 atol 清零 | -| `test_det_gemm.py` | **B** | `load_contract()["accuracy"]...` | **优先迁**：改为 `resolve_tolerance(..., forward/gradient_accuracy)` | -| `test_deterministic_attention_cuda.py` | **B/C** | 部分用 suite；仍见 `5e-2/2e-2` 字面量 | 字面量改为 resolve；invariance 保持 equal | - -### 1.2 WS1 算子测试 — 私有阈值（应迁，按优先级） - -| 优先级 | 文件 | op_class 建议 | 现状摘要 | 何时必须迁 | -|--------|------|---------------|----------|------------| -| **P0** | `test_batch_invariant_logp.py` | `logprob` | 大量 `1e-6`…`1e-2` 私有；含 bwd | 接 C3/C4/C8 证据前 | -| **P0** | `test_linear_logp.py` | `logprob` | `1e-5`…`1.5e-1` 混用；bf16 松阈值 | 同上；链级改用三聚合 API | -| **P0** | `test_logp.py` / `test_deterministic_logp.py` | `logprob` | 私有 atol | 关 #148 residual / C8 前 | -| **P0** | `test_rms_norm.py` | `reduction` | `1e-5`…`8e-2`；bwd 混用 | C8 RMSNorm 证据前 | -| **P0** | `test_triton_batch_invariant_attention.py` | `attention` | 混 `1e-5` 与 `5e-2/2e-2` | C8 Attention 证据前 | -| **P0** | `test_attention.py` | `attention` | native GT；`1e-4`/`2e-6` 等 | 与 contract `attention` 行对齐 | -| **P1** | `test_kv_cache_attention.py` | `attention` | 含 `2e-6` 等；#152 相关 | **C6/C7 前必须**消私有 decode 阈值 | -| **P1** | `test_issue151_embedding_lm_head_invariance.py` | emb + lm_head + logp | bf16 `5e-2` 手写 | C8 emb/lm_head 证据前 | -| **P1** | `test_lm_head.py` | `reduction` | 多 equal；grad `1e-5` 私有 | 迁 grad → `gradient_accuracy` | -| **P1** | `test_embedding.py` | `elementwise` | 多为 equal | 若有 tolerance 路径再 resolve | -| **P1** | `test_rope.py` | `elementwise` | `1e-3`…`2e-2` | C5 RoPE 证据前 | -| **P1** | `test_matmul.py` | `reduction` | 私有 `1e-4/1e-5` | 与 det_gemm 统一 | -| **P2** | `test_pack.py` | `elementwise` | 几乎 equal；gradcheck `1e-6` | packing 纳入 #150 时 | -| **P2** | `test_grpo_loss.py` / `test_ratio_kl.py` | （loss，契约暂无独立 class） | `1e-4` 等 | 若进 chain 则扩展 op_class 或显式 N/A | -| **P3** | `test_attention_correctness.py` | 非 BI EXIT | FA/SDPA 私有表 | **不迁入 WS1 SSOT**；文档标明 out of WS1 claim | -| **P3** | `test_op_accuracy.py` | 杂项 harness | `1e-3` | 废弃或改走 `check_operator` + contract | - -### 1.3 非 WS1 门禁（低优先级 / 不阻塞 #267） - -| 文件 | 状态 | 说明 | -|------|------|------| -| `test_deepspeed_training_worker.py` | **E** | 训练 worker；`atol=1e-5` 编排级 | -| `test_stateless_training_contract.py` | **E** | 契约字段/数值 smoke | -| `test_rl_kernel_loss_step.py` | **E** | 端到端 loss 步 | -| `test_sampler_temperature.py` | **E** | 采样 | -| `test_weight_sync_bridge.py` 等 | **D/E** | bridge / IPC | -| `test_vllm_rollout_sampler.py` | **D/E** | vLLM mock | -| `test_alignment_model_wrappers.py` | **D/E** | wrapper 行为 | -| `test_rl_batch_fixture.py` | **D** | fixture 身份 | -| `test_stateless_executor.py` / `*_hf_integration*` | **D/E** | 执行器集成 | - -这些**不**作为 #266 Full WS1 EXIT 的数值证据来源；C10/C11 不得引用其私有阈值刷绿。 - ---- - -## 2. 按 #266 子 issue 的「何时必须迁」 - -| 子 issue | 阻塞迁移范围 | 完成信号 | -|----------|--------------|----------| -| **C1 #267** | 契约 + resolver + schema/报告测试 | **实现完成；待 CI / issue evidence** | -| **C3 #269** forward harness | 所有 **forward_accuracy / forward_invariance** 的 op 单测证据路径 | 无 private forward atol 作为 gate | -| **C4 #270** grad harness | 所有 **gradient_*** 证据路径 | 无「grad 抄 forward 字面量」 | -| **C5 #271** RoPE/elementwise | `test_rope.py`、activation/swiglu residual | audit 报告阈值均来自 resolve | -| **C6/C7 #272/#273** KV | `test_kv_cache_attention.py` 及后续 kv harness | **禁止** `_DECODE_ATOL` 类私有常量 | -| **C8 #274** closed-op 矩阵 | rmsnorm / gemm / attn / logp / emb / lm_head 测试 | 每格 `requested/actual backend` + resolve 阈值 | -| **C10 #276** 全模型 gate | 仅用 resolver + 三聚合；禁止任何测试内字面量阈值 | gate 报告无 private tol 字段 | -| **C11 #277** CI | CI 只跑 resolve 路径 | fail-closed | - -**规则：** 某 op 的 PR 若声称「满足 #266/C8」，则该 PR 触达的 assert **必须**来自 `resolve_tolerance` / 聚合 API，而不是文件顶部的魔法数。 - ---- - -## 3. 文件级迁移清单（可勾选） - -### 3.1 P0 — 直接挡 C3/C4/C8 - -- [ ] `tests/test_batch_invariant_logp.py` — fwd/bwd accuracy + invariance 拆 judgment -- [ ] `tests/test_linear_logp.py` — 同上；训推/链级改用三聚合 -- [ ] `tests/test_logp.py` -- [ ] `tests/test_deterministic_logp.py` -- [ ] `tests/test_rms_norm.py` -- [ ] `tests/test_triton_batch_invariant_attention.py` -- [ ] `tests/test_attention.py` -- [ ] `tests/test_det_gemm.py` — 去掉 `contract["accuracy"]` 直读 - -### 3.2 P1 — C5/C6/C7/C8 residual - -- [ ] `tests/test_kv_cache_attention.py` -- [ ] `tests/test_issue151_embedding_lm_head_invariance.py` -- [ ] `tests/test_lm_head.py` -- [ ] `tests/test_embedding.py`（若有 non-bitwise 路径） -- [ ] `tests/test_rope.py` -- [ ] `tests/test_matmul.py` -- [ ] `tests/test_deterministic_attention_cuda.py` 中剩余字面量 -- [ ] `tests/test_swiglu.py` 中任何 residual 字面量 - -### 3.3 P2 — 进 chain 时 - -- [ ] `tests/test_pack.py` -- [ ] `tests/test_grpo_loss.py` / `tests/test_ratio_kl.py`（先扩 contract op_class 或标 N/A） -- [ ] `scripts/check_operator.py` 报告字段确认只回传 resolve 结果（已间接） - -### 3.4 明确不迁入 WS1 SSOT - -- [x] `tests/test_attention_correctness.py` — 生产 FA；文档标注非 EXIT -- [x] bridge / vLLM / DeepSpeed / sampler 类 **E** 组 - ---- - -## 4. 推荐落地动作（每个测试文件） - -1. **分类每条 assert** - - identity / batch-invariance → `forward_invariance` 或 `gradient_invariance` + `torch.equal` - - vs fp32 gold → `*_accuracy` - - train vs infer logp → 三聚合，不用单点 atol 冒充 -2. **删除模块级 `_ATOL` / `_RTOL`** -3. **dtype 参数化** 时用 `resolve_tolerance(..., dtype=dtype)`，禁止 bf16 写死 `5e-2` -4. **报告**（若有）写上 `comparison_lhs_role` / `comparison_rhs_role`（从 spec 取） -5. **禁止** 为让 invariance 通过而调大 atol - -### 4.1 与契约行不一致时怎么办 - -| 情况 | 动作 | -|------|------| -| 测试私有更松，契约更紧 → 测试红 | **修 kernel** 或开 Blocker；**禁止**在测试放宽 | -| 测试私有更紧，契约更松 | 迁到契约后可能变绿；可保留额外严格 assert 但须标注 *non-gate* | -| 需要新 op_class（如 `grpo_loss`） | 先改 `tolerance_contract.json` + schema 测试，再迁测试 | -| decode vs prefill 无法 bitwise | 用 contract 已声明的 semantic 行 / 三聚合；**不要**私设 `_DECODE_ATOL` | - ---- - -## 5. 工具与 CI 建议（后续，非 C1 范围） - -| 建议 | 作用 | -|------|------| -| 简单 lint：`tests/**/*.py` 禁止 `atol=\d`（allowlist 契约测试与 FA 测试） | 防回流 | -| `pytest` marker：`ws1_gate` 仅收集 resolve 路径 | C11 门禁清晰 | -| 在 `check_operator.py` 输出中强制打印 `judgment` + roles | 证据可检索 | - -C1 **不**强制上 lint；C8/C10 前建议至少做 allowlist 扫描。 - ---- - -## 6. 现状一句话 - -| 层 | 状态 | -|----|------| -| **契约 + resolver（gtest 核心）** | 已就绪；待 CI / issue evidence（#267） | -| **op_checks 接入** | 已按 judgment 分叉，并持久化 roles / provenance | -| **存量 op 单测** | **多数仍私有 atol**（上表 P0/P1） | -| **#266 EXIT** | 依赖后续把 P0/P1 迁完，而不是只合 C1 | - -**C1 的价值是「唯一入口已存在」；清单的价值是「知道还欠哪些文件」。** -未完成 P0/P1 迁移前，**不得**声称「全仓测试已统一走 WS1 数值契约」。 - ---- - -## 7. 修订记录 - -| 日期 | 说明 | -|------|------| -| 2026-08-11 | 初版：C1 落地后基于 `tests/` 扫描的迁移清单与优先级 | diff --git a/docs/design/ws1-numerical-contract.md b/docs/design/ws1-numerical-contract.md deleted file mode 100644 index 76410fcb..00000000 --- a/docs/design/ws1-numerical-contract.md +++ /dev/null @@ -1,187 +0,0 @@ -# WS1 Numerical Contract (C1 / #267) - -> **Parent:** [#266](https://github.com/RL-Align/RL-Kernel/issues/266) -> **Issue:** [#267](https://github.com/RL-Align/RL-Kernel/issues/267) -> **SSOT files:** `rl_engine/kernels/gtest/tolerance_contract.json`, `rl_engine/kernels/gtest/tolerance.py` - -This document freezes the **sole** numerical judgment source for WS1 ablations and -gates. New gates must obtain thresholds only through the shared resolver APIs; private -`atol` / `rtol` constants are forbidden. - -## Scope boundary - -**Allowed WS1 claim after full exit (#266):** single-GPU model-level train–inference -consistency for full Qwen3-8B Dense under required CUDA BF16 and Triton-on-CUDA BF16 -profiles (in-repo BI stack). - -**Not claimed here:** multi-GPU (WS2), real vLLM vs Megatron / vime product alignment -(WS3), or kernel bug fixes (open Blockers). - -## Dtype policy - -| Field | Lock | -| --- | --- | -| `execution_dtype` | **BF16** (mandatory) | -| `accumulation_dtype` | **FP32** | -| `reference_dtype` | **FP32** | -| `output_dtype.default` | follows execution | -| logprob aggregates compute dtype | **FP32** | -| FP8 | **out of scope** (request → hard fail) | -| FP16 | optional; rows complete when declared | -| TF32 (reference + candidate) | **disabled** (repo-wide single policy) | -| Backend profiles | `cuda_bf16`, `triton_cuda_bf16` (same thresholds) | -| Backend-private tolerance relaxation | **forbidden** | - -Execution, accumulation, output, and reference dtypes resolve **independently** via -`resolve_dtype_policy()`. - -## Four judgments - -| Judgment | What it compares | Default mode | -| --- | --- | --- | -| `forward_accuracy` | BF16 candidate vs FP32 reference outputs | tolerance | -| `forward_invariance` | transformed vs canonical config, same backend/dtype/logical workload | **bitwise** (`atol=0`, `rtol=0`) | -| `gradient_accuracy` | candidate gradient/VJP vs FP32 reference gradient/VJP | tolerance (independent of forward) | -| `gradient_invariance` | transformed vs canonical gradients, same logical workload | **bitwise** (`atol=0`, `rtol=0`) | - -Every declared-applicable `(judgment, op_class, dtype)` tuple must resolve. Missing -applicable cells hard-fail. Explicit `not_applicable` / `out_of_scope` is allowed only -when present in the schema. Use `resolve_tolerance_support()` to persist the explicit -support status; requesting thresholds for an N/A or out-of-scope cell still hard-fails. - -**Batch/Chunk invariance** (issue #150 / C10 matrix) **must** use the invariance -judgments in bitwise mode. Nonzero tolerance cannot satisfy that gate. - -Op classes: `elementwise`, `reduction`, `logprob`, `attention`. - -## Comparison roles - -Reports must record `comparison_lhs_role` / `comparison_rhs_role`. A bare `baseline` -field is forbidden. C2 `singleton_aggregate` is an **execution/aggregation mode**, not -a comparison role. - -| Report kind | `comparison_lhs_role` | `comparison_rhs_role` | -| --- | --- | --- | -| `forward_accuracy` | `bf16_candidate` | `fp32_reference` | -| `forward_invariance` | `transformed_config` | `canonical_config` | -| `train_infer_logprob_parity` | `training_style_teacher_forcing` | `inference_style_rollout_decode` | -| `gradient_accuracy` | `bf16_candidate` | `fp32_reference` | -| `gradient_invariance` | `transformed_config` | `canonical_config` | - -Direction is locked so train/infer preserves: - -```text -dlogp = train_logp - rollout_logp -ratio0 = exp(dlogp) -``` - -Swapping lhs/rhs without a different declared contract row hard-fails. - -API: `resolve_comparison_roles()`, `assert_comparison_roles()`. - -Aggregate callers must also provide `contract`, `report_kind`, and both roles; the -compute API validates the direction before calculating any metric. gtest reports -persist these roles on every output verdict. Backend reports must include -`BackendProvenance` (requested/actual backend, all four dtypes, and TF32 state), -which `validate_backend_provenance()` checks against the selected profile. - -## Chain-level logprob aggregates - -These three metrics are the **only** chain-level logprob / ablation aggregates for WS1 -pass/fail. Gradients use independent `gradient_*` tensor verdicts and **do not** use -these aggregates. - -Computed in FP32 on **active selected tokens only**: - -```text -dlogp = comparison_lhs_logp - comparison_rhs_logp -max_abs_dlogp = max(abs(dlogp)) -approx_kl0 = mean(exp(dlogp) - 1 - dlogp) -clipfrac0 = mean(1[exp(dlogp) outside clip_interval]) -``` - -Rules: - -- **All three** must pass (`require_all=true`). -- Empty active-token set → hard fail. -- NaN / Inf in `dlogp`, `ratio0`, or aggregates → hard fail. -- Clip interval is pinned by the workload manifest (C2); the contract stores a default - and the field name `clip_interval`. - -API: `compute_logprob_aggregates()`, `judge_logprob_aggregates()`, -`resolve_chain_aggregate_thresholds()`. - -## Resolver usage - -```python -from rl_engine.kernels.gtest.tolerance import ( - load_contract, - resolve_tolerance, - resolve_dtype_policy, - compute_logprob_aggregates, - judge_logprob_aggregates, - default_clip_interval, -) - -contract = load_contract() -policy = resolve_dtype_policy(contract) - -fwd = resolve_tolerance( - contract, - judgment="forward_accuracy", - op_class="logprob", - dtype="bfloat16", - backend_profile="cuda_bf16", -) -bwd = resolve_tolerance( - contract, - judgment="gradient_accuracy", - op_class="logprob", - dtype="bfloat16", - backend_profile="triton_cuda_bf16", # same thresholds as cuda_bf16 -) -inv = resolve_tolerance( - contract, - judgment="forward_invariance", - op_class="attention", - dtype="bfloat16", -) -# inv.mode == "bitwise", inv.atol == 0.0, inv.rtol == 0.0 -``` - -`op_checks.run_operator_suite` resolves **forward_accuracy** for outputs and -**gradient_accuracy** for gradients. - -## Compatibility keys - -For older tests that still dig into: - -- `contract["accuracy"]["default"][op_class][dtype]` — mirror of `forward_accuracy` -- `contract["batch_invariance"]` — `{atol: 0, rtol: 0}` - -New code should call the resolvers above. Schema validation fails if the compatibility -mirror drifts from `forward_accuracy` or if invariance rows leave bitwise mode. - -## Related issues - -| ID | Role | -| --- | --- | -| #266 | WS1 closeout parent | -| #267 | This contract (C1) | -| #268 | Full-model workload / clip interval pin in manifest | -| #269 / #270 | Forward / gradient invariance harnesses | -| #276 | Full-model train/infer gate consuming this contract | -| #154 / #108 | Historical contract owners (superseded remaining work → C1) | - -## Migration of existing tests - -Most operator tests still use **private** `atol` / `rtol` literals. That is expected -after C1: the SSOT exists, but call sites have not all moved. - -See the full inventory, priority, and “when it must migrate” map: - -- [WS1 gtest 阈值迁移清单](ws1-gtest-migration-checklist.md) - -How to register ops and run the CLI (post-#267): - -- [gtest usage guide](../contributing/gtest-usage.md) From 81ddd652aa9d16d4c9b52925fed7ace36f3a4606 Mon Sep 17 00:00:00 2001 From: maxiaosong1124 Date: Wed, 12 Aug 2026 01:43:15 +0800 Subject: [PATCH 4/4] fix(ws1): address tolerance contract review --- docs/contributing/testing.md | 1 - rl_engine/kernels/gtest/__init__.py | 8 ++ rl_engine/kernels/gtest/op_checks.py | 2 +- rl_engine/kernels/gtest/tolerance.py | 36 ++++++-- tests/test_op_checks.py | 2 +- tests/test_tolerance_contract.py | 123 ++++++++++++++++++++++----- 6 files changed, 137 insertions(+), 35 deletions(-) diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index 7d8aff90..749b203e 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -14,7 +14,6 @@ Full usage (register `OP_SPECS`, build inputs, CLI flags, and the WS1 four-judgm tolerance contract after #267): - **[gtest usage guide](gtest-usage.md)** (operator CLI + `OP_SPECS` + contract; English) -- **[gtest 使用指南](gtest-usage.zh-CN.md)**（算子 CLI、`OP_SPECS` 与数值合同；中文） ## Dispatch Tests diff --git a/rl_engine/kernels/gtest/__init__.py b/rl_engine/kernels/gtest/__init__.py index 61b43218..a12db99e 100644 --- a/rl_engine/kernels/gtest/__init__.py +++ b/rl_engine/kernels/gtest/__init__.py @@ -4,8 +4,12 @@ from .op_checks import CandidateSpec, OperatorCase, run_operator_suite from .tolerance import ( BackendProvenance, + ContractError, ContractResolveError, ContractSchemaError, + load_contract, + resolve_dtype_policy, + resolve_tolerance, resolve_tolerance_support, validate_backend_provenance, ) @@ -15,8 +19,12 @@ "OperatorCase", "run_operator_suite", "BackendProvenance", + "ContractError", "ContractResolveError", "ContractSchemaError", + "load_contract", + "resolve_tolerance", + "resolve_dtype_policy", "resolve_tolerance_support", "validate_backend_provenance", ] diff --git a/rl_engine/kernels/gtest/op_checks.py b/rl_engine/kernels/gtest/op_checks.py index 085b5f89..1bada7b3 100644 --- a/rl_engine/kernels/gtest/op_checks.py +++ b/rl_engine/kernels/gtest/op_checks.py @@ -155,7 +155,7 @@ def _run_candidate( validate_backend_provenance(contract, candidate.provenance) if candidate.backend != candidate.provenance.actual_backend: raise ContractResolveError( - f"candidate backend {candidate.backend!r} disagrees with reported actual backend " + f"candidate backend {candidate.backend!r} disagrees with reported actual_backend " f"{candidate.provenance.actual_backend!r}" ) for case in cases: diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index 230dd95d..0d2ae2e7 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -381,7 +381,7 @@ def resolve_tolerance( "backend_private_tolerance_relaxation must remain false under WS1 C1" ) - if dtype_name in OUT_OF_SCOPE_DTYPES or policy.fp8 == "out_of_scope" and dtype_name == "float8": + if dtype_name in OUT_OF_SCOPE_DTYPES: raise ContractResolveError( f"dtype {dtype_name!r} is out of scope for WS1 (FP8 requests hard-fail)" ) @@ -781,12 +781,10 @@ def _validate_judgments(judgments: Mapping[str, Any]) -> None: "applicable", "not_applicable", }: - # BF16/FP32 must be explicitly applicable (or explicit N/A). - if status != "applicable": - raise ContractSchemaError( - f"mandatory dtype cell must be applicable: " - f"{judgment}/{op_class}/{dtype_name} status={status!r}" - ) + raise ContractSchemaError( + f"mandatory dtype cell must be applicable: " + f"{judgment}/{op_class}/{dtype_name} status={status!r}" + ) if status in {"applicable", "optional"}: for thr in ("atol", "rtol", "mode"): if thr not in cell: @@ -810,6 +808,8 @@ def _validate_chain_aggregates(root: Mapping[str, Any]) -> None: "require_all", "nan_inf_policy", "empty_active_token_set", + "active_token_policy", + "clip_interval_field", "dlogp_definition", "default_clip_interval", "sole_chain_level_logprob_metrics", @@ -823,6 +823,12 @@ def _validate_chain_aggregates(root: Mapping[str, Any]) -> None: raise ContractSchemaError("nan_inf_policy must be hard_fail") if root["empty_active_token_set"] != "hard_fail": raise ContractSchemaError("empty_active_token_set must be hard_fail") + if root["active_token_policy"] != "active selected tokens only": + raise ContractSchemaError("active_token_policy must be 'active selected tokens only'") + if root["clip_interval_field"] != "clip_interval": + raise ContractSchemaError("clip_interval_field must be 'clip_interval'") + if root["dlogp_definition"] != "comparison_lhs_logp - comparison_rhs_logp": + raise ContractSchemaError("dlogp_definition does not match implementation") if not root["require_all"]: raise ContractSchemaError("require_all must be true for chain logprob aggregates") sole = list(root["sole_chain_level_logprob_metrics"]) @@ -843,6 +849,15 @@ def _validate_chain_aggregates(root: Mapping[str, Any]) -> None: for dtype_name in ("bfloat16", "float32"): if dtype_name not in by_dtype or "threshold" not in by_dtype[dtype_name]: raise ContractSchemaError(f"metric {name!r} missing threshold for {dtype_name}") + expected_formula = { + "max_abs_dlogp": "max(abs(dlogp))", + "approx_kl0": "mean(exp(dlogp) - 1 - dlogp)", + "clipfrac0": "mean(1[exp(dlogp) outside clip_interval])", + }[name] + if metrics[name].get("formula") != expected_formula: + raise ContractSchemaError(f"metric {name!r} formula does not match implementation") + if metrics[name].get("pass_rule") != "value <= threshold": + raise ContractSchemaError(f"metric {name!r} pass_rule must be 'value <= threshold'") def _validate_compat_views(contract: Mapping[str, Any]) -> None: @@ -883,6 +898,7 @@ def _lookup_cell( dtype_name: str, arch_key: str | None, ) -> Mapping[str, Any] | None: + base = judgment_root.get("by_op_class", {}).get(op_class, {}).get(dtype_name) if arch_key is not None: arch_cell = ( judgment_root.get("arch_overrides", {}) @@ -891,8 +907,10 @@ def _lookup_cell( .get(dtype_name) ) if arch_cell is not None: - return arch_cell - return judgment_root.get("by_op_class", {}).get(op_class, {}).get(dtype_name) + if base is None: + return arch_cell + return {**base, **arch_cell} + return base def _dtype_name(dtype: str | Any) -> str: diff --git a/tests/test_op_checks.py b/tests/test_op_checks.py index 35de05a5..de2ceb22 100644 --- a/tests/test_op_checks.py +++ b/tests/test_op_checks.py @@ -225,7 +225,7 @@ def test_ws1_report_rejects_backend_provenance_mismatch(): provenance = BackendProvenance( backend_profile="cuda_bf16", requested_backend="cuda", - actual_backend="triton", + actual_backend="cuda", execution_dtype="bfloat16", accumulation_dtype="float32", output_dtype="bfloat16", diff --git a/tests/test_tolerance_contract.py b/tests/test_tolerance_contract.py index 5f711061..4fcbe23f 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -113,22 +113,28 @@ def test_invariance_rows_are_bitwise_zero(): def test_cuda_and_triton_profiles_share_thresholds(): contract = load_contract() - for profile in ("cuda_bf16", "triton_cuda_bf16"): - a = resolve_tolerance( - contract, - judgment="forward_accuracy", - op_class="logprob", - dtype="bfloat16", - backend_profile=profile, - ) - b = resolve_tolerance( - contract, - judgment="forward_accuracy", - op_class="logprob", - dtype="bfloat16", - backend_profile="cuda_bf16" if profile != "cuda_bf16" else "triton_cuda_bf16", - ) - assert a.atol == b.atol and a.rtol == b.rtol and a.mode == b.mode + for judgment in JUDGMENTS: + for op_class in OP_CLASSES: + for dtype_name in ("float32", "bfloat16"): + cuda = resolve_tolerance( + contract, + judgment=judgment, + op_class=op_class, + dtype=dtype_name, + backend_profile="cuda_bf16", + ) + triton = resolve_tolerance( + contract, + judgment=judgment, + op_class=op_class, + dtype=dtype_name, + backend_profile="triton_cuda_bf16", + ) + assert (cuda.atol, cuda.rtol, cuda.mode) == ( + triton.atol, + triton.rtol, + triton.mode, + ) def test_unknown_backend_profile_hard_fails(): @@ -313,9 +319,15 @@ def test_resolve_tolerance_attaches_roles(): def test_chain_aggregate_named_resolve(): contract = load_contract() + expected = { + "max_abs_dlogp": {"bfloat16": 5.0e-2, "float32": 1.0e-5}, + "approx_kl0": {"bfloat16": 5.0e-2, "float32": 1.0e-5}, + "clipfrac0": {"bfloat16": 0.0, "float32": 0.0}, + } + assert set(expected) == set(CHAIN_AGGREGATE_METRICS) for metric in CHAIN_AGGREGATE_METRICS: - thr = resolve_chain_aggregate_thresholds(contract, metric, "bfloat16") - assert thr >= 0.0 + for dtype, value in expected[metric].items(): + assert resolve_chain_aggregate_thresholds(contract, metric, dtype) == value with pytest.raises(ContractResolveError, match="unknown chain aggregate"): resolve_chain_aggregate_thresholds(contract, "mean_abs_dlogp", "bfloat16") @@ -398,6 +410,21 @@ def test_nan_inf_hard_fail(): comparison_rhs_role="inference_style_rollout_decode", ) + # Finite dlogp can still overflow exp(dlogp), which is a separate hard-fail. + lhs = torch.tensor([200.0, 0.0], dtype=torch.float32) + rhs = torch.zeros(2) + with pytest.raises(ContractResolveError, match="ratio0"): + compute_logprob_aggregates( + lhs, + rhs, + torch.ones(2, dtype=torch.bool), + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + lhs = torch.tensor([float("inf"), 0.0]) with pytest.raises(ContractResolveError, match="NaN/Inf"): compute_logprob_aggregates( @@ -412,6 +439,50 @@ def test_nan_inf_hard_fail(): ) +def test_inactive_nan_is_ignored(): + agg = compute_logprob_aggregates( + torch.tensor([0.0, float("nan")]), + torch.zeros(2), + torch.tensor([True, False]), + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + assert agg.active_token_count == 1 + assert agg.max_abs_dlogp == 0.0 + + +def test_clipfrac0_counts_ratios_outside_the_interval(): + agg = compute_logprob_aggregates( + torch.tensor([0.0, 1.0, -1.0]), + torch.zeros(3), + torch.ones(3, dtype=torch.bool), + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(0.8, 1.2), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + assert math.isclose(agg.clipfrac0, 2.0 / 3.0, rel_tol=0.0, abs_tol=1e-6) + + +def test_clip_interval_endpoints_count_as_inside(): + lo, hi = 0.5, 2.0 + agg = compute_logprob_aggregates( + torch.tensor([math.log(lo), math.log(hi)]), + torch.zeros(2), + torch.ones(2, dtype=torch.bool), + contract=load_contract(), + report_kind="train_infer_logprob_parity", + clip_interval=(lo, hi), + comparison_lhs_role="training_style_teacher_forcing", + comparison_rhs_role="inference_style_rollout_decode", + ) + assert agg.clipfrac0 == 0.0 + + def test_judge_requires_all_three_aggregates(): contract = load_contract() clip = default_clip_interval(contract) @@ -434,10 +505,11 @@ def test_judge_requires_all_three_aggregates(): assert {m.metric for m in verdict.metrics} == set(CHAIN_AGGREGATE_METRICS) assert all(m.passed for m in verdict.metrics) - # Large drift fails max_abs_dlogp / approx_kl0 / possibly clipfrac. - lhs = torch.tensor([0.0, 1.0]) - rhs = torch.zeros(2) - mask = torch.ones(2, dtype=torch.bool) + # A small in-interval drift fails only max_abs_dlogp. This proves the + # overall verdict requires all three metrics, rather than any one metric. + lhs = torch.tensor([0.1]) + rhs = torch.zeros(1) + mask = torch.ones(1, dtype=torch.bool) agg = compute_logprob_aggregates( lhs, rhs, @@ -450,7 +522,12 @@ def test_judge_requires_all_three_aggregates(): ) verdict = judge_logprob_aggregates(agg, contract, execution_dtype="bfloat16") assert not verdict.passed - assert any(not m.passed for m in verdict.metrics) + by_metric = {metric.metric: metric.passed for metric in verdict.metrics} + assert by_metric == { + "max_abs_dlogp": False, + "approx_kl0": True, + "clipfrac0": True, + } def test_compat_accuracy_mirrors_forward_accuracy():