diff --git a/docs/contributing/gtest-usage.md b/docs/contributing/gtest-usage.md new file mode 100644 index 00000000..40438c77 --- /dev/null +++ b/docs/contributing/gtest-usage.md @@ -0,0 +1,391 @@ +# 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 + +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 | +| Profiles | `cuda_bf16` and `triton_cuda_bf16` share **the same** thresholds | + +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. + +`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=...)` + +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 | +|-----|---------| +| [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..749b203e 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -2,6 +2,19 @@ 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) + ## Dispatch Tests ```bash @@ -14,6 +27,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/rl_engine/kernels/gtest/__init__.py b/rl_engine/kernels/gtest/__init__.py index c3fc3665..a12db99e 100644 --- a/rl_engine/kernels/gtest/__init__.py +++ b/rl_engine/kernels/gtest/__init__.py @@ -2,9 +2,29 @@ # Copyright (c) 2026 RL-Kernel Contributors 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, +) __all__ = [ "CandidateSpec", "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 efea31a3..1bada7b3 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..0d2ae2e7 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -1,20 +1,995 @@ # 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: + 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", + }: + 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", + "active_token_policy", + "clip_interval_field", + "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 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"]) + 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}") + 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: + """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: + 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", {}) + .get(arch_key, {}) + .get(op_class, {}) + .get(dtype_name) + ) + if arch_cell is not None: + if base is None: + return arch_cell + return {**base, **arch_cell} + return base + + +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..de2ceb22 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="cuda", + 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..4fcbe23f 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,497 @@ 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 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(): + 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() + 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: + 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") + + +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", + ) + + # 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( + 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_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) + # 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) + + # 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, + 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 + 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(): + 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)