diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..bc7009cf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,10 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run Mismatch Framework Tests (CPU-safe) + run: | + python -m pytest tests/test_mismatch_framework.py tests/test_mismatch_logprob_adapter.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/rl_engine/mismatch/README.md b/rl_engine/mismatch/README.md new file mode 100644 index 00000000..8446577c --- /dev/null +++ b/rl_engine/mismatch/README.md @@ -0,0 +1,121 @@ + + + +# `rl_engine/mismatch` — detecting and reporting training-inference mismatch + +The rollout policy `π_old` and the training policy `π_θ` compute logprobs for the +**same tokens with the same weights** and still disagree. This package turns +"which of the dozens of possible causes is it" into named **factors**, each a +switch that can be flipped one at a time and attributed to a side. + +The output is a report: which factors were measured, which could not be and why, +and the few most suspicious modules ranked. + +## Why the tail, not the mean + +`dlogp = log π_θ − log π_old` enters the GRPO objective through `ρ = exp(dlogp)`, +which the objective clips at `1 ± ε`. With `ε = 0.2`, any token past +`|dlogp| > ln(1.2) ≈ 0.182` has its **gradient signal discarded** — and not +random tokens, the most mismatched ones. + +A healthy `dlogp_mean` is `0.002–0.008` (dense) or `0.01–0.03` (large MoE), all +far below that edge. **Judging on the mean alone always concludes "everything is +fine."** Hence `dlogp_p99` / `dlogp_max` / `clip_fraction` / `worst_token`, and a +diagnosis matrix that converges on `clip_fraction`. + +## Four arms, then four gates + +A factor expands into four arms, not on/off — only a one-sided swap identifies a +side, and only a two-sided swap proves the reference itself is sound: + +| arm | rollout | training | what it buys | +|---|---|---|---| +| `both_native` | native | native | the baseline the others are measured against | +| `both_reference` | reference | reference | **self-check gate** — must be bitwise identical, or this factor's conclusions are void | +| `training_reference_only` | native | reference | deviation gone ⇒ training side is the source | +| `rollout_reference_only` | reference | native | deviation gone ⇒ rollout side is the source | + +A factor with no reference implementation is a **parameter sweep** instead: one +arm per allowed value. A sweep measures; it cannot conclude, because nothing was +swapped and so there is no side to attribute to. + +Before any verdict, four gates run. **"Not measured" and "measured and clean" +are different things**, and confusing them is the mistake an attribution +framework is most likely to make: + +| gate | fails when | verdict | +|---|---|---| +| 1 · did it apply | any arm is not `APPLIED` | `VARIANT_DID_NOT_APPLY`, with the resolution trace | +| 2 · evidence | `required_evidence` incomplete | `INSUFFICIENT_EVIDENCE` | +| 3 · shards | fewer logprob shards than `world_size` | `INSUFFICIENT_EVIDENCE` | +| 4 · guards | a pitfall guard failed | `INSUFFICIENT_EVIDENCE` | + +`SwitchStatus.FELL_BACK` is why gate 1 exists: the reference was requested, the +engine silently reverted to native, and "the deviation did not change" then reads +as a clean `NOT_THIS_FACTOR`. + +## Noise floors + +A result only means something at a floor that can resolve it. Each step down +adds exactly one new noise source, so a failure points at a known suspect set. A +floor that has not passed blocks the next. + +| floor | configuration | new noise source | +|---|---|---| +| `SINGLE_LAYER_ANCHOR` | 1 layer, single device, determinism on | none — failing bitwise here is an **operator bug**, not mismatch | +| `FULL_MODEL_SINGLE_GPU` | all layers, single device | accumulation over depth | +| `SHARDED_SINGLE_NODE` | TP + SP on one node | reduction order — the first floor with *real* mismatch | +| `PRODUCTION` | target TP/CP/PP, determinism off, decode | everything else; the only floor readable against `EXPECTED_RANGES` | + +## Layout + +``` +mismatch/ +├── schema/ pure data types, frozen, no behaviour +├── pipeline/ registry → planner → runner → diagnosis → report +├── engines/ the two sides under test: megatron.py, vllm.py +├── reference_adapters/ delivering pinned settings, and reading them back +├── model_meta/ per-model correspondence and call chain (qwen3.py) +├── operator_checks/ plugins, one directory per operator +├── docs/ tutorials +└── __main__.py CLI, and the only module that imports plugins +``` + +`engines/` holds **`megatron.py` and `vllm.py` and nothing else** — the two +policies as they really run, shared across operators. Anything that merely +satisfies `ScoringBackend` is a harness, not a side under test, and lives in +`tests/` (see `tests/mismatch_cpu_backend.py`). The line is role, not protocol. + +Three dependency rules keep the plugin seam open: + +1. `schema/` never imports `pipeline/`; inside `schema/`, `values.py` imports + nothing from the project. +2. `pipeline/` never imports `operator_checks/` — it sees only what the registry + hands it. Break this and adding an operator becomes changing the framework. +3. Only `__main__` imports `operator_checks/`, which triggers registration. + +## Running it + +```bash +python -m rl_engine.mismatch list # operators and their factors +python -m rl_engine.mismatch plan --gpu-count 2 # expand into cases, cheapest first +python -m rl_engine.mismatch plan --json +``` + +The attention and logprob adapters are wired end to end; GEMM still raises +`NotImplementedError`. Attention fails closed when an engine does not report +the actual Split-KV plan set, CP block manifest, collective trace, or RoPE +evidence. Logprob remains the smaller worked example of the adapter layer. + +## Adding to this package + +**Adding an operator is adding a directory; adding a factor is adding a file.** +No existing file changes, apart from one line in `__main__._OPERATOR_PACKAGES` +for a new operator. If your change needs an edit inside `pipeline/`, a global +dict, or another operator's directory, the framework is missing an abstraction — +raise it rather than patching around it. + +| you want to | read | +|---|---| +| add a kernel's factor | [`docs/add-a-kernel-factor.md`](docs/add-a-kernel-factor.md) | +| add a communication feature | [`docs/add-a-comm-feature.md`](docs/add-a-comm-feature.md) | diff --git a/rl_engine/mismatch/__init__.py b/rl_engine/mismatch/__init__.py new file mode 100644 index 00000000..b99d452e --- /dev/null +++ b/rl_engine/mismatch/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-inference mismatch diagnosis, one factor at a time. + +Rollout and training compute logprobs for the same tokens with the same weights +and still disagree. This turns "which of the dozens of possible causes is it" +into switches that can be flipped one at a time and attributed to a side. + +Three dependency rules keep the plugin seam open: + +1. ``schema/`` never imports ``pipeline/``, and inside ``schema/`` the imports go + one way, with ``values.py`` importing nothing from the project. +2. ``pipeline/`` never imports ``operator_checks/``; it sees only what the + registry hands it. Break this and adding an operator becomes changing the + framework. +3. Only ``__main__`` imports ``operator_checks/``, to trigger self-registration. + +See ``README.md`` for the layout and ``docs/`` for how to add a factor. +""" + +from rl_engine.mismatch import pipeline, schema + +__all__ = [ + "pipeline", + "schema", +] diff --git a/rl_engine/mismatch/__main__.py b/rl_engine/mismatch/__main__.py new file mode 100644 index 00000000..8fd3c566 --- /dev/null +++ b/rl_engine/mismatch/__main__.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Command line entry point, and the only place that imports operator plugins. + +Importing them here is what triggers self-registration, and it keeps the +dependency arrow pointing one way: ``pipeline/`` never reaches into +``operator_checks/``. +""" + +from __future__ import annotations + +import argparse +import importlib +import json +import sys +from typing import Sequence + +from rl_engine.mismatch.pipeline import ( + OPERATOR_CHECKS, + build_variants, + missing_prerequisites, + order_cases_by_rebind_cost, + reject_contradictory_factors, +) +from rl_engine.mismatch.schema import NoiseFloor + +# An operator that is not listed here does not exist as far as the framework +# is concerned. +_OPERATOR_PACKAGES: tuple[str, ...] = ( + "rl_engine.mismatch.operator_checks.gemm", + "rl_engine.mismatch.operator_checks.attention", + "rl_engine.mismatch.operator_checks.logprob", +) + + +def load_operator_plugins(packages: Sequence[str] = _OPERATOR_PACKAGES) -> tuple[str, ...]: + """Import each plugin package so its decorator runs.""" + + for package in packages: + importlib.import_module(package) + return OPERATOR_CHECKS.operators() + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m rl_engine.mismatch", + description="Diagnose training-inference mismatch, one factor at a time.", + ) + sub = parser.add_subparsers(dest="command", required=True) + + listing = sub.add_parser("list", help="list registered operators and their factors") + listing.add_argument("--operator", default=None) + + plan = sub.add_parser("plan", help="expand factors into variants without running anything") + plan.add_argument("--operator", default=None) + plan.add_argument( + "--noise-floor", + default=NoiseFloor.SINGLE_LAYER_ANCHOR.value, + choices=[floor.value for floor in NoiseFloor], + ) + plan.add_argument("--gpu-count", type=int, default=0) + plan.add_argument("--json", action="store_true") + + return parser + + +def command_list(operator: str | None) -> int: + operators = load_operator_plugins() + if not operators: + print( + "no operator plugins registered.\n" + "The framework ships without operators: add one under " + "rl_engine/mismatch/operator_checks// and list it in " + "__main__._OPERATOR_PACKAGES.\n" + "See rl_engine/mismatch/operator_checks/__init__.py for the layout." + ) + return 0 + + for name in operators: + if operator is not None and name != operator: + continue + factors = OPERATOR_CHECKS.factors_for(name) + print(f"{name}: {len(factors)} factors") + for factor in factors: + print(f" {factor.id:<40} {factor.category.value}") + return 0 + + +def command_plan(operator: str | None, noise_floor: str, gpu_count: int, as_json: bool) -> int: + load_operator_plugins() + factors = OPERATOR_CHECKS.factors_for(operator) + if not factors: + print("nothing to plan: no operator plugins are registered.") + return 0 + + reject_contradictory_factors(factors) + + runnable = [] + skipped = [] + for factor in factors: + unmet = missing_prerequisites(factor, gpu_count=gpu_count) + if unmet: + skipped.append((factor, unmet)) + else: + runnable.append(factor) + + cases = [(factor, variant) for factor in runnable for variant in build_variants(factor)] + ordered = order_cases_by_rebind_cost(cases) + + if as_json: + payload = { + "noise_floor": noise_floor, + "runnable_factors": [factor.id for factor in runnable], + "skipped": {factor.id: [item.reason for item in unmet] for factor, unmet in skipped}, + "cases": [ + { + "factor": factor.id, + "variant": variant.name, + "rebind_cost": factor.switch.rebind_cost.value, + "switch_values": dict(variant.switch_values), + } + for factor, variant in ordered + ], + } + print(json.dumps(payload, indent=2, sort_keys=True)) + return 0 + + print(f"noise floor: {noise_floor}") + print(f"runnable factors: {len(runnable)} cases: {len(ordered)}") + if skipped: + print("\nskipped (prerequisites not met):") + for factor, unmet in skipped: + for item in unmet: + print(f" {factor.id}: {item.reason}") + print("\ncases in execution order (cheapest rebuild first):") + for factor, variant in ordered: + print(f" [{factor.switch.rebind_cost.value:<22}] {factor.id} :: {variant.name}") + return 0 + + +def main(argv: Sequence[str] | None = None) -> int: + args = build_parser().parse_args(argv) + if args.command == "list": + return command_list(args.operator) + if args.command == "plan": + return command_plan(args.operator, args.noise_floor, args.gpu_count, args.json) + return 1 + + +if __name__ == "__main__": # pragma: no cover + sys.exit(main()) diff --git a/rl_engine/mismatch/docs/README.md b/rl_engine/mismatch/docs/README.md new file mode 100644 index 00000000..bd454436 --- /dev/null +++ b/rl_engine/mismatch/docs/README.md @@ -0,0 +1,11 @@ + + + +# Tutorials + +Concepts live in [`../README.md`](../README.md); these are the how-to. + +| tutorial | when | +|---|---| +| [add-a-kernel-factor.md](add-a-kernel-factor.md) | a kernel computes something different on the two sides | +| [add-a-comm-feature.md](add-a-comm-feature.md) | the suspect is a collective: a reduction order, a rewrite, a CP merge, a backend | diff --git a/rl_engine/mismatch/docs/add-a-comm-feature.md b/rl_engine/mismatch/docs/add-a-comm-feature.md new file mode 100644 index 00000000..cf4cf08f --- /dev/null +++ b/rl_engine/mismatch/docs/add-a-comm-feature.md @@ -0,0 +1,151 @@ + + + +# Add a communication feature + +Mismatch comes from floating-point addition not being associative, and the +accumulation order is almost entirely decided by collective communication — so +collectives are declared objects here, not an implementation detail. Six factors +across gemm, attention, logprob and MoE are instances of one semantic model; +written separately they drift apart. + +This covers only what differs from +[add-a-kernel-factor.md](add-a-kernel-factor.md), which you should read first. +Worked example: `gemm.forward_reduce`. + +## Describe the collective, do not just name it + +```python +ORDERED_REDUCE_SCATTER = CollectiveContract( + op=CollectiveOp.REDUCE_SCATTER, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.GLOBAL_RANK_INDEX, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.STABLE_ACROSS_TOPOLOGY, + backend="rl_kernel", +) +``` + +| field | what it decides | +|---|---| +| `op` | which collective. `NONE` is a real value — record the single-device path rather than leaving it blank | +| `reduction_order` | **the direct root of mismatch.** `ARRIVAL` and `NCCL_ALGORITHM` are unstable; the `GLOBAL_*_INDEX` orders are fixed | +| `accumulate_precision` + `downcast_at` | how much error the accumulation keeps | +| `determinism` | how strong a reproducibility guarantee this offers | +| `backend` | `nccl` / `vllm_custom_ipc` / `mnnvl` / `transformer_engine` / `rl_kernel` | + +**One combination is rejected before anything runs**: claiming +`STABLE_ACROSS_TOPOLOGY` while reducing with `NCCL_ALGORITHM` or `ARRIVAL` +produces numbers that mean nothing. Do not weaken the claim to get past +`reject_contradictory_factors()`; fix whichever half is wrong. + +## Pin the contract so the planner can see it + +`declared_collectives()` finds your contract among the reference's +`required_settings`, so it is pinned like any other setting, with the channel +that actually delivers it: + +```python +required_settings=( + RequiredSetting("forward_reduce_contract", ORDERED_REDUCE_SCATTER, + SettingChannel.CALL_ARG, + readback="module.last_collective_contract"), + RequiredSetting("NCCL_ALGO", "Ring", SettingChannel.ENV_VAR, + readback="os.environ", guards="nccl_algo_unpinned"), +) +``` + +Communication is one of the few places where `SELF_WRITTEN` is the honest +answer: neither TE nor FlashInfer exposes a reduction whose order is fixed across +topologies. Say that in the PR rather than leaving the tier unexplained. + +## What differs in the factor declaration + +See `operator_checks/gemm/factors/forward_reduce.py`. Four things are specific to +a comm factor: + +- **Indexed paths reach into the collective** — `collectives[0].reduction_order` + resolves through the tuple, so no operator-specific comparison code is needed. +- **`backend` is `RECORD_ONLY`.** Two backends may legitimately differ; what must + agree is the order. Comparing it buries the real finding. +- **`min_gpu_count=2` and `PROCESS_GROUP_REBUILD`.** The factor is identically + zero on one device and belongs at `SHARDED_SINGLE_NODE` or above. +- **`call_sites`** records one factor acting in several physical places — + attention's O-linear, the MLP's down-linear and the MoE output are all row + parallel linears with the same accumulation-order problem. One factor, three + sites. + +`compare_contracts()` adds a check you do not declare: two sides promising +different `DeterminismLevel`s emit `DETERMINISM_INCOMPATIBLE`, because comparing +against something not reproducible across runs measures the weaker side's noise +rather than the gap. + +## The cheapest strong check + +An implementation claiming `STABLE_ACROSS_TOPOLOGY` must produce bitwise +identical results when NCCL uses a different algorithm. No cross-framework +comparison, just reruns — the highest value per minute in the framework. + +```python +repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")} +``` + +The runner expands the cartesian product and asserts bitwise agreement; a +disagreement sets `ERROR`, which gate 1 turns into `VARIANT_DID_NOT_APPLY`. + +This protects **the premise of the self-check gate**: `both_reference` can only +anchor the other arms if the fixed-order implementation really did fix the order. +Without it, `REFERENCE_ITSELF_IS_BROKEN` is itself untrustworthy. + +To use `repeat_under` you pass the arms explicitly through +`MismatchFactor.variants` — a non-empty tuple is returned as-is by +`build_variants`. + +## Observe what actually ran + +`build_contract()` says what was asked for; `observe_collectives()` says what +happened. vLLM switching between custom IPC, MNNVL and NCCL by world size and +topology is exactly where the two disagree, and only the second is evidence. It +feeds `COLLECTIVE_CONTRACT`, which gate 2 requires. + +## Rewrites + +"Mathematically identical, unequal in floating point" is its own type; +`schema/collectives.py` declares two: + +```python +ALL_REDUCE_AS_SCATTER_GATHER # all_reduce -> reduce_scatter + all_gather +ALL_TO_ALL_AS_GATHER_SLICE # all_to_all -> all_gather + local slice +``` + +`preserves_bitwise` is `False` on both and always will be — that is the entire +problem. Megatron applying the first with sequence parallelism on while the +rollout side does not *is* `gemm.forward_reduce`. If your factor is "one side +rewrites this collective", add a constant next to those two rather than +describing the rewrite in prose. + +## Adding an enum value + +A new `CollectiveOp`, `ParallelDim` or `ReductionOrder` **is** a framework +change. Justified when the semantics cannot be expressed by the existing values — +a new parallel dimension, or a fixed order keyed on something that is neither +rank, block, nor vocab shard. Not justified for a new backend (that is the +`backend` string) or a new library version (that is `LibraryPin`). + +If you do add one: put it in `schema/collectives.py`; extend +`_NON_DETERMINISTIC_ORDERS` in `pipeline/planner.py` if the order is not stable, +so the contradiction check keeps working; add a test; and say in the PR why an +existing value could not express it. + +## Checklist + +- [ ] Every `CollectiveContract` field comes from what the code does, not what the config requests. +- [ ] `determinism` and `reduction_order` do not contradict each other. +- [ ] `backend` is `RECORD_ONLY`; `group_size` is `MUST_MATCH_BITWISE`. +- [ ] `min_gpu_count` ≥ 2, floor `SHARDED_SINGLE_NODE` or above. +- [ ] `NCCL_ALGO` / `NCCL_PROTO` pinned **with readback**. +- [ ] Any topology-independence claim is backed by a `repeat_under` arm. +- [ ] `observe_collectives()` returns the trace; `COLLECTIVE_CONTRACT` is required evidence. +- [ ] `call_sites` lists every place this one factor acts. diff --git a/rl_engine/mismatch/docs/add-a-kernel-factor.md b/rl_engine/mismatch/docs/add-a-kernel-factor.md new file mode 100644 index 00000000..d706e7ea --- /dev/null +++ b/rl_engine/mismatch/docs/add-a-kernel-factor.md @@ -0,0 +1,192 @@ + + + +# Add a kernel's mismatch factor + +You suspect a kernel of computing something different on the two sides and want +that measured and attributed instead of argued about. Worked example: +`attn.rope_fusion`. + +Read [`../README.md`](../README.md) first — the four arms, the four gates and the +noise floors are assumed here. You will write **declarations only**; the +framework expands the arms, runs them, gates them, and concludes. + +## Three decisions before writing anything + +**Sweep or swap?** `reference is None` makes it a sweep (one arm per allowed +value); a `ReferenceImplementation` makes it a swap (the four arms). There is no +`kind` field — derivable state is state that can disagree with itself. RoPE is a +swap; `operator_checks/logprob/` has the worked sweep. + +**Which reference?** `ReferenceAuthority` is a decision order: + +``` +FP64_ORACLE slow, exact — gold standard at the lowest floor + ↑ validates +SHARED_BACKEND TransformerEngine / FlashInfer — look here FIRST + ↑ falls back to +SELF_WRITTEN only when the first two have a real semantic hole +``` + +A PR adding a `SELF_WRITTEN` reference must say why the first two cannot cover +it. RoPE is covered by TE and FlashInfer, so it is `SHARED_BACKEND`. + +**Which floor?** The lowest one where the factor is not identically zero. A +`PROCESS_GROUP_REBUILD` switch is identical on a single device, so running it at +the anchor floor only burns machine time. + +## Step 1 — the directory + +Skip to step 3 if the operator exists; then you are only adding one file. + +``` +operator_checks/attention/ +├── __init__.py operator name + discover_factors +├── adapter.py the four operator-level methods +├── _common.py reference implementations, contract helpers +└── factors/ + ├── __init__.py + └── rope_fusion.py +``` + +**A factor file's name is its id with the operator prefix stripped.** +`discover_factors()` enforces it, so renaming an id without renaming the file +fails at import rather than silently dropping the factor. + +## Step 2 — `_common.py` + +```python +TE_ROPE_REFERENCE = ReferenceImplementation( + name="transformer_engine", + tier=ReferenceAuthority.SHARED_BACKEND, + training_impl="transformer_engine.pytorch.attention.rope.apply_rotary_pos_emb", + rollout_impl="flashinfer.rope.apply_rope", + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ), + required_settings=( + RequiredSetting( + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0", + SettingChannel.ENV_VAR, readback="os.environ", + ), + ), + pinned_libraries=(LibraryPin("transformer_engine", "2.9.0.dev0", commit="8260f49"),), +) +``` + +Three fields decide more than they look like they do: + +- **`covers_paths` defines the shape of the self-check gate**: every path this + reference covers must agree bitwise on the same sequence. Two paths put the + gate across the two sides; a reference that also covers `ROLLOUT_DECODE` puts + it *inside* the rollout side, and then no decode stub is needed on the training + side. +- **`required_settings` are not documentation.** Each is delivered by its channel + and verified by readback. `readback=None` can only ever be `UNOBSERVABLE`. +- **`pinned_libraries` is required.** TE and FlashInfer change kernel selection + across versions, so the same factor can reach the opposite conclusion on a + different one. + +## Step 3 — the factor file + +One `FACTOR` constant, no behaviour. See +`operator_checks/attention/factors/rope_fusion.py` for the full declaration; the +fields worth thinking about: + +**`comparison_rules`** — keys are dotted paths from the contract root +(`precision.accumulate`, `collectives[0].reduction_order`, `extra.rope_theta`). +The framework indexes both contracts by path, so you never write a "collect the +comparable fields" function. + +| tier | for | consequence | +|---|---|---| +| `MUST_MATCH_BITWISE` | identity: shapes, dtypes, TP size | differing ⇒ the case is void, not a finding | +| `MUST_MATCH_SEMANTICALLY` | may be implemented differently, must mean the same | differing ⇒ a `SEMANTIC_MISMATCH` | +| `RECORD_ONLY` | representation that differs by construction | recorded, never compared | + +`RECORD_ONLY` exists so structural differences do not drown the real problems. +Declare a packed-QKV-style field `MUST_MATCH_*` and your factor disappears under +false positives. The registry also rejects two factors declaring the same +contract field at different tiers — if that fires, one declaration is wrong. + +**`required_evidence`** — three items apply to every factor; operator-specific +ones are plain string constants (`POSITION_CACHE`, `LSE_EXPORT`, …) so adding an +operator never means editing an enum. Missing evidence is gate 2, which is *not* +the same verdict as "measured and clean". + +**`prerequisites`** — a declared whitelist. The planner turns each unmet item +into a reason, so `plan` prints *why* a factor was skipped. + +**`pitfalls`** — `symptom` and `actual_cause` are separate because a pitfall is +one precisely when its appearance points at the wrong cause. `guard_runs_at` +should be the lowest floor that can run the check. + +## Step 4 — `adapter.py` + +The four methods are operator-level, not factor-level: reading configuration back +from an engine is the same logic for all of an operator's factors. See +`operator_checks/attention/adapter.py` for the signatures. + +`resolve_implementation` is where a factor most often dies quietly. Returning a +bare `None` produces `FELL_BACK` with nothing to investigate, and a fallen-back +arm whose deviation "did not change" reads exactly like a clean +`NOT_THIS_FACTOR`. **Return the trace even when resolution fails.** + +## Step 5 — register + +```python +@OPERATOR_CHECKS.register +class AttentionChecks: + operator = "attention" + + def declare_factors(self): + return discover_factors(__package__) + + build_contract = staticmethod(adapter.build_contract) + ... +``` + +Then add the package to `__main__._OPERATOR_PACKAGES` — the one edit outside your +own directory. Adding the next factor drops a file into `factors/` and changes +neither file. + +## Step 6 — check it, without a GPU + +```bash +python -m rl_engine.mismatch list +# attention: 1 factors +# attn.rope_fusion kernel_implementation +``` + +`plan` on a machine without the prerequisites names every unmet one rather than +silently omitting the factor: + +``` +skipped (prerequisites not met): + attn.rope_fusion: operator 'rope' is not dispatchable + attn.rope_fusion: package 'transformer_engine>=2.0' is not installed +``` + +Where they are met, the same declaration expands into the four arms ordered +cheapest-rebuild-first. Then add a test: `tests/mismatch_cpu_backend.py` can +inject a one-sided bias, silently ignore a switch, and be unstable across +environments — enough to prove your factor attributes the right side and that a +fallback is reported rather than mistaken for a clean result. + +## What you do not have to write + +`build_variants()`, `compare_contracts()`, `diagnose()`, `missing_prerequisites()`, +`order_cases_by_rebind_cost()`, and `repeat_under` + `assert_order_is_topology_independent()` +for topology-invariance reruns. + +## Checklist + +- [ ] File name equals the factor id minus the operator prefix. +- [ ] `question` is one line and says what the factor *answers*. +- [ ] Representation-only fields are `RECORD_ONLY`, not `MUST_MATCH_*`. +- [ ] A `SELF_WRITTEN` reference is justified in the PR body. +- [ ] Every `RequiredSetting` has a `readback`, or `UNOBSERVABLE` is accepted. +- [ ] `pinned_libraries` names an exact version. +- [ ] Pitfalls are `KnownPitfall` values, not comments. +- [ ] `list` and `plan` show what you expect. diff --git a/rl_engine/mismatch/engines/__init__.py b/rl_engine/mismatch/engines/__init__.py new file mode 100644 index 00000000..33a835b2 --- /dev/null +++ b/rl_engine/mismatch/engines/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The two sides under test: engine lifetime and configuration readback. + +This package holds ``megatron.py`` and ``vllm.py`` and nothing else. Adding an +operator never adds a file here; all of attention's factors use the one +``vllm.py``. + +Anything that merely satisfies ``ScoringBackend`` is a harness, not a side under +test, and lives in ``tests/`` -- see ``tests/mismatch_cpu_backend.py``. +""" + +from rl_engine.mismatch.engines.megatron import MegatronBackend +from rl_engine.mismatch.engines.vllm import VllmBackend + +__all__ = ["MegatronBackend", "VllmBackend"] diff --git a/rl_engine/mismatch/engines/megatron.py b/rl_engine/mismatch/engines/megatron.py new file mode 100644 index 00000000..f2f788e0 --- /dev/null +++ b/rl_engine/mismatch/engines/megatron.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The training side: Megatron-LM. Not implemented. + +Settings this backend must pin: + + AttnBackend explicit, never "auto" + NVTE_ALLOW_NONDETERMINISTIC_ALGO "0" + CUBLAS_WORKSPACE_CONFIG ":4096:8" + NCCL_ALGO / NCCL_PROTO pinned + torch.backends.cuda.matmul.allow_tf32 False + ...allow_bf16_reduced_precision_reduction False + torch.use_deterministic_algorithms True + lm_head dtype fp32 + +``AttnBackend=auto`` picks a different kernel per shape, so requested and actual +must be recorded separately. ``score()`` must leave the model untouched: a kernel +that mutates weights in place makes ``both_reference`` fail bitwise at random and +the blame lands on the reference. Under TP/CP no single rank holds the whole +logprob vector, so shards are collected per rank and counted against +``world_size``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, Sequence + +from rl_engine.mismatch.schema import ComparisonIdentity, LogprobShard, PolicyRole, ReuseKey + + +@dataclass +class MegatronBackend: + role: PolicyRole = PolicyRole.TRAINING + model: Any = None + effective_config: dict[str, Any] = field(default_factory=dict) + + def score( + self, + role: PolicyRole, + identity: ComparisonIdentity, + switch_values: Mapping[str, Any], + replacement: Callable[..., Any] | None, + ) -> tuple[Sequence[float], Mapping[str, Any]]: + raise NotImplementedError( + "Run one forward over prompt + response, gather the selected " + "logprobs, and assert the model state fingerprint is unchanged." + ) + + def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: + raise NotImplementedError( + "Group switches by tier: determinism env -> process, TP/CP/PP -> " + "process_group, dtype and kernel choice -> engine, batch -> request." + ) + + def read_effective_config(self) -> Mapping[str, Any]: + raise NotImplementedError("Read each pinned setting off the live config.") + + def collect_logprob_shards(self) -> tuple[LogprobShard, ...]: + raise NotImplementedError("Return one shard per rank.") + + +__all__ = ["MegatronBackend"] diff --git a/rl_engine/mismatch/engines/vllm.py b/rl_engine/mismatch/engines/vllm.py new file mode 100644 index 00000000..08678388 --- /dev/null +++ b/rl_engine/mismatch/engines/vllm.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The rollout side: vLLM. Not implemented. + +Settings this backend must pin, and where it reads each one back: + + enable_chunked_prefill False llm_engine.scheduler_config.chunked_prefill_enabled + max_num_batched_tokens > len(prompt + response) + llm_engine.scheduler_config.max_num_batched_tokens + long_prefill_token_threshold 0 llm_engine.scheduler_config.long_prefill_token_threshold + enable_prefix_caching False llm_engine.cache_config.enable_prefix_caching + +``enable_chunked_prefill`` defaults to True, so a full-sequence prefill is +chunked unless it is turned off. ``prompt_logprobs`` already skips the prefix +cache, so unless prefix caching is off too, the full-prefill and decode paths run +against different cache state and are not comparable. Position 0 of +``prompt_logprobs`` is None, which is where the shift-by-one convention has to be +asserted against the training side rather than assumed. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, Sequence + +from rl_engine.mismatch.schema import ComparisonIdentity, PolicyRole, ReuseKey + + +@dataclass +class VllmBackend: + role: PolicyRole = PolicyRole.ROLLOUT + engine: Any = None + effective_config: dict[str, Any] = field(default_factory=dict) + + def score( + self, + role: PolicyRole, + identity: ComparisonIdentity, + switch_values: Mapping[str, Any], + replacement: Callable[..., Any] | None, + ) -> tuple[Sequence[float], Mapping[str, Any]]: + raise NotImplementedError( + "Build the engine with the settings pinned in this module's docstring, " + "run the requested ExecutionPath, and return (logprobs, readback)." + ) + + def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: + raise NotImplementedError( + "Group switches by tier: determinism env -> process, world size and " + "TP/CP -> process_group, dtype and backend and KV layout -> engine, " + "batch and sequence -> request." + ) + + def read_effective_config(self) -> Mapping[str, Any]: + raise NotImplementedError( + "Read each pinned setting off the live engine. A setting with no " + "readback path can only be recorded UNOBSERVABLE." + ) + + +__all__ = ["VllmBackend"] diff --git a/rl_engine/mismatch/model_meta/__init__.py b/rl_engine/mismatch/model_meta/__init__.py new file mode 100644 index 00000000..9d977f5b --- /dev/null +++ b/rl_engine/mismatch/model_meta/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Model shape and the module correspondence table. + +Filled once per model and shared by every operator: "Megatron's ``linear_fc1`` +corresponds to vLLM's ``gate_up_proj``" is the same fact for attention, gemm and +logprob alike. Swapping in another model means redoing it, which is why it lives +here rather than in any ``operator_checks/``. +""" + +from rl_engine.mismatch.model_meta.qwen3 import QWEN3_CORRESPONDENCES, QWEN3_EDGES + +__all__ = [ + "QWEN3_CORRESPONDENCES", + "QWEN3_EDGES", +] diff --git a/rl_engine/mismatch/model_meta/qwen3.py b/rl_engine/mismatch/model_meta/qwen3.py new file mode 100644 index 00000000..a7673c47 --- /dev/null +++ b/rl_engine/mismatch/model_meta/qwen3.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Qwen3 module correspondence and call chain. + +Every ``equivalence`` carries a ``verified_by``: unproven, "filtering false +positives" quietly becomes "hiding real findings". +""" + +from __future__ import annotations + +from rl_engine.mismatch.schema import ModuleCorrespondence, PropagationEdge + +QWEN3_CORRESPONDENCES: tuple[ModuleCorrespondence, ...] = ( + ModuleCorrespondence( + semantic_name="attn.qkv", + training_module="megatron.core.transformer.attention.linear_qkv", + rollout_module="vllm.model_executor.models.qwen3.qkv_proj", + equivalence="concat_on_dim0", + verified_by="tests/test_mismatch_model_meta.py::test_qwen3_qkv_concat_equivalence", + ), + ModuleCorrespondence( + semantic_name="attn.out", + training_module="megatron.core.transformer.attention.linear_proj", + rollout_module="vllm.model_executor.models.qwen3.o_proj", + ), + ModuleCorrespondence( + semantic_name="mlp.gate_up", + training_module="megatron.core.transformer.mlp.linear_fc1", + rollout_module="vllm.model_executor.models.qwen3.gate_up_proj", + equivalence="concat_on_dim0", + verified_by="tests/test_mismatch_model_meta.py::test_qwen3_gate_up_concat_equivalence", + ), + ModuleCorrespondence( + semantic_name="mlp.down", + training_module="megatron.core.transformer.mlp.linear_fc2", + rollout_module="vllm.model_executor.models.qwen3.down_proj", + ), + ModuleCorrespondence( + semantic_name="norm.input", + training_module="megatron.core.transformer.transformer_layer.input_layernorm", + rollout_module="vllm.model_executor.models.qwen3.input_layernorm", + ), + ModuleCorrespondence( + semantic_name="lm_head", + training_module="megatron.core.models.gpt.gpt_model.output_layer", + rollout_module="vllm.model_executor.models.qwen3.lm_head", + ), +) + + +QWEN3_EDGES: tuple[PropagationEdge, ...] = ( + PropagationEdge(upstream="norm.input", downstream="attn.qkv"), + PropagationEdge(upstream="attn.qkv", downstream="attn.out"), + PropagationEdge(upstream="attn.out", downstream="mlp.gate_up"), + PropagationEdge(upstream="mlp.gate_up", downstream="mlp.down"), + PropagationEdge(upstream="mlp.down", downstream="lm_head"), +) + + +QWEN3_8B_SHAPE = "L=36,H=4096,Hq=32,Hkv=8,D=128" +QWEN3_0B5_SHAPE = "L=24,H=896,Hq=14,Hkv=2,D=64" +QWEN3_SINGLE_LAYER_SHAPE = "L=1,H=896,Hq=14,Hkv=2,D=64" + + +__all__ = [ + "QWEN3_0B5_SHAPE", + "QWEN3_8B_SHAPE", + "QWEN3_CORRESPONDENCES", + "QWEN3_EDGES", + "QWEN3_SINGLE_LAYER_SHAPE", +] diff --git a/rl_engine/mismatch/operator_checks/__init__.py b/rl_engine/mismatch/operator_checks/__init__.py new file mode 100644 index 00000000..9964e360 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Operator plugins, one directory per operator. + +Importing this package registers nothing. ``__main__`` imports the individual +operator packages, and that import triggers self-registration -- the framework +does not know which operators exist. + +One operator's layout, and the conventions that are enforced:: + + operator_checks/attention/ + |-- __init__.py operator name + discover_factors + |-- adapter.py the four operator-level methods + |-- _common.py shared reference implementations and contract helpers + `-- factors/ one file per factor, holding one FACTOR constant + +A factor file's name must equal its factor id with the operator prefix stripped; +``discover_factors()`` fails at import otherwise. The four methods stay in +``adapter.py`` because they are operator-level, not factor-level. + +Writing one: ``docs/add-a-kernel-factor.md``, or ``docs/add-a-comm-feature.md`` +when the suspect is a collective. +""" diff --git a/rl_engine/mismatch/operator_checks/attention/__init__.py b/rl_engine/mismatch/operator_checks/attention/__init__.py new file mode 100644 index 00000000..cac2bf21 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/__init__.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The attention operator plugin.""" + +from rl_engine.mismatch.operator_checks.attention import adapter +from rl_engine.mismatch.pipeline import OPERATOR_CHECKS, discover_factors + + +@OPERATOR_CHECKS.register +class AttentionChecks: + operator = "attention" + + def declare_factors(self): + return discover_factors(__package__) + + build_contract = staticmethod(adapter.build_contract) + read_effective_config = staticmethod(adapter.read_effective_config) + observe_collectives = staticmethod(adapter.observe_collectives) + resolve_implementation = staticmethod(adapter.resolve_implementation) diff --git a/rl_engine/mismatch/operator_checks/attention/_common.py b/rl_engine/mismatch/operator_checks/attention/_common.py new file mode 100644 index 00000000..b60e990f --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/_common.py @@ -0,0 +1,638 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Attention references and strict runtime-provenance normalization.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import replace +from typing import Any, Mapping, Sequence + +from rl_engine.mismatch.schema import ( + CollectiveContract, + CollectiveOp, + DeterminismLevel, + DowncastPoint, + ExecutionPath, + LibraryPin, + ParallelDim, + Precision, + ReductionOrder, + ReferenceAuthority, + ReferenceImplementation, + RequiredSetting, + SettingChannel, +) + +ATTENTION_LSE_DOMAIN = "attention" +ATTENTION_MERGE_STATE = "out_lse" +SPLIT_KV_PLAN_EVIDENCE = "split_kv_runtime_plan_set" +CP_BLOCK_MANIFEST_EVIDENCE = "cp_block_manifest" +ATTENTION_LSE_EVIDENCE = "attention_lse_export" +POST_ROPE_QK_EVIDENCE = "post_rope_qk_digest" + +DTYPES: dict[str, Precision] = { + "bf16": Precision.BF16, + "bfloat16": Precision.BF16, + "fp16": Precision.FP16, + "float16": Precision.FP16, + "fp32": Precision.FP32, + "float32": Precision.FP32, +} + +DOWNCAST_POINTS: dict[str, DowncastPoint] = { + "never": DowncastPoint.NEVER, + "per_block": DowncastPoint.PER_BLOCK, + "per_partial": DowncastPoint.PER_PARTIAL, + "final_write": DowncastPoint.FINAL_WRITE, +} + +COLLECTIVE_OPS: dict[str, CollectiveOp] = {item.value: item for item in CollectiveOp} +REDUCTION_ORDERS: dict[str, ReductionOrder] = {item.value: item for item in ReductionOrder} +DETERMINISM_LEVELS: dict[str, DeterminismLevel] = { + item.value: item for item in DeterminismLevel +} + + +class AttentionContractError(ValueError): + """Raised when runtime metadata cannot prove an Attention contract.""" + + +REFERENCE_CP_MERGE = CollectiveContract( + op=CollectiveOp.POINT_TO_POINT, + group=ParallelDim.CONTEXT, + group_size=2, + reduction_order=ReductionOrder.GLOBAL_BLOCK_INDEX, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.STABLE_ACROSS_TOPOLOGY, + backend="p2p_nccl_reference", +) + +TE_ROPE_REFERENCE = ReferenceImplementation( + name="transformer_engine", + tier=ReferenceAuthority.SHARED_BACKEND, + training_impl="transformer_engine.pytorch.attention.rope.apply_rotary_pos_emb", + rollout_impl="flashinfer.rope.apply_rope", + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ), + required_settings=( + RequiredSetting( + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", + "0", + SettingChannel.ENV_VAR, + readback="os.environ", + ), + ), + pinned_libraries=(LibraryPin("transformer_engine", "2.9.0.dev0", commit="8260f49"),), +) + +SPLIT_KV_REFERENCE = ReferenceImplementation( + name="rl_kernel", + tier=ReferenceAuthority.SELF_WRITTEN, + training_impl=( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + rollout_impl=( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ExecutionPath.ROLLOUT_CHUNKED_PREFILL, + ), + pinned_libraries=(LibraryPin("torch", "2.6.0"),), +) + +CP_MERGE_REFERENCE = ReferenceImplementation( + name="p2p_nccl_reference", + tier=ReferenceAuthority.SELF_WRITTEN, + training_impl=( + "rl_engine.kernels.ops.cuda.attention.cp_comm.P2PNCCLAttentionCPCommunication" + ), + rollout_impl=( + "rl_engine.kernels.ops.cuda.attention.cp_comm.P2PNCCLAttentionCPCommunication" + ), + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ExecutionPath.ROLLOUT_CHUNKED_PREFILL, + ), + required_settings=( + RequiredSetting( + "attn.cp_collective", + REFERENCE_CP_MERGE, + SettingChannel.CALL_ARG, + readback="dispatch.provenance['cp_collective']", + ), + ), + pinned_libraries=(LibraryPin("nccl", "2.21.5"),), +) + + +def precision(value: Any, field: str) -> Precision: + if isinstance(value, Precision): + return value + key = str(value).lower() + if key not in DTYPES: + raise AttentionContractError(f"{field} must be one of {tuple(DTYPES)}, got {value!r}") + return DTYPES[key] + + +def downcast_point(value: Any, field: str) -> DowncastPoint: + if isinstance(value, DowncastPoint): + return value + key = str(value).lower() + if key not in DOWNCAST_POINTS: + raise AttentionContractError( + f"{field} must be one of {tuple(DOWNCAST_POINTS)}, got {value!r}" + ) + return DOWNCAST_POINTS[key] + + +def positive_int(value: Any, field: str) -> int: + if isinstance(value, bool): + raise AttentionContractError(f"{field} must be a positive integer, got {value!r}") + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise AttentionContractError( + f"{field} must be a positive integer, got {value!r}" + ) from exc + if parsed <= 0: + raise AttentionContractError(f"{field} must be a positive integer, got {value!r}") + return parsed + + +def optional_positive_int(value: Any, field: str) -> int | None: + return None if value is None else positive_int(value, field) + + +def reference_selected(value: Any, role: Any, reference_name: str, field: str) -> bool: + if value in (None, "native"): + return False + if value == reference_name: + return True + role_value = getattr(role, "value", str(role)) + if value in (f"{reference_name}@training", f"{reference_name}@rollout"): + return value.endswith(f"@{role_value}") + raise AttentionContractError( + f"unknown {field} value {value!r}; expected 'native', {reference_name!r}, " + f"'{reference_name}@training' or '{reference_name}@rollout'" + ) + + +def normalize_collective( + raw: Any, + *, + cp_world_size: int, +) -> tuple[CollectiveContract, ...]: + """Normalize the collective that actually ran; absence stays unknown.""" + + if raw is None: + return () + if isinstance(raw, CollectiveContract): + contract = raw + elif isinstance(raw, Mapping): + op = _enum_member(COLLECTIVE_OPS, raw.get("op"), "attn.cp_collective.op") + order = _enum_member( + REDUCTION_ORDERS, + raw.get("reduction_order"), + "attn.cp_collective.reduction_order", + ) + determinism = _enum_member( + DETERMINISM_LEVELS, + raw.get("determinism"), + "attn.cp_collective.determinism", + ) + contract = CollectiveContract( + op=op, + group=ParallelDim.CONTEXT, + group_size=positive_int( + raw.get("group_size", cp_world_size), "attn.cp_collective.group_size" + ), + reduction_order=order, + accumulate_precision=precision( + raw.get("accumulate_precision"), + "attn.cp_collective.accumulate_precision", + ), + downcast_at=downcast_point( + raw.get("downcast_at"), "attn.cp_collective.downcast_at" + ), + determinism=determinism, + backend=_non_empty_string(raw.get("backend"), "attn.cp_collective.backend"), + ) + else: + raise AttentionContractError( + "attn.cp_collective must be a CollectiveContract or mapping" + ) + + if contract.group is not ParallelDim.CONTEXT: + raise AttentionContractError("Attention CP collective must use the context group") + if contract.group_size != cp_world_size: + raise AttentionContractError( + "Attention CP collective group_size must equal attn.cp_world_size" + ) + if contract.accumulate_precision is not Precision.FP32: + raise AttentionContractError("Attention (out, lse) merge must accumulate in fp32") + return (contract,) + + +def reference_collective(cp_world_size: int) -> tuple[CollectiveContract, ...]: + return (replace(REFERENCE_CP_MERGE, group_size=cp_world_size),) + + +def normalize_split_kv_plan_set(raw: Any) -> Mapping[str, Any] | None: + """Validate and canonicalize complete batch/TP/CP/owner runtime plans. + + A requested policy is deliberately not accepted here. Every coordinate must + carry the actual logical boundaries and numerical merge semantics. + """ + + if raw is None: + return None + if not isinstance(raw, Mapping): + raise AttentionContractError("attn.actual_split_kv_plan_set must be a mapping") + + batch_size = positive_int(raw.get("batch_size"), "split_kv.batch_size") + tp_world_size = positive_int(raw.get("tp_world_size"), "split_kv.tp_world_size") + cp_world_size = positive_int(raw.get("cp_world_size"), "split_kv.cp_world_size") + totals = _int_tuple(raw.get("total_kv_tokens"), "split_kv.total_kv_tokens") + if len(totals) != batch_size or any(value <= 0 for value in totals): + raise AttentionContractError( + "split_kv.total_kv_tokens must contain one positive length per batch item" + ) + + entries_raw = raw.get("entries") + if not isinstance(entries_raw, Sequence) or isinstance(entries_raw, (str, bytes)): + raise AttentionContractError("split_kv.entries must be a sequence") + + expected_coordinates = { + (batch, tp, cp, owner) + for batch in range(batch_size) + for tp in range(tp_world_size) + for cp in range(cp_world_size) + for owner in range(cp_world_size) + } + entries: list[tuple[Any, ...]] = [] + seen: set[tuple[int, int, int, int]] = set() + by_owner: dict[tuple[int, int], list[tuple[Any, ...]]] = defaultdict(list) + + for index, entry_raw in enumerate(entries_raw): + if not isinstance(entry_raw, Mapping): + raise AttentionContractError(f"split_kv.entries[{index}] must be a mapping") + coordinate = ( + _bounded_int(entry_raw.get("batch_index"), batch_size, "batch_index"), + _bounded_int(entry_raw.get("tp_rank"), tp_world_size, "tp_rank"), + _bounded_int(entry_raw.get("cp_rank"), cp_world_size, "cp_rank"), + _bounded_int(entry_raw.get("owner_cp_rank"), cp_world_size, "owner_cp_rank"), + ) + if coordinate in seen: + raise AttentionContractError(f"duplicate Split-KV coordinate {coordinate}") + seen.add(coordinate) + + expected_range = _range_pair( + entry_raw.get("expected_kv_range"), + f"split_kv.entries[{index}].expected_kv_range", + ) + if expected_range[1] > totals[coordinate[0]]: + raise AttentionContractError("Split-KV expected range exceeds total_kv_tokens") + + requested_mode = _mode(entry_raw.get("requested_split_kv_policy"), "requested mode") + actual_mode = _mode(entry_raw.get("actual_split_kv_policy"), "actual mode") + requested_size = optional_positive_int( + entry_raw.get("requested_split_kv_size"), "requested_split_kv_size" + ) + actual_size = optional_positive_int( + entry_raw.get("actual_split_kv_size"), "actual_split_kv_size" + ) + boundaries = _boundaries( + entry_raw.get("actual_split_boundaries"), expected_range, index + ) + reported_count = entry_raw.get("actual_split_kv_count") + if reported_count is not None and positive_int( + reported_count, "actual_split_kv_count" + ) != len(boundaries): + raise AttentionContractError( + "actual_split_kv_count must equal the number of actual boundaries" + ) + merge_order = _enum_member( + REDUCTION_ORDERS, + entry_raw.get("split_kv_merge_order"), + "split_kv_merge_order", + ) + accumulate = precision( + entry_raw.get("split_kv_accum_dtype"), "split_kv_accum_dtype" + ) + downcast = downcast_point( + entry_raw.get("split_kv_downcast_at"), "split_kv_downcast_at" + ) + backend = _non_empty_string( + entry_raw.get("split_kv_backend"), "split_kv_backend" + ) + source = _non_empty_string( + entry_raw.get("split_kv_plan_source"), "split_kv_plan_source" + ) + fallback = entry_raw.get("split_kv_fallback") + if not isinstance(fallback, bool): + raise AttentionContractError("split_kv_fallback must be a bool") + fallback_reason = entry_raw.get("split_kv_fallback_reason") + if fallback and (not isinstance(fallback_reason, str) or not fallback_reason.strip()): + raise AttentionContractError( + "a Split-KV fallback must include a non-empty fallback_reason" + ) + if not fallback and fallback_reason is not None: + raise AttentionContractError( + "split_kv_fallback_reason must be None when fallback is false" + ) + + _validate_split_mode_sizes( + requested_mode, + requested_size, + actual_mode, + actual_size, + boundaries, + fallback, + ) + if merge_order is not ReductionOrder.GLOBAL_BLOCK_INDEX: + raise AttentionContractError( + "Split-KV partials must merge in global_block_index order" + ) + if accumulate is not Precision.FP32: + raise AttentionContractError("Split-KV partials must accumulate in fp32") + if downcast is not DowncastPoint.FINAL_WRITE: + raise AttentionContractError("Split-KV partials may downcast only at final_write") + + canonical = ( + *coordinate, + expected_range, + requested_mode, + requested_size, + actual_mode, + actual_size, + boundaries, + merge_order, + accumulate, + downcast, + backend, + source, + fallback, + fallback_reason, + ) + entries.append(canonical) + by_owner[(coordinate[0], coordinate[3])].append(canonical) + + missing = expected_coordinates - seen + extra = seen - expected_coordinates + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan coverage is incomplete; " + f"missing={sorted(missing)}, extra={sorted(extra)}" + ) + + _validate_owner_ranges(entries, batch_size, tp_world_size, cp_world_size, totals) + for owner, owner_entries in by_owner.items(): + reference = _plan_semantics(owner_entries[0]) + if any(_plan_semantics(entry) != reference for entry in owner_entries[1:]): + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={owner[0]}, owner_cp={owner[1]}" + ) + + ordered = tuple(sorted(entries, key=lambda item: item[:4])) + return { + "batch_size": batch_size, + "tp_world_size": tp_world_size, + "cp_world_size": cp_world_size, + "total_kv_tokens": totals, + "coordinates": tuple(item[:4] for item in ordered), + "owner_ranges": tuple((item[:4], item[4]) for item in ordered), + "boundaries": tuple((item[:4], item[9]) for item in ordered), + "merge_order": tuple((item[:4], item[10]) for item in ordered), + "accumulate_precision": tuple((item[:4], item[11]) for item in ordered), + "downcast_at": tuple((item[:4], item[12]) for item in ordered), + "fallback": tuple((item[:4], item[15], item[16]) for item in ordered), + "backend": tuple((item[:4], item[13]) for item in ordered), + "source": tuple((item[:4], item[14]) for item in ordered), + "canonical": tuple(_cross_side_plan_semantics(item) for item in ordered), + } + + +def normalize_cp_block_manifest( + raw: Any, + *, + tp_world_size: int, + cp_world_size: int, +) -> tuple[tuple[int, int, int, int, int], ...] | None: + """Validate global block ownership and gap-free KV range coverage.""" + + if raw is None: + return None + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)) or not raw: + raise AttentionContractError("attn.cp_block_manifest must be a non-empty sequence") + + blocks: list[tuple[int, int, int, int, int]] = [] + for index, item in enumerate(raw): + if not isinstance(item, Mapping): + raise AttentionContractError(f"attn.cp_block_manifest[{index}] must be a mapping") + block_index = _non_negative_int(item.get("global_block_index"), "global_block_index") + start, end = _range_pair( + (item.get("kv_block_start"), item.get("kv_block_end")), + "KV block range", + ) + owner_cp = _bounded_int(item.get("owner_cp_rank"), cp_world_size, "owner_cp_rank") + owner_tp = _bounded_int(item.get("owner_tp_rank"), tp_world_size, "owner_tp_rank") + blocks.append((block_index, start, end, owner_cp, owner_tp)) + + ordered = tuple(sorted(blocks)) + if len({block[0] for block in ordered}) != len(ordered): + raise AttentionContractError("CP block manifest contains duplicate global_block_index") + if tuple(block[0] for block in ordered) != tuple(range(len(ordered))): + raise AttentionContractError("CP global_block_index values must be contiguous from zero") + previous_end = ordered[0][1] + for _, start, end, _, _ in ordered: + if start != previous_end: + raise AttentionContractError("CP block KV ranges must be gap-free and non-overlapping") + previous_end = end + return ordered + + +def _enum_member(values: Mapping[str, Any], value: Any, field: str) -> Any: + key = value.value if hasattr(value, "value") else str(value).lower() + if key not in values: + raise AttentionContractError(f"{field} must be one of {tuple(values)}, got {value!r}") + return values[key] + + +def _non_empty_string(value: Any, field: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise AttentionContractError(f"{field} must be a non-empty string") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer") + return value + + +def _bounded_int(value: Any, upper: int, field: str) -> int: + parsed = _non_negative_int(value, field) + if parsed >= upper: + raise AttentionContractError(f"{field}={parsed} must be smaller than {upper}") + return parsed + + +def _int_tuple(value: Any, field: str) -> tuple[int, ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise AttentionContractError(f"{field} must be a sequence of integers") + result = tuple(value) + if any(isinstance(item, bool) or not isinstance(item, int) for item in result): + raise AttentionContractError(f"{field} must be a sequence of integers") + return result + + +def _range_pair(value: Any, field: str) -> tuple[int, int]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)) or len(value) != 2: + raise AttentionContractError(f"{field} must be a (start, end) pair") + start, end = value + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError(f"{field} must satisfy 0 <= start < end") + return start, end + + +def _boundaries( + raw: Any, + expected_range: tuple[int, int], + entry_index: int, +) -> tuple[tuple[int, int], ...]: + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)) or not raw: + raise AttentionContractError( + f"split_kv.entries[{entry_index}] must report actual_split_boundaries" + ) + boundaries = tuple( + _range_pair(item, f"split_kv.entries[{entry_index}].actual_split_boundaries") + for item in raw + ) + previous_end = expected_range[0] + for start, end in boundaries: + if start != previous_end: + raise AttentionContractError( + "actual Split-KV boundaries must be gap-free in logical KV order" + ) + previous_end = end + if previous_end != expected_range[1]: + raise AttentionContractError( + "actual Split-KV boundaries must exactly cover expected_kv_range" + ) + return boundaries + + +def _mode(value: Any, field: str) -> str: + if value not in ("disabled", "fixed", "auto"): + raise AttentionContractError( + f"{field} must be 'disabled', 'fixed', or 'auto', got {value!r}" + ) + return value + + +def _validate_split_mode_sizes( + requested_mode: str, + requested_size: int | None, + actual_mode: str, + actual_size: int | None, + boundaries: tuple[tuple[int, int], ...], + fallback: bool, +) -> None: + if (requested_mode == "fixed") != (requested_size is not None): + raise AttentionContractError( + "requested fixed Split-KV mode and requested_split_kv_size must appear together" + ) + if (actual_mode == "fixed") != (actual_size is not None): + raise AttentionContractError( + "actual fixed Split-KV mode and actual_split_kv_size must appear together" + ) + if actual_mode == "disabled" and len(boundaries) != 1: + raise AttentionContractError("disabled Split-KV must report exactly one boundary") + if actual_mode == "fixed" and actual_size is not None: + widths = tuple(end - start for start, end in boundaries) + if any(width != actual_size for width in widths[:-1]) or widths[-1] > actual_size: + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_kv_size except at the tail" + ) + if not fallback and (requested_mode, requested_size) != (actual_mode, actual_size): + raise AttentionContractError( + "actual Split-KV mode/size may differ from the request only for a fallback" + ) + + +def _validate_owner_ranges( + entries: Sequence[tuple[Any, ...]], + batch_size: int, + tp_world_size: int, + cp_world_size: int, + totals: Sequence[int], +) -> None: + by_coordinate = {entry[:4]: entry for entry in entries} + for batch in range(batch_size): + for tp in range(tp_world_size): + for cp in range(cp_world_size): + previous_end = 0 + for owner in range(cp_world_size): + expected_range = by_coordinate[(batch, tp, cp, owner)][4] + if expected_range[0] != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = expected_range[1] + if previous_end != totals[batch]: + raise AttentionContractError( + "Split-KV owner ranges must cover total_kv_tokens" + ) + + +def _plan_semantics(entry: tuple[Any, ...]) -> tuple[Any, ...]: + # Drop TP/CP consumer coordinates and provenance labels. Within one side, + # every consumer of an owner range must execute the same numerical plan. + return (entry[0], entry[3], *entry[4:13], *entry[15:]) + + +def _cross_side_plan_semantics(entry: tuple[Any, ...]) -> tuple[Any, ...]: + # Backend and source are provenance, not numerical semantics. + return (*entry[:13], *entry[15:]) + + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "ATTENTION_LSE_EVIDENCE", + "ATTENTION_MERGE_STATE", + "CP_BLOCK_MANIFEST_EVIDENCE", + "CP_MERGE_REFERENCE", + "DOWNCAST_POINTS", + "POST_ROPE_QK_EVIDENCE", + "REFERENCE_CP_MERGE", + "SPLIT_KV_PLAN_EVIDENCE", + "SPLIT_KV_REFERENCE", + "TE_ROPE_REFERENCE", + "AttentionContractError", + "downcast_point", + "normalize_collective", + "normalize_cp_block_manifest", + "normalize_split_kv_plan_set", + "positive_int", + "precision", + "reference_collective", + "reference_selected", +] diff --git a/rl_engine/mismatch/operator_checks/attention/adapter.py b/rl_engine/mismatch/operator_checks/attention/adapter.py new file mode 100644 index 00000000..6f104314 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/adapter.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Attention's operator-level contract, readback and implementation adapter.""" + +from __future__ import annotations + +import importlib +from typing import Any, Callable, Mapping + +from rl_engine.mismatch.operator_checks.attention._common import ( + ATTENTION_LSE_DOMAIN, + ATTENTION_MERGE_STATE, + CP_MERGE_REFERENCE, + SPLIT_KV_REFERENCE, + AttentionContractError, + downcast_point, + normalize_collective, + normalize_cp_block_manifest, + normalize_split_kv_plan_set, + positive_int, + precision, + reference_collective, + reference_selected, +) +from rl_engine.mismatch.schema import ( + CollectiveContract, + ImplementationResolution, + OperatorContract, + PolicyRole, + Precision, + PrecisionProfile, + RejectedCandidate, +) + + +# One public error type for adapter and contract normalization failures. Callers +# should not have to know which validation layer rejected the runtime record. +AttentionAdapterError = AttentionContractError + + +def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: + """Build a contract from effective runtime state, never requested policy alone. + + Actual Split-KV plans, CP ownership and post-RoPE digests are intentionally + absent when the engine did not report them. Their factor rules then produce + ``REQUIRED_FIELD_MISSING`` instead of allowing a requested setting to pass as + runtime evidence. + """ + + if not isinstance(role, PolicyRole): + raise AttentionAdapterError(f"role must be a PolicyRole, got {role!r}") + if not isinstance(switch_values, Mapping): + raise AttentionAdapterError("Attention effective config must be a mapping") + + compute = precision(switch_values.get("attn.compute_dtype", "bf16"), "attn.compute_dtype") + accumulate = precision( + switch_values.get("attn.accumulate_dtype", "fp32"), "attn.accumulate_dtype" + ) + if accumulate is not Precision.FP32: + raise AttentionAdapterError( + "Attention softmax, Split-KV and CP (out, lse) merges must accumulate in fp32" + ) + downcast = downcast_point( + _role_value( + switch_values, + role, + common="attn.downcast_at", + training="attn.training_downcast_at", + rollout="attn.rollout_downcast_at", + default="final_write", + ), + "attn.downcast_at", + ) + + batch_size = positive_int(switch_values.get("attn.batch_size", 1), "attn.batch_size") + tp_world_size = positive_int( + switch_values.get("attn.tp_world_size", 1), "attn.tp_world_size" + ) + cp_world_size = positive_int( + switch_values.get("attn.cp_world_size", 1), "attn.cp_world_size" + ) + + split_reference = reference_selected( + switch_values.get("attn.split_kv"), role, SPLIT_KV_REFERENCE.name, "attn.split_kv" + ) + cp_reference = reference_selected( + switch_values.get("attn.cp_merge"), role, CP_MERGE_REFERENCE.name, "attn.cp_merge" + ) + + raw_collective = switch_values.get("attn.cp_collective") + if raw_collective is None and cp_reference: + collectives = reference_collective(cp_world_size) + else: + collectives = normalize_collective(raw_collective, cp_world_size=cp_world_size) + + plan = normalize_split_kv_plan_set( + switch_values.get("attn.actual_split_kv_plan_set") + ) + if plan is not None: + topology = (plan["batch_size"], plan["tp_world_size"], plan["cp_world_size"]) + expected = (batch_size, tp_world_size, cp_world_size) + if topology != expected: + raise AttentionAdapterError( + "Split-KV runtime plan topology does not match the Attention invocation: " + f"plan={topology}, attention={expected}" + ) + + manifest = normalize_cp_block_manifest( + switch_values.get("attn.cp_block_manifest"), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + + lse_domain = switch_values.get("attn.lse_domain") + export_lse = switch_values.get("attn.export_lse") + merge_state = switch_values.get("attn.merge_state") + if lse_domain is not None and lse_domain != ATTENTION_LSE_DOMAIN: + raise AttentionAdapterError( + f"attn.lse_domain must be {ATTENTION_LSE_DOMAIN!r}, got {lse_domain!r}" + ) + if export_lse is not None and not isinstance(export_lse, bool): + raise AttentionAdapterError("attn.export_lse must be a bool") + if merge_state is not None and merge_state != ATTENTION_MERGE_STATE: + raise AttentionAdapterError( + f"attn.merge_state must be {ATTENTION_MERGE_STATE!r}, got {merge_state!r}" + ) + + extra: dict[str, Any] = { + "batch_size": batch_size, + "tp_world_size": tp_world_size, + "cp_world_size": cp_world_size, + "requested_split_kv_policy": switch_values.get("attn.requested_split_kv_policy"), + "requested_split_kv_size": switch_values.get("attn.requested_split_kv_size"), + "split_kv_reference_selected": split_reference, + "cp_reference_selected": cp_reference, + "fusion_boundary": switch_values.get("attn.fusion_boundary"), + } + + _copy_if_present(extra, switch_values, "rope_theta", "attn.rope_theta") + _copy_if_present(extra, switch_values, "position_ids_digest", "attn.position_ids_digest") + _copy_if_present(extra, switch_values, "post_rope_qk_digest", "attn.post_rope_qk_digest") + _copy_if_present(extra, switch_values, "q_rope_state", "attn.q_rope_state") + _copy_if_present(extra, switch_values, "k_rope_state", "attn.k_rope_state") + _copy_if_present(extra, switch_values, "k_cache_rope_state", "attn.k_cache_rope_state") + + if plan is not None: + extra.update( + { + "total_kv_tokens": plan["total_kv_tokens"], + "split_kv_coordinates": plan["coordinates"], + "split_kv_owner_ranges": plan["owner_ranges"], + "split_kv_boundaries": plan["boundaries"], + "split_kv_merge_order": plan["merge_order"], + "split_kv_accumulate_precision": plan["accumulate_precision"], + "split_kv_downcast_at": plan["downcast_at"], + "split_kv_fallback": plan["fallback"], + "split_kv_runtime_plan_set": plan["canonical"], + "split_kv_backend": plan["backend"], + "split_kv_plan_source": plan["source"], + } + ) + if manifest is not None: + extra["cp_block_manifest"] = manifest + extra["cp_owner_ranges"] = tuple( + (block_index, start, end, owner_cp, owner_tp) + for block_index, start, end, owner_cp, owner_tp in manifest + ) + if lse_domain is not None: + extra["lse_domain"] = lse_domain + if export_lse is not None: + extra["export_lse"] = export_lse + if merge_state is not None: + extra["merge_state"] = merge_state + + return OperatorContract( + operator="attention", + role=role, + precision=PrecisionProfile( + compute=compute, + accumulate=accumulate, + softmax_accumulate=accumulate, + downcast_at=downcast, + ), + collectives=collectives, + extra=extra, + ) + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read the state that ran, rejecting requested-only configuration objects.""" + + adapter_role = getattr(adapter, "role", None) + if adapter_role is not None: + try: + normalized_role = PolicyRole(adapter_role) + except (TypeError, ValueError) as exc: + raise AttentionAdapterError(f"invalid adapter role {adapter_role!r}") from exc + if normalized_role is not role: + raise AttentionAdapterError( + f"adapter plays {normalized_role.value!r} but was queried as {role.value!r}" + ) + + reader = getattr(adapter, "read_effective_config", None) + if callable(reader): + value = reader() + elif isinstance(adapter, Mapping): + value = adapter + else: + value = getattr(adapter, "effective_config", None) + if not isinstance(value, Mapping): + raise AttentionAdapterError( + f"cannot read effective Attention config from {type(adapter).__name__}: expected " + "read_effective_config(), a mapping, or an effective_config mapping" + ) + + config = dict(value) + requested_only = config.pop("requested_config", None) + if requested_only is not None and not any(key.startswith("attn.") for key in config): + raise AttentionAdapterError( + "engine returned requested_config but no effective attn.* runtime readback" + ) + return config + + +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """Return the CP collective that actually ran, or no evidence when absent.""" + + config = read_effective_config(role, adapter) + cp_world_size = positive_int( + config.get("attn.cp_world_size", 1), "attn.cp_world_size" + ) + return normalize_collective(config.get("attn.cp_collective"), cp_world_size=cp_world_size) + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve every candidate in order and preserve every rejection reason.""" + + candidates = _implementation_candidates(factor_id, role, impl_name) + rejected: list[RejectedCandidate] = [] + for candidate in candidates: + parsed = _import_target(candidate) + if parsed is None: + rejected.append( + RejectedCandidate(name=candidate, reason="not a dotted or module:attribute path") + ) + continue + module_name, attribute = parsed + try: + module = importlib.import_module(module_name) + except (ImportError, OSError) as exc: + rejected.append( + RejectedCandidate(name=candidate, reason=f"import failed: {exc}") + ) + continue + + resolved: Any = module + for part in attribute.split("."): + resolved = getattr(resolved, part, None) + if resolved is None: + break + if resolved is None: + rejected.append( + RejectedCandidate( + name=candidate, + reason=f"{module_name} has no attribute {attribute!r}", + ) + ) + continue + if isinstance(resolved, type): + try: + resolved = resolved() + except Exception as exc: # noqa: BLE001 - recorded in provenance + rejected.append( + RejectedCandidate( + name=candidate, + reason=f"instantiation failed: {exc}", + ) + ) + continue + if not callable(resolved): + rejected.append( + RejectedCandidate(name=candidate, reason="resolved object is not callable") + ) + continue + return resolved, ImplementationResolution( + requested=impl_name, + resolved=candidate, + rejected=tuple(rejected), + ) + + return None, ImplementationResolution( + requested=impl_name, + resolved=None, + rejected=tuple(rejected), + ) + + +def _role_value( + values: Mapping[str, Any], + role: PolicyRole, + *, + common: str, + training: str, + rollout: str, + default: Any, +) -> Any: + role_key = training if role is PolicyRole.TRAINING else rollout + return values.get(role_key, values.get(common, default)) + + +def _copy_if_present( + target: dict[str, Any], source: Mapping[str, Any], target_key: str, source_key: str +) -> None: + if source_key in source: + target[target_key] = source[source_key] + + +def _implementation_candidates( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[str, ...]: + if factor_id == "attn.rope_fusion" and role is PolicyRole.ROLLOUT: + fallback = "vllm.model_executor.layers.rotary_embedding.get_rope" + return (impl_name, fallback) if impl_name != fallback else (impl_name,) + return (impl_name,) + + +def _import_target(value: str) -> tuple[str, str] | None: + if ":" in value: + module_name, attribute = value.rsplit(":", 1) + elif "." in value: + module_name, _, attribute = value.rpartition(".") + else: + return None + if not module_name or not attribute: + return None + return module_name, attribute + + +__all__ = [ + "AttentionAdapterError", + "build_contract", + "observe_collectives", + "read_effective_config", + "resolve_implementation", +] diff --git a/rl_engine/mismatch/operator_checks/attention/factors/__init__.py b/rl_engine/mismatch/operator_checks/attention/factors/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/rl_engine/mismatch/operator_checks/attention/factors/cp_merge.py b/rl_engine/mismatch/operator_checks/attention/factors/cp_merge.py new file mode 100644 index 00000000..34a1a6a4 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/factors/cp_merge.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""attn.cp_merge -- CP ownership, communication and (out, lse) merge order.""" + +from __future__ import annotations + +from rl_engine.mismatch.operator_checks.attention._common import ( + ATTENTION_LSE_EVIDENCE, + CP_BLOCK_MANIFEST_EVIDENCE, + CP_MERGE_REFERENCE, +) +from rl_engine.mismatch.schema import ( + COLLECTIVE_CONTRACT, + ComparisonRule, + Evidence, + FactorCategory, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +FACTOR = MismatchFactor( + id="attn.cp_merge", + operator="attention", + category=FactorCategory.SHARDING_AND_REDUCTION, + question=( + "Does CP Attention drift because block ownership, communication, or the FP32 " + "attention-domain (out, lse) merge order differs?" + ), + switch=Switch( + path="attn.cp_merge", + rebind_cost=RebindCost.PROCESS_GROUP_REBUILD, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=("native", CP_MERGE_REFERENCE.name), + ), + comparison_rules={ + "extra.tp_world_size": ComparisonRule.MUST_MATCH_BITWISE, + "extra.cp_world_size": ComparisonRule.MUST_MATCH_BITWISE, + "extra.cp_block_manifest": ComparisonRule.MUST_MATCH_BITWISE, + "extra.cp_owner_ranges": ComparisonRule.MUST_MATCH_BITWISE, + "extra.lse_domain": ComparisonRule.MUST_MATCH_BITWISE, + "extra.export_lse": ComparisonRule.MUST_MATCH_BITWISE, + "extra.merge_state": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].group": ComparisonRule.MUST_MATCH_BITWISE, + "collectives[0].group_size": ComparisonRule.MUST_MATCH_BITWISE, + "collectives[0].op": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].reduction_order": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].accumulate_precision": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].backend": ComparisonRule.RECORD_ONLY, + }, + prerequisites=Prerequisites( + required_ops=("p2p_nccl_attention_reference",), + min_gpu_count=2, + ), + required_evidence=( + Evidence.EFFECTIVE_CONFIG_READBACK.value, + COLLECTIVE_CONTRACT, + CP_BLOCK_MANIFEST_EVIDENCE, + ATTENTION_LSE_EVIDENCE, + ), + reference=CP_MERGE_REFERENCE, + pitfalls=( + KnownPitfall( + id="cp_arrival_order_merge", + mode=FailureMode.SILENT_FALSE_NEGATIVE, + symptom="the CP result is stable on one topology but moves on another", + actual_cause=( + "partial (out, lse) states were merged in arrival/NCCL order instead of " + "logical global block order" + ), + guard=( + "exchange an authoritative block manifest, sort by global_block_index, " + "then merge (out, lse) in fp32" + ), + guard_runs_at=NoiseFloor.SHARDED_SINGLE_NODE, + ), + ), +) diff --git a/rl_engine/mismatch/operator_checks/attention/factors/precision_downcast.py b/rl_engine/mismatch/operator_checks/attention/factors/precision_downcast.py new file mode 100644 index 00000000..0a0958cd --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/factors/precision_downcast.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""attn.precision_downcast -- compute precision and final write boundary.""" + +from __future__ import annotations + +from rl_engine.mismatch.schema import ( + ComparisonRule, + Evidence, + FactorCategory, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +FACTOR = MismatchFactor( + id="attn.precision_downcast", + operator="attention", + category=FactorCategory.OUTPUT_NUMERICS, + question=( + "Does Attention drift because the compute dtype differs or an FP32 partial is " + "downcast before the final output write?" + ), + switch=Switch( + path="attn.training_downcast_at", + rebind_cost=RebindCost.PER_REQUEST, + applies_to=(PolicyRole.TRAINING,), + allowed_values=("final_write", "per_partial", "per_block"), + ), + comparison_rules={ + "precision.compute": ComparisonRule.MUST_MATCH_BITWISE, + "precision.accumulate": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "precision.softmax_accumulate": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + }, + prerequisites=Prerequisites(required_ops=("attention",)), + required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value,), + pitfalls=( + KnownPitfall( + id="partial_state_downcast_hidden", + mode=FailureMode.SILENT_FALSE_NEGATIVE, + symptom="the final output dtype matches but CP/Split-KV drift remains", + actual_cause="one path rounded each partial before the online-softmax merge", + guard="capture partial out/lse dtypes and require exactly one final-write downcast", + guard_runs_at=NoiseFloor.SINGLE_LAYER_ANCHOR, + ), + ), +) diff --git a/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py new file mode 100644 index 00000000..a48f5de9 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""attn.rope_fusion -- RoPE fusion and tensor state.""" + +from __future__ import annotations + +from rl_engine.mismatch.operator_checks.attention._common import ( + POST_ROPE_QK_EVIDENCE, + TE_ROPE_REFERENCE, +) +from rl_engine.mismatch.schema import ( + POSITION_CACHE, + ComparisonRule, + Evidence, + FactorCategory, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +FACTOR = MismatchFactor( + id="attn.rope_fusion", + operator="attention", + category=FactorCategory.KERNEL_IMPLEMENTATION, + question=( + "Does the deviation come from fused vs small-operator vs sin/cos-cached " + "RoPE, or from position_ids / theta / the cast boundary?" + ), + switch=Switch( + path="attn.rope_fusion", + rebind_cost=RebindCost.ENGINE_REBUILD, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=("native", "transformer_engine"), + ), + comparison_rules={ + "extra.rope_theta": ComparisonRule.MUST_MATCH_BITWISE, + "extra.position_ids_digest": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.post_rope_qk_digest": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.q_rope_state": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.k_rope_state": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.k_cache_rope_state": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + # Fused vs unfused is what this factor ablates, so comparing it would + # fail every arm by construction. + "extra.fusion_boundary": ComparisonRule.RECORD_ONLY, + }, + prerequisites=Prerequisites( + required_ops=("rope",), + required_packages=("transformer_engine>=2.0",), + ), + required_evidence=( + Evidence.EFFECTIVE_CONFIG_READBACK.value, + POSITION_CACHE, + POST_ROPE_QK_EVIDENCE, + ), + reference=TE_ROPE_REFERENCE, + pitfalls=( + KnownPitfall( + id="rope_hook_not_covered", + mode=FailureMode.MISSING_INSTRUMENTATION, + symptom="RoPE looks perfectly consistent between the two sides", + actual_cause="the hook never attached, so nothing was captured at all", + guard="dump post-RoPE Q/K on both sides and compare bitwise before ablating", + guard_runs_at=NoiseFloor.SINGLE_LAYER_ANCHOR, + ), + ), +) diff --git a/rl_engine/mismatch/operator_checks/attention/factors/split_kv.py b/rl_engine/mismatch/operator_checks/attention/factors/split_kv.py new file mode 100644 index 00000000..e40aec55 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/factors/split_kv.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""attn.split_kv -- actual per-batch/TP/CP/owner Split-KV schedule.""" + +from __future__ import annotations + +from rl_engine.mismatch.operator_checks.attention._common import ( + SPLIT_KV_PLAN_EVIDENCE, + SPLIT_KV_REFERENCE, +) +from rl_engine.mismatch.schema import ( + BATCH_PLACEMENT, + ComparisonRule, + Evidence, + FactorCategory, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +FACTOR = MismatchFactor( + id="attn.split_kv", + operator="attention", + category=FactorCategory.SHARDING_AND_REDUCTION, + question=( + "Does Attention drift because train and rollout executed different logical " + "Split-KV boundaries or fallback schedules?" + ), + switch=Switch( + path="attn.split_kv", + rebind_cost=RebindCost.ENGINE_REBUILD, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=("native", SPLIT_KV_REFERENCE.name), + ), + comparison_rules={ + "extra.batch_size": ComparisonRule.MUST_MATCH_BITWISE, + "extra.tp_world_size": ComparisonRule.MUST_MATCH_BITWISE, + "extra.cp_world_size": ComparisonRule.MUST_MATCH_BITWISE, + "extra.total_kv_tokens": ComparisonRule.MUST_MATCH_BITWISE, + "extra.split_kv_coordinates": ComparisonRule.MUST_MATCH_BITWISE, + "extra.split_kv_owner_ranges": ComparisonRule.MUST_MATCH_BITWISE, + "extra.split_kv_boundaries": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.split_kv_merge_order": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.split_kv_accumulate_precision": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.split_kv_downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.split_kv_fallback": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.split_kv_runtime_plan_set": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.requested_split_kv_policy": ComparisonRule.RECORD_ONLY, + "extra.requested_split_kv_size": ComparisonRule.RECORD_ONLY, + "extra.split_kv_backend": ComparisonRule.RECORD_ONLY, + "extra.split_kv_plan_source": ComparisonRule.RECORD_ONLY, + }, + prerequisites=Prerequisites( + required_ops=("cp_attention",), + min_gpu_count=2, + ), + required_evidence=( + Evidence.EFFECTIVE_CONFIG_READBACK.value, + BATCH_PLACEMENT, + SPLIT_KV_PLAN_EVIDENCE, + ), + reference=SPLIT_KV_REFERENCE, + pitfalls=( + KnownPitfall( + id="requested_split_kv_is_not_execution", + mode=FailureMode.SILENT_FALSE_NEGATIVE, + symptom="matching Split-KV policy scalars make the case look clean", + actual_cause=( + "runtime shape selection or fallback produced different actual boundaries " + "on one batch/rank/owner" + ), + guard=( + "require the complete batch x TP x CP x owner runtime plan set and compare " + "actual boundaries" + ), + guard_runs_at=NoiseFloor.SHARDED_SINGLE_NODE, + ), + ), +) diff --git a/rl_engine/mismatch/operator_checks/gemm/__init__.py b/rl_engine/mismatch/operator_checks/gemm/__init__.py new file mode 100644 index 00000000..41d6ba31 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/gemm/__init__.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The gemm operator plugin.""" + +from rl_engine.mismatch.operator_checks.gemm import adapter +from rl_engine.mismatch.pipeline import OPERATOR_CHECKS, discover_factors + + +@OPERATOR_CHECKS.register +class GemmChecks: + operator = "gemm" + + def declare_factors(self): + return discover_factors(__package__) + + build_contract = staticmethod(adapter.build_contract) + read_effective_config = staticmethod(adapter.read_effective_config) + observe_collectives = staticmethod(adapter.observe_collectives) + resolve_implementation = staticmethod(adapter.resolve_implementation) diff --git a/rl_engine/mismatch/operator_checks/gemm/_common.py b/rl_engine/mismatch/operator_checks/gemm/_common.py new file mode 100644 index 00000000..754c4710 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/gemm/_common.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Collective contracts and the ordered reference, shared by gemm's factors.""" + +from __future__ import annotations + +from rl_engine.mismatch.schema import ( + CollectiveContract, + CollectiveOp, + DeterminismLevel, + DowncastPoint, + ExecutionPath, + LibraryPin, + ParallelDim, + Precision, + ReductionOrder, + ReferenceAuthority, + ReferenceImplementation, + RequiredSetting, + SettingChannel, +) + +TP_SIZE = 2 + +ORDERED_REDUCE_SCATTER = CollectiveContract( + op=CollectiveOp.REDUCE_SCATTER, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.GLOBAL_RANK_INDEX, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.STABLE_ACROSS_TOPOLOGY, + backend="rl_kernel", +) + +# Megatron with sequence parallelism on: all_reduce rewritten as +# reduce_scatter + all_gather, ordered by whatever NCCL picks. +NATIVE_TRAINING_REDUCE = CollectiveContract( + op=CollectiveOp.REDUCE_SCATTER, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.NCCL_ALGORITHM, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.NONE, + backend="nccl", +) + +# vLLM without sequence parallelism: a plain all_reduce, backend chosen at +# runtime from world size and topology. +NATIVE_ROLLOUT_REDUCE = CollectiveContract( + op=CollectiveOp.ALL_REDUCE, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.NCCL_ALGORITHM, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.PER_PARTIAL, + determinism=DeterminismLevel.NONE, + backend="vllm_custom_ipc", +) + +# SELF_WRITTEN because neither TE nor FlashInfer exposes a reduction whose order +# is fixed across topologies. +DETERMINISTIC_REDUCE_REFERENCE = ReferenceImplementation( + name="rl_kernel", + tier=ReferenceAuthority.SELF_WRITTEN, + training_impl="rl_engine.kernels.collectives.ordered_reduce_scatter", + rollout_impl="rl_engine.kernels.collectives.ordered_reduce_scatter", + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ), + required_settings=( + RequiredSetting( + "forward_reduce_contract", + ORDERED_REDUCE_SCATTER, + SettingChannel.CALL_ARG, + readback="module.last_collective_contract", + ), + RequiredSetting( + "NCCL_ALGO", + "Ring", + SettingChannel.ENV_VAR, + readback="os.environ", + guards="nccl_algo_unpinned", + ), + RequiredSetting( + "NCCL_PROTO", + "Simple", + SettingChannel.ENV_VAR, + readback="os.environ", + guards="nccl_algo_unpinned", + ), + ), + pinned_libraries=(LibraryPin("torch", "2.6.0"),), +) diff --git a/rl_engine/mismatch/operator_checks/gemm/adapter.py b/rl_engine/mismatch/operator_checks/gemm/adapter.py new file mode 100644 index 00000000..81b0a5c1 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/gemm/adapter.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""GEMM's four operator-level methods. Not implemented.""" + +from __future__ import annotations + +from typing import Any, Callable, Mapping + +from rl_engine.mismatch.schema import ( + CollectiveContract, + ImplementationResolution, + OperatorContract, + PolicyRole, +) + + +def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: + """Return this side's contract, with the collective the switch selects first. + + The factors index into ``collectives`` by position, so the order has to stay + stable. The two native sides differ by construction -- Megatron with sequence + parallelism rewrites all_reduce into reduce_scatter + all_gather, vLLM does a + plain all_reduce -- and that difference is the factor, not a defect. + """ + + raise NotImplementedError + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + raise NotImplementedError( + "Read actual values off the live engine, including which all-reduce " + "backend vLLM really chose." + ) + + +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """Return the collective trace for this run. + + vLLM switches between custom IPC, MNNVL and NCCL by world size and topology, + so only what ran is evidence. + """ + + raise NotImplementedError + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve an arm's implementation name, returning the trace even on failure. + + ``rl_engine.kernels.collectives.ordered_reduce_scatter`` does not exist yet, + so this must fail loudly enough for gate 1 to report + ``VARIANT_DID_NOT_APPLY`` rather than letting a silent fallback read as a + clean result. + """ + + raise NotImplementedError diff --git a/rl_engine/mismatch/operator_checks/gemm/factors/__init__.py b/rl_engine/mismatch/operator_checks/gemm/factors/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py b/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py new file mode 100644 index 00000000..3722080a --- /dev/null +++ b/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""gemm.forward_reduce -- RowParallel forward reduction order.""" + +from __future__ import annotations + +from rl_engine.mismatch.operator_checks.gemm._common import DETERMINISTIC_REDUCE_REFERENCE +from rl_engine.mismatch.schema import ( + COLLECTIVE_CONTRACT, + ComparisonRule, + Evidence, + ExpectedOutcome, + FactorCategory, + FactorVariant, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +_REF = DETERMINISTIC_REDUCE_REFERENCE + +# The standard four arms, spelled out so the self-check gate can also carry +# repeat_under: an implementation claiming topology independence must survive +# NCCL choosing a different algorithm. +_VARIANTS = ( + FactorVariant( + name="both_native", + switch_values={"gemm.forward_reduce": "native"}, + why="baseline: each side on its own framework's native reduction", + ), + FactorVariant( + name="both_reference", + switch_values={"gemm.forward_reduce": _REF.name}, + replace_on={ + PolicyRole.ROLLOUT: _REF.rollout_impl, + PolicyRole.TRAINING: _REF.training_impl, + }, + expected=ExpectedOutcome.BITWISE_IDENTICAL, + repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")}, + why="self-check gate, plus: a fixed order must survive a different NCCL algorithm", + ), + FactorVariant( + name="training_reference_only", + switch_values={"gemm.forward_reduce": f"{_REF.name}@training"}, + replace_on={PolicyRole.TRAINING: _REF.training_impl}, + why="swap the training side only: if the deviation goes, that side is the source", + ), + FactorVariant( + name="rollout_reference_only", + switch_values={"gemm.forward_reduce": f"{_REF.name}@rollout"}, + replace_on={PolicyRole.ROLLOUT: _REF.rollout_impl}, + why="swap the rollout side only: if the deviation goes, that side is the source", + ), +) + +FACTOR = MismatchFactor( + id="gemm.forward_reduce", + operator="gemm", + category=FactorCategory.SHARDING_AND_REDUCTION, + question=( + "Does the RowParallel forward reduction differ because sequence " + "parallelism rewrites all_reduce into reduce_scatter + all_gather?" + ), + switch=Switch( + path="gemm.forward_reduce", + rebind_cost=RebindCost.PROCESS_GROUP_REBUILD, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=("native", "rl_kernel"), + ), + comparison_rules={ + "collectives[0].op": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].reduction_order": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].accumulate_precision": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].group_size": ComparisonRule.MUST_MATCH_BITWISE, + # Two backends may legitimately differ; what must agree is the order. + "collectives[0].backend": ComparisonRule.RECORD_ONLY, + }, + prerequisites=Prerequisites(required_ops=("ordered_reduce_scatter",), min_gpu_count=2), + required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value, COLLECTIVE_CONTRACT), + reference=_REF, + # All three are row parallel linears eating the same accumulation order. + call_sites=("attention.o_linear", "mlp.down_linear", "moe.output"), + variants=_VARIANTS, + pitfalls=( + KnownPitfall( + id="nccl_algo_unpinned", + mode=FailureMode.SILENT_FALSE_NEGATIVE, + symptom="the reduction-order conclusion looks stable", + actual_cause="NCCL picks ring or tree per run, so the conclusion is noise", + guard="pin NCCL_ALGO/NCCL_PROTO and rerun; results must be bitwise identical", + guard_runs_at=NoiseFloor.SHARDED_SINGLE_NODE, + ), + ), +) diff --git a/rl_engine/mismatch/operator_checks/logprob/__init__.py b/rl_engine/mismatch/operator_checks/logprob/__init__.py new file mode 100644 index 00000000..1019ceb4 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/__init__.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The logprob operator plugin.""" + +from rl_engine.mismatch.operator_checks.logprob import adapter +from rl_engine.mismatch.pipeline import OPERATOR_CHECKS, discover_factors + + +@OPERATOR_CHECKS.register +class LogprobChecks: + operator = "logprob" + + def declare_factors(self): + return discover_factors(__package__) + + build_contract = staticmethod(adapter.build_contract) + read_effective_config = staticmethod(adapter.read_effective_config) + observe_collectives = staticmethod(adapter.observe_collectives) + resolve_implementation = staticmethod(adapter.resolve_implementation) diff --git a/rl_engine/mismatch/operator_checks/logprob/_common.py b/rl_engine/mismatch/operator_checks/logprob/_common.py new file mode 100644 index 00000000..c490c784 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/_common.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Values shared by logprob's factors: contracts, shard maps, the reference. + +The swap reference is WS2's deterministic vocab-parallel selected-logprob +(issue #241 PR3): per-shard (max, sumexp) partials merged in global vocab-shard +order, fp32 accumulation, one downcast at the final write. The sweep factors +need no reference and scan a parameter instead. +""" + +from __future__ import annotations + +from rl_engine.mismatch.schema import ( + CollectiveContract, + CollectiveOp, + DeterminismLevel, + DowncastPoint, + ExecutionPath, + LibraryPin, + ParallelDim, + Precision, + ReductionOrder, + ReferenceAuthority, + ReferenceImplementation, + RequiredSetting, + SettingChannel, +) + +# vLLM computes logits at the model dtype; Megatron can keep the head in fp32. +HEAD_DTYPES: dict[str, Precision] = { + "bf16": Precision.BF16, + "fp32": Precision.FP32, +} + +DOWNCAST_POINTS: dict[str, DowncastPoint] = { + "final_write": DowncastPoint.FINAL_WRITE, + "per_partial": DowncastPoint.PER_PARTIAL, +} + +TP_SIZE = 2 + +QWEN3_REAL_VOCAB = 151936 +QWEN3_PADDED_VOCAB = 152064 + + +def even_vocab_shard_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: + """The even split both frameworks produce when padding already divides evenly. + + MCore's ``VocabUtility`` and vLLM's ``_get_indices`` can still disagree once + per-shard padding rules differ, which is why the effective map is read back + per side rather than assumed equal. + """ + + shard = padded_vocab // tp_world_size + return tuple( + (rank * shard, padded_vocab if rank == tp_world_size - 1 else (rank + 1) * shard) + for rank in range(tp_world_size) + ) + + +# WS2 reference: all_gather the per-shard (max, sumexp) partials, then every +# rank merges them locally in vocab-shard-index order. The gather concatenates +# by rank index, so the merge order is fixed regardless of NCCL's choices. +REFERENCE_LSE_MERGE = CollectiveContract( + op=CollectiveOp.ALL_GATHER, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.STABLE_ACROSS_TOPOLOGY, + backend="rl_kernel", +) + +# Megatron's vocab-parallel cross entropy: all_reduce of the partial max and +# partial sumexp over the TP group, ordered by whatever NCCL picks. +NATIVE_TRAINING_LSE_MERGE = CollectiveContract( + op=CollectiveOp.ALL_REDUCE, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.NCCL_ALGORITHM, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.NONE, + backend="nccl", +) + +# vLLM: gather the full logits to one rank and reduce locally in one pass -- +# a different floating-point association than merging per-shard partials. +NATIVE_ROLLOUT_LSE_MERGE = CollectiveContract( + op=CollectiveOp.ALL_GATHER, + group=ParallelDim.TENSOR, + group_size=TP_SIZE, + reduction_order=ReductionOrder.NCCL_ALGORITHM, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.NONE, + backend="vllm_custom_ipc", +) + +# SELF_WRITTEN because neither TE nor FlashInfer offers a vocab-parallel +# selected-logprob whose partial-LSE merge order is fixed across topologies. +DETERMINISTIC_LSE_REFERENCE = ReferenceImplementation( + name="rl_kernel", + tier=ReferenceAuthority.SELF_WRITTEN, + training_impl="rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp", + rollout_impl="rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp", + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ), + required_settings=( + # The op takes its typed LogprobContract as a call argument and echoes + # the resolved reduction semantics back through dispatch provenance. + RequiredSetting( + "logp.reduction_contract", + REFERENCE_LSE_MERGE, + SettingChannel.CALL_ARG, + readback="dispatch.provenance['contract']['reduction']", + ), + ), + pinned_libraries=(LibraryPin("torch", "2.6.0"),), +) diff --git a/rl_engine/mismatch/operator_checks/logprob/adapter.py b/rl_engine/mismatch/operator_checks/logprob/adapter.py new file mode 100644 index 00000000..60cf9544 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/adapter.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Logprob's four operator-level methods.""" + +from __future__ import annotations + +import importlib +from dataclasses import replace +from typing import Any, Callable, Mapping + +from rl_engine.mismatch.operator_checks.logprob._common import ( + DETERMINISTIC_LSE_REFERENCE, + DOWNCAST_POINTS, + HEAD_DTYPES, + NATIVE_ROLLOUT_LSE_MERGE, + NATIVE_TRAINING_LSE_MERGE, + QWEN3_PADDED_VOCAB, + QWEN3_REAL_VOCAB, + REFERENCE_LSE_MERGE, + TP_SIZE, + even_vocab_shard_bounds, +) +from rl_engine.mismatch.schema import ( + CollectiveContract, + DowncastPoint, + ImplementationResolution, + OperatorContract, + PolicyRole, + Precision, + PrecisionProfile, + RejectedCandidate, + positive_int, +) + +# Both sides run the model at bf16; only the training head can deviate from it. +_MODEL_DTYPE = Precision.BF16 + + +class LogprobAdapterError(ValueError): + """A switch value or engine adapter this plugin cannot interpret.""" + + +def _merge_choice(value: Any, role: PolicyRole) -> str: + """Map a ``logp.lse_merge`` switch value to this side's implementation. + + ``@training`` / ``@rollout`` are the one-sided swap arms. + """ + + if value in (None, "native"): + return "native" + name = DETERMINISTIC_LSE_REFERENCE.name + if value == name: + return "reference" + if value == f"{name}@training": + return "reference" if role is PolicyRole.TRAINING else "native" + if value == f"{name}@rollout": + return "reference" if role is PolicyRole.ROLLOUT else "native" + raise LogprobAdapterError( + f"unknown logp.lse_merge value {value!r}; expected 'native', {name!r}, " + f"'{name}@training' or '{name}@rollout'" + ) + + +def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: + """Return this side's contract, with the vocab shard map under ``extra``. + + Only the training side can vary the head dtype -- vLLM computes logits at the + model dtype. Under TP, MCore's ``VocabUtility`` and vLLM's ``_get_indices`` + disagree on the shard boundaries for Qwen3, so comparing partial results + without recording both maps is meaningless. + """ + + tp_world_size = positive_int(switch_values.get("logp.tp_world_size", TP_SIZE)) + + head_key = switch_values.get("logp.head_dtype", "bf16") + if head_key not in HEAD_DTYPES: + raise LogprobAdapterError( + f"unknown logp.head_dtype value {head_key!r}; expected one of {tuple(HEAD_DTYPES)}" + ) + downcast_key = switch_values.get("logp.downcast_at", "final_write") + if downcast_key not in DOWNCAST_POINTS: + raise LogprobAdapterError( + f"unknown logp.downcast_at value {downcast_key!r}; " + f"expected one of {tuple(DOWNCAST_POINTS)}" + ) + if role is PolicyRole.TRAINING: + lm_head = HEAD_DTYPES[head_key] + downcast_at = DOWNCAST_POINTS[downcast_key] + else: + lm_head = _MODEL_DTYPE + downcast_at = DowncastPoint.FINAL_WRITE + + merge = _merge_choice(switch_values.get("logp.lse_merge"), role) + collectives: tuple[CollectiveContract, ...] + if tp_world_size == 1: + collectives = () + elif merge == "reference": + collectives = (replace(REFERENCE_LSE_MERGE, group_size=tp_world_size),) + elif role is PolicyRole.TRAINING: + collectives = (replace(NATIVE_TRAINING_LSE_MERGE, group_size=tp_world_size),) + else: + collectives = (replace(NATIVE_ROLLOUT_LSE_MERGE, group_size=tp_world_size),) + + return OperatorContract( + operator="logprob", + role=role, + precision=PrecisionProfile( + compute=_MODEL_DTYPE, + accumulate=Precision.FP32, + downcast_at=downcast_at, + lm_head=lm_head, + ), + collectives=collectives, + extra={ + "vocab_shard_map": even_vocab_shard_bounds(QWEN3_PADDED_VOCAB, tp_world_size), + "real_vocab_size": QWEN3_REAL_VOCAB, + "padded_vocab_size": QWEN3_PADDED_VOCAB, + "logprobs_mode": ( + "vocab_parallel_cross_entropy" if role is PolicyRole.TRAINING else "raw_logprobs" + ), + "lse_export": merge == "reference", + }, + ) + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read the head dtype, transform chain and shard boundaries off the engine. + + Accepts an engine backend exposing ``read_effective_config()``, a plain + mapping already read back, or an object carrying ``effective_config``. What + comes back is actual state, never the requested switch values. + """ + + adapter_role = getattr(adapter, "role", None) + if adapter_role is not None and adapter_role is not role: + raise LogprobAdapterError( + f"adapter plays {adapter_role.value!r} but was queried as {role.value!r}" + ) + + reader = getattr(adapter, "read_effective_config", None) + if callable(reader): + return dict(reader()) + if isinstance(adapter, Mapping): + return dict(adapter) + config = getattr(adapter, "effective_config", None) + if config is not None: + return dict(config) + raise LogprobAdapterError( + f"cannot read an effective config off {type(adapter).__name__}: expected " + "read_effective_config(), a mapping, or an effective_config attribute" + ) + + +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """Return the collective trace: empty at TP=1, the partial-LSE merge otherwise. + + Which contract applies is decided by the effective config read off the + engine, not by what was requested. + """ + + config = read_effective_config(role, adapter) + tp_world_size = positive_int(config.get("logp.tp_world_size", 1)) + if tp_world_size == 1: + return () + merge = _merge_choice(config.get("logp.lse_merge"), role) + if merge == "reference": + return (replace(REFERENCE_LSE_MERGE, group_size=tp_world_size),) + if role is PolicyRole.TRAINING: + return (replace(NATIVE_TRAINING_LSE_MERGE, group_size=tp_world_size),) + return (replace(NATIVE_ROLLOUT_LSE_MERGE, group_size=tp_world_size),) + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve an arm's dotted import path, returning the trace even on failure. + + The WS2 vocab-parallel reference ships behind issue #241 PR3, so on a tree + without it the rejection record is the finding: ``FELL_BACK`` with the + import error, not a silent ``None``. + """ + + rejected: list[RejectedCandidate] = [] + if "." not in impl_name: + rejected.append(RejectedCandidate(name=impl_name, reason="not a dotted import path")) + return None, ImplementationResolution( + requested=impl_name, resolved=None, rejected=tuple(rejected) + ) + + module_name, _, attribute = impl_name.rpartition(".") + try: + module = importlib.import_module(module_name) + except ImportError as exc: + rejected.append(RejectedCandidate(name=impl_name, reason=f"import failed: {exc}")) + return None, ImplementationResolution( + requested=impl_name, resolved=None, rejected=tuple(rejected) + ) + + resolved = getattr(module, attribute, None) + if resolved is None: + rejected.append( + RejectedCandidate( + name=impl_name, reason=f"{module_name} has no attribute {attribute!r}" + ) + ) + return None, ImplementationResolution( + requested=impl_name, resolved=None, rejected=tuple(rejected) + ) + if isinstance(resolved, type): + try: + resolved = resolved() + except Exception as exc: # noqa: BLE001 - the reason goes into the trace + rejected.append( + RejectedCandidate(name=impl_name, reason=f"instantiation failed: {exc}") + ) + return None, ImplementationResolution( + requested=impl_name, resolved=None, rejected=tuple(rejected) + ) + if not callable(resolved): + rejected.append(RejectedCandidate(name=impl_name, reason="resolved object is not callable")) + return None, ImplementationResolution( + requested=impl_name, resolved=None, rejected=tuple(rejected) + ) + + return resolved, ImplementationResolution( + requested=impl_name, resolved=impl_name, rejected=tuple(rejected) + ) diff --git a/rl_engine/mismatch/operator_checks/logprob/factors/__init__.py b/rl_engine/mismatch/operator_checks/logprob/factors/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/rl_engine/mismatch/operator_checks/logprob/factors/lse_merge_order.py b/rl_engine/mismatch/operator_checks/logprob/factors/lse_merge_order.py new file mode 100644 index 00000000..d5b1da35 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/factors/lse_merge_order.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""logp.lse_merge_order -- how the per-shard LSE partials are merged under TP. + +An implementation swap: training all_reduces (max, sumexp) partials in NCCL +order while rollout gathers full logits and reduces locally, and the reference +replaces both with WS2's fixed vocab-shard-order merge. +""" + +from __future__ import annotations + +from rl_engine.mismatch.operator_checks.logprob._common import DETERMINISTIC_LSE_REFERENCE +from rl_engine.mismatch.schema import ( + COLLECTIVE_CONTRACT, + LSE_EXPORT, + VOCAB_SHARD_MAP, + ComparisonRule, + Evidence, + ExpectedOutcome, + FactorCategory, + FactorVariant, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +_REF = DETERMINISTIC_LSE_REFERENCE + +_VARIANTS = ( + FactorVariant( + name="both_native", + switch_values={"logp.lse_merge": "native"}, + why="baseline: each side merges its LSE partials the way its framework does", + ), + FactorVariant( + name="both_reference", + switch_values={"logp.lse_merge": _REF.name}, + replace_on={ + PolicyRole.ROLLOUT: _REF.rollout_impl, + PolicyRole.TRAINING: _REF.training_impl, + }, + expected=ExpectedOutcome.BITWISE_IDENTICAL, + repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")}, + why="self-check gate: a shard-order merge must survive a different NCCL algorithm", + ), + FactorVariant( + name="training_reference_only", + switch_values={"logp.lse_merge": f"{_REF.name}@training"}, + replace_on={PolicyRole.TRAINING: _REF.training_impl}, + why="swap the training side only: if the deviation goes, that side is the source", + ), + FactorVariant( + name="rollout_reference_only", + switch_values={"logp.lse_merge": f"{_REF.name}@rollout"}, + replace_on={PolicyRole.ROLLOUT: _REF.rollout_impl}, + why="swap the rollout side only: if the deviation goes, that side is the source", + ), +) + +FACTOR = MismatchFactor( + id="logp.lse_merge_order", + operator="logprob", + category=FactorCategory.SHARDING_AND_REDUCTION, + question=( + "Does the logprob deviation come from the two sides merging their " + "partial-LSE shards in different floating-point orders under TP?" + ), + switch=Switch( + path="logp.lse_merge", + rebind_cost=RebindCost.PROCESS_GROUP_REBUILD, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=("native", "rl_kernel"), + ), + comparison_rules={ + # Same tiers as gemm.forward_reduce for the shared collective paths; + # the registry rejects one path declared at two different tiers. + "collectives[0].op": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].reduction_order": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].accumulate_precision": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "collectives[0].group_size": ComparisonRule.MUST_MATCH_BITWISE, + "collectives[0].backend": ComparisonRule.RECORD_ONLY, + # Disagreeing shard maps make partial comparison meaningless: void, not + # a finding. + "extra.vocab_shard_map": ComparisonRule.MUST_MATCH_BITWISE, + "extra.lse_export": ComparisonRule.RECORD_ONLY, + }, + prerequisites=Prerequisites(required_ops=("vocab_parallel_logp",), min_gpu_count=2), + required_evidence=( + Evidence.EFFECTIVE_CONFIG_READBACK.value, + COLLECTIVE_CONTRACT, + VOCAB_SHARD_MAP, + LSE_EXPORT, + ), + reference=_REF, + variants=_VARIANTS, + pitfalls=( + KnownPitfall( + id="padded_vocab_in_lse", + mode=FailureMode.SILENT_FALSE_NEGATIVE, + symptom="a one-sided dlogp bias on every token, blamed on the merge order", + actual_cause=( + "padded vocab columns leak into one side's local sumexp, inflating " + "its LSE denominator -- the merge order was never the problem" + ), + guard="assert exp(logp) sums to 1 over the real vocabulary on one anchor token", + guard_runs_at=NoiseFloor.SINGLE_LAYER_ANCHOR, + ), + ), +) diff --git a/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py b/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py new file mode 100644 index 00000000..2c2abfa1 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""logp.precision_downcast -- lm_head dtype and where fp32 is written back. + +A parameter sweep: nothing is replaced, the head dtype is scanned. That is what +``reference=None`` says. +""" + +from __future__ import annotations + +from rl_engine.mismatch.operator_checks.logprob._common import HEAD_DTYPES +from rl_engine.mismatch.schema import ( + MODEL_SHAPE, + ComparisonRule, + Evidence, + FactorCategory, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +FACTOR = MismatchFactor( + id="logp.precision_downcast", + operator="logprob", + category=FactorCategory.OUTPUT_NUMERICS, + question=( + "Is the deviation the lm_head GEMM running at the model dtype, or where " + "the fp32 accumulator is written back?" + ), + switch=Switch( + path="logp.head_dtype", + rebind_cost=RebindCost.PER_REQUEST, + # vLLM computes logits at the model dtype, so only training can vary. + applies_to=(PolicyRole.TRAINING,), + allowed_values=tuple(HEAD_DTYPES), + ), + comparison_rules={ + "precision.lm_head": ComparisonRule.MUST_MATCH_BITWISE, + "precision.accumulate": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "extra.logprobs_mode": ComparisonRule.RECORD_ONLY, + }, + prerequisites=Prerequisites(required_ops=("lm_head",)), + required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value, MODEL_SHAPE), + pitfalls=( + KnownPitfall( + id="head_dtype_tail_only", + mode=FailureMode.SILENT_FALSE_NEGATIVE, + symptom="dlogp_mean sits far below the clip edge, so the run reads as clean", + actual_cause=( + "a bf16 head is a rounding error on most tokens and a large one on " + "the few with the flattest distribution -- the damage is all tail" + ), + guard="judge this factor on clip_fraction and dlogp_p99, never on the mean", + guard_runs_at=NoiseFloor.SINGLE_LAYER_ANCHOR, + ), + ), +) diff --git a/rl_engine/mismatch/pipeline/__init__.py b/rl_engine/mismatch/pipeline/__init__.py new file mode 100644 index 00000000..128cd99a --- /dev/null +++ b/rl_engine/mismatch/pipeline/__init__.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The execution pipeline: free functions only, no state. + +One file per step, in the order they run:: + + registry plugin registration and factor discovery + planner filter, statically reject, expand variants, order by cost + runner execution loop, reuse decisions, variant repeats + diagnosis four gates plus the matrix + report filter false positives, trace root causes, emit the report + +``comparison`` holds the declaration-driven contract comparison the runner uses. +""" + +from rl_engine.mismatch.pipeline.comparison import compare_contracts, resolve_field_path +from rl_engine.mismatch.pipeline.diagnosis import CONVERGENCE_RATIO, diagnose +from rl_engine.mismatch.pipeline.planner import ( + ContradictoryFactor, + UnmetPrerequisite, + build_variants, + missing_prerequisites, + order_cases_by_rebind_cost, + reject_contradictory_factors, + suggested_floor_is_lowest, +) +from rl_engine.mismatch.pipeline.registry import ( + OPERATOR_CHECKS, + FactorDiscoveryError, + OperatorChecks, + PluginRegistry, + RegistrationError, + discover_factors, +) +from rl_engine.mismatch.pipeline.report import ( + build_report, + filter_known_equivalences, + render_summary, + trace_root_causes, +) +from rl_engine.mismatch.pipeline.runner import ( + ReadOnlyViolation, + RunContext, + ScoringBackend, + assert_comparison_is_read_only, + assert_order_is_topology_independent, + compute_metrics, + expand_repeats, + run_variant, +) + +__all__ = [ + "compare_contracts", + "resolve_field_path", + "CONVERGENCE_RATIO", + "diagnose", + "ContradictoryFactor", + "UnmetPrerequisite", + "build_variants", + "missing_prerequisites", + "order_cases_by_rebind_cost", + "reject_contradictory_factors", + "suggested_floor_is_lowest", + "OPERATOR_CHECKS", + "FactorDiscoveryError", + "OperatorChecks", + "PluginRegistry", + "RegistrationError", + "discover_factors", + "build_report", + "filter_known_equivalences", + "render_summary", + "trace_root_causes", + "ReadOnlyViolation", + "RunContext", + "ScoringBackend", + "assert_comparison_is_read_only", + "assert_order_is_topology_independent", + "compute_metrics", + "expand_repeats", + "run_variant", +] diff --git a/rl_engine/mismatch/pipeline/comparison.py b/rl_engine/mismatch/pipeline/comparison.py new file mode 100644 index 00000000..6779fa39 --- /dev/null +++ b/rl_engine/mismatch/pipeline/comparison.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Contract comparison, driven entirely by declarations. + +The loop is generic; the per-operator part is putting fields in the right place +inside ``build_contract()``. There are no operator branches here. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import fields, is_dataclass +from typing import Any, Mapping, Sequence + +from rl_engine.mismatch.schema import ( + ComparisonIssue, + ComparisonIssueCode, + ComparisonRule, + DeterminismLevel, + MismatchFactor, + OperatorContract, + PolicyRole, +) + +_MISSING = object() +_INDEX = re.compile(r"^(?P[^\[]+)\[(?P\d+)\]$") + + +def resolve_field_path(contract: OperatorContract, path: str) -> Any: + """Index a contract by a dotted path such as ``collectives[0].reduction_order``. + + Absent paths return a sentinel, so a missing field is reported not raised. + """ + + current: Any = contract + for part in path.split("."): + matched = _INDEX.match(part) + index = None + if matched: + part = matched.group("name") + index = int(matched.group("index")) + + if isinstance(current, Mapping): + if part not in current: + return _MISSING + current = current[part] + elif is_dataclass(current) and any(f.name == part for f in fields(current)): + current = getattr(current, part) + elif hasattr(current, part): + current = getattr(current, part) + else: + return _MISSING + + if index is not None: + if not isinstance(current, Sequence) or index >= len(current): + return _MISSING + current = current[index] + + return current + + +def _values_equal(left: Any, right: Any, *, bitwise: bool) -> bool: + if isinstance(left, float) and isinstance(right, float): + if math.isnan(left) and math.isnan(right): + return True + return left == right if bitwise else math.isclose(left, right, rel_tol=0.0, abs_tol=0.0) + return left == right + + +def compare_contracts( + rollout: OperatorContract, + training: OperatorContract, + factors: Sequence[MismatchFactor], +) -> tuple[ComparisonIssue, ...]: + """Compare the two contracts field by field, per ``comparison_rules``. + + ``RECORD_ONLY`` fields are never compared: that tier exists so structural + differences like packed QKV do not drown the real problems. + """ + + rules: dict[str, ComparisonRule] = {} + for factor in factors: + rules.update(factor.comparison_rules) + + issues: list[ComparisonIssue] = [] + for path, rule in sorted(rules.items()): + if rule is ComparisonRule.RECORD_ONLY: + continue + + left = resolve_field_path(rollout, path) + right = resolve_field_path(training, path) + + if left is _MISSING or right is _MISSING: + issues.append( + ComparisonIssue( + code=ComparisonIssueCode.REQUIRED_FIELD_MISSING, + rule=rule, + field_path=path, + values={ + PolicyRole.ROLLOUT: None if left is _MISSING else left, + PolicyRole.TRAINING: None if right is _MISSING else right, + }, + message=( + f"{path!r} is declared {rule.value!r} but is absent from " + f"{'rollout' if left is _MISSING else 'training'}'s contract" + ), + ) + ) + continue + + bitwise = rule is ComparisonRule.MUST_MATCH_BITWISE + if _values_equal(left, right, bitwise=bitwise): + continue + + code = ( + ComparisonIssueCode.BITWISE_MISMATCH + if bitwise + else ComparisonIssueCode.SEMANTIC_MISMATCH + ) + issues.append( + ComparisonIssue( + code=code, + rule=rule, + field_path=path, + values={PolicyRole.ROLLOUT: left, PolicyRole.TRAINING: right}, + message=f"{path!r} differs: rollout={left!r} training={right!r}", + ) + ) + + issues.extend(_determinism_issues(rollout, training)) + return tuple(issues) + + +def _determinism_issues( + rollout: OperatorContract, training: OperatorContract +) -> tuple[ComparisonIssue, ...]: + """One side claiming a stronger reproducibility guarantee than the other. + + Comparing against an implementation that is not reproducible across runs + measures the weaker side's noise, not the gap between the two. + """ + + strength = { + DeterminismLevel.NONE: 0, + DeterminismLevel.STABLE_WITHIN_PROCESS: 1, + DeterminismLevel.STABLE_ACROSS_RUNS: 2, + DeterminismLevel.STABLE_ACROSS_TOPOLOGY: 3, + } + issues: list[ComparisonIssue] = [] + # strict=False: the two sides may declare different numbers of collectives. + paired = zip(rollout.collectives, training.collectives, strict=False) + for index, (left, right) in enumerate(paired): + if strength[left.determinism] != strength[right.determinism]: + issues.append( + ComparisonIssue( + code=ComparisonIssueCode.DETERMINISM_INCOMPATIBLE, + rule=ComparisonRule.MUST_MATCH_SEMANTICALLY, + field_path=f"collectives[{index}].determinism", + values={ + PolicyRole.ROLLOUT: left.determinism, + PolicyRole.TRAINING: right.determinism, + }, + message=( + f"collectives[{index}]: rollout guarantees " + f"{left.determinism.value!r} while training guarantees " + f"{right.determinism.value!r}" + ), + ) + ) + return tuple(issues) + + +__all__ = ["compare_contracts", "resolve_field_path"] diff --git a/rl_engine/mismatch/pipeline/diagnosis.py b/rl_engine/mismatch/pipeline/diagnosis.py new file mode 100644 index 00000000..9650ae47 --- /dev/null +++ b/rl_engine/mismatch/pipeline/diagnosis.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Four gates plus the diagnosis matrix. + +The framework draws the conclusion; nobody reads the numbers by hand. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +from rl_engine.mismatch.schema import ( + Diagnosis, + ExpectedOutcome, + FactorReport, + KnownPitfall, + MismatchFactor, + MismatchMetrics, + NoiseFloor, + SwitchStatus, + VariantResult, + is_silent_failure, + missing_evidence, + tolerance_floor, +) + +CONVERGENCE_RATIO = 0.1 + + +@dataclass(frozen=True) +class _Outcome: + diagnosis: Diagnosis + reason: str + + +def diagnose( + factor: MismatchFactor, + variants: Sequence[VariantResult], + *, + noise_floor: NoiseFloor, + model_family: str = "dense", + failed_guards: Sequence[KnownPitfall] = (), +) -> FactorReport: + """Run four gates, then the matrix. + + "Not measured" and "measured and clean" are different things, and confusing + them is the mistake an attribution framework is most likely to make. Nothing + reaches the matrix without passing the gates. + """ + + outcome = ( + _gate_variants_applied(variants) + or _gate_evidence_complete(factor, variants) + or _gate_shards_complete(variants) + or _gate_guards_passed(failed_guards) + or _run_matrix(variants, noise_floor=noise_floor, model_family=model_family) + ) + + return FactorReport( + factor_id=factor.id, + noise_floor=noise_floor, + variants=tuple(variants), + diagnosis=outcome.diagnosis, + diagnosis_reason=outcome.reason, + ) + + +def _gate_variants_applied(variants: Sequence[VariantResult]) -> _Outcome | None: + """Gate 1: did every variant actually take effect? + + ``FELL_BACK`` is the dangerous one: the engine silently reverted to native, + and "the deviation did not change" then reads as ``NOT_THIS_FACTOR`` -- a + false negative that looks exactly like a clean result. + """ + + for result in variants: + if result.status is SwitchStatus.APPLIED: + continue + detail = "" + if result.resolution is not None: + rejected = ", ".join( + f"{candidate.name} ({candidate.reason})" for candidate in result.resolution.rejected + ) + detail = f"; tried: {rejected or 'nothing'}" + return _Outcome( + Diagnosis.VARIANT_DID_NOT_APPLY, + f"variant {result.variant.name!r} is {result.status.value!r}{detail}", + ) + return None + + +def _gate_evidence_complete( + factor: MismatchFactor, variants: Sequence[VariantResult] +) -> _Outcome | None: + """Gate 2: is the required evidence present?""" + + for result in variants: + absent = missing_evidence(result.evidence, factor.required_evidence) + if absent: + return _Outcome( + Diagnosis.INSUFFICIENT_EVIDENCE, + f"variant {result.variant.name!r} is missing evidence: {sorted(absent)}", + ) + return None + + +def _gate_shards_complete(variants: Sequence[VariantResult]) -> _Outcome | None: + """Gate 3: were all logprob shards collected? + + One vocab shard short and the LSE denominator loses a chunk, so logp comes + out systematically high with nothing to show for it. + """ + + for result in variants: + shards = result.logprob_shards + if not shards: + continue + world_size = shards[0].world_size + if len({shard.rank for shard in shards}) != world_size: + return _Outcome( + Diagnosis.INSUFFICIENT_EVIDENCE, + f"variant {result.variant.name!r} collected {len(shards)} of " + f"{world_size} logprob shards", + ) + return None + + +def _gate_guards_passed(failed_guards: Sequence[KnownPitfall]) -> _Outcome | None: + """Gate 4: did every pitfall guard pass?""" + + if failed_guards: + names = ", ".join(guard.id for guard in failed_guards) + return _Outcome(Diagnosis.INSUFFICIENT_EVIDENCE, f"pitfall guards did not pass: {names}") + return None + + +def _run_matrix( + variants: Sequence[VariantResult], + *, + noise_floor: NoiseFloor, + model_family: str, +) -> _Outcome: + """The matrix proper, once all four gates are clear.""" + + by_name = {result.variant.name: result for result in variants} + + baseline = by_name.get("both_native") + if baseline is None or baseline.metrics is None: + return _Outcome( + Diagnosis.INSUFFICIENT_EVIDENCE, "no both_native baseline to compare against" + ) + + self_check = by_name.get("both_reference") + if self_check is not None: + if not _is_bitwise_identical(self_check): + return _Outcome( + Diagnosis.REFERENCE_ITSELF_IS_BROKEN, + "both_reference is not bitwise identical; fix the reference before " + "trusting any conclusion from this factor", + ) + + training_only = by_name.get("training_reference_only") + rollout_only = by_name.get("rollout_reference_only") + if training_only is None or rollout_only is None: + return _Outcome( + Diagnosis.INSUFFICIENT_EVIDENCE, + "one-sided variants are missing; a two-sided swap alone cannot attribute a side", + ) + + floor = tolerance_floor(model_family, noise_floor) + training_converged = _converged(baseline.metrics, training_only.metrics, floor) + rollout_converged = _converged(baseline.metrics, rollout_only.metrics, floor) + + if training_converged and not rollout_converged: + return _Outcome(Diagnosis.CAUSED_BY_TRAINING_SIDE, "only the training-side swap converged") + if rollout_converged and not training_converged: + return _Outcome(Diagnosis.CAUSED_BY_ROLLOUT_SIDE, "only the rollout-side swap converged") + if training_converged and rollout_converged: + return _Outcome( + Diagnosis.CAUSED_BY_BOTH_SIDES, + "both one-sided swaps converged; the reference is the only anchor", + ) + + # Neither converged. Check the tail before calling it clean: the mean stays + # far below the clip edge at every production floor. + for candidate in (training_only, rollout_only): + if candidate.metrics is not None and is_silent_failure(candidate.metrics): + return _Outcome( + Diagnosis.COUPLED_WITH_OTHER_FACTORS, + "neither swap converged and the tail is past the clip edge; " + "this factor interacts with another", + ) + + return _Outcome(Diagnosis.NOT_THIS_FACTOR, "neither one-sided swap moved the deviation") + + +def _is_bitwise_identical(result: VariantResult) -> bool: + if result.variant.expected is not ExpectedOutcome.BITWISE_IDENTICAL: + return True + if result.metrics is None: + return False + return result.metrics.dlogp_max == 0.0 + + +def _converged( + baseline: MismatchMetrics, candidate: MismatchMetrics | None, tol_floor: float +) -> bool: + """Convergence is judged on ``clip_fraction``, not ``dlogp_mean``. + + At the production floor the mean is always far below the clip edge, so using + it would call nearly every factor ``NOT_THIS_FACTOR``. + """ + + if candidate is None: + return False + if baseline.clip_fraction > 0.0: + return candidate.clip_fraction <= CONVERGENCE_RATIO * baseline.clip_fraction + return candidate.dlogp_mean <= max(CONVERGENCE_RATIO * baseline.dlogp_mean, tol_floor) + + +__all__ = ["CONVERGENCE_RATIO", "diagnose"] diff --git a/rl_engine/mismatch/pipeline/planner.py b/rl_engine/mismatch/pipeline/planner.py new file mode 100644 index 00000000..3cf26297 --- /dev/null +++ b/rl_engine/mismatch/pipeline/planner.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Planning: filter, statically reject, expand variants, order by rebuild cost.""" + +from __future__ import annotations + +import importlib.util +from dataclasses import dataclass +from typing import Sequence + +from rl_engine.mismatch.schema import ( + DeterminismLevel, + ExpectedOutcome, + FactorVariant, + MismatchFactor, + NoiseFloor, + PolicyRole, + RebindCost, + ReductionOrder, + declared_collectives, + requires_fixed_order, +) + +_NON_DETERMINISTIC_ORDERS = (ReductionOrder.NCCL_ALGORITHM, ReductionOrder.ARRIVAL) + +_REBIND_ORDER = { + RebindCost.PER_REQUEST: 0, + RebindCost.ENGINE_REBUILD: 1, + RebindCost.PROCESS_GROUP_REBUILD: 2, + RebindCost.PROCESS_RESTART: 3, +} + + +class ContradictoryFactor(ValueError): + """A factor's declaration contradicts itself; rejected before anything runs.""" + + +@dataclass(frozen=True) +class UnmetPrerequisite: + """One reason a factor cannot run today.""" + + factor_id: str + reason: str + + +def reject_contradictory_factors(factors: Sequence[MismatchFactor]) -> None: + """Reject a factor whose declaration contradicts itself. Nothing executes. + + Claiming topology independence while reducing non-deterministically produces + numbers that mean nothing, and that is knowable before running. + """ + + for factor in factors: + for contract in declared_collectives(factor): + if requires_fixed_order(contract) and contract.reduction_order in ( + _NON_DETERMINISTIC_ORDERS + ): + raise ContradictoryFactor( + f"{factor.id}: claims {DeterminismLevel.STABLE_ACROSS_TOPOLOGY.value} " + f"but reduces with {contract.reduction_order.value}" + ) + + +def missing_prerequisites( + factor: MismatchFactor, + *, + available_ops: frozenset[str] = frozenset(), + gpu_count: int = 0, + model_traits: frozenset[str] = frozenset(), +) -> tuple[UnmetPrerequisite, ...]: + """What this factor is still missing: operators, devices, packages, traits.""" + + needs = factor.prerequisites + unmet: list[UnmetPrerequisite] = [] + + for op in needs.required_ops: + if op not in available_ops: + unmet.append(UnmetPrerequisite(factor.id, f"operator {op!r} is not dispatchable")) + + if gpu_count < needs.min_gpu_count: + unmet.append( + UnmetPrerequisite( + factor.id, f"needs {needs.min_gpu_count} devices, {gpu_count} available" + ) + ) + + for requirement in needs.required_packages: + package = requirement.split(">")[0].split("=")[0].split("<")[0].strip() + if importlib.util.find_spec(package.replace("-", "_")) is None: + unmet.append(UnmetPrerequisite(factor.id, f"package {requirement!r} is not installed")) + + for trait in needs.required_model_traits: + if trait not in model_traits: + unmet.append(UnmetPrerequisite(factor.id, f"model does not have trait {trait!r}")) + + for blocker in needs.blocked_by: + unmet.append(UnmetPrerequisite(factor.id, f"blocked by {blocker}")) + + return tuple(unmet) + + +def build_variants(factor: MismatchFactor) -> tuple[FactorVariant, ...]: + """Expand one factor into its variants. + + Four arms rather than on/off: only a one-sided swap tells you which side is + at fault, and only a two-sided swap proves the reference itself is sound. + """ + + if factor.variants: + return tuple(factor.variants) + + path = factor.switch.path + reference = factor.reference + + if reference is None: # a parameter sweep + allowed = factor.switch.allowed_values or () + return tuple( + FactorVariant( + name=f"value_{value}", + switch_values={path: value}, + why=f"sweep {path} = {value!r}", + ) + for value in allowed + ) + + variants = [ + FactorVariant( + name="both_native", + switch_values={path: "native"}, + why="baseline: each side on its own framework's native implementation", + ), + FactorVariant( + name="both_reference", + switch_values={path: reference.name}, + replace_on={ + PolicyRole.ROLLOUT: reference.rollout_impl, + PolicyRole.TRAINING: reference.training_impl, + }, + expected=ExpectedOutcome.BITWISE_IDENTICAL, + why=( + "self-check gate: both sides on one implementation must agree bitwise, " + "otherwise every conclusion from this factor is void" + ), + ), + FactorVariant( + name="training_reference_only", + switch_values={path: f"{reference.name}@training"}, + replace_on={PolicyRole.TRAINING: reference.training_impl}, + why="swap the training side only: if the deviation goes, that side is the source", + ), + FactorVariant( + name="rollout_reference_only", + switch_values={path: f"{reference.name}@rollout"}, + replace_on={PolicyRole.ROLLOUT: reference.rollout_impl}, + why="swap the rollout side only: if the deviation goes, that side is the source", + ), + ] + + if reference.fp64_oracle is not None: + variants.append( + FactorVariant( + name="fp64_oracle", + switch_values={path: "fp64_oracle"}, + replace_on={ + PolicyRole.ROLLOUT: reference.fp64_oracle, + PolicyRole.TRAINING: reference.fp64_oracle, + }, + expected=ExpectedOutcome.BITWISE_IDENTICAL, + why="gold standard: proves the reference used by both_reference is itself correct", + ) + ) + + return tuple(variants) + + +def order_cases_by_rebind_cost( + cases: Sequence[tuple[MismatchFactor, FactorVariant]], +) -> tuple[tuple[MismatchFactor, FactorVariant], ...]: + """Order cases so rebuild cost never decreases, maximising reuse. + + Required, not a nicety: 160 cases in a random order restart the process for + every one of them in the worst case. + """ + + return tuple(sorted(cases, key=lambda item: _REBIND_ORDER[item[0].switch.rebind_cost])) + + +def suggested_floor_is_lowest(factor: MismatchFactor, floor: NoiseFloor) -> bool: + """Whether this factor can show anything at the given floor. + + A process-level switch is identical on a single device, so running it at the + anchor floor wastes machine time. + """ + + if factor.switch.rebind_cost is RebindCost.PROCESS_GROUP_REBUILD: + return floor in (NoiseFloor.SHARDED_SINGLE_NODE, NoiseFloor.PRODUCTION) + return True + + +__all__ = [ + "ContradictoryFactor", + "UnmetPrerequisite", + "build_variants", + "missing_prerequisites", + "order_cases_by_rebind_cost", + "reject_contradictory_factors", + "suggested_floor_is_lowest", +] diff --git a/rl_engine/mismatch/pipeline/registry.py b/rl_engine/mismatch/pipeline/registry.py new file mode 100644 index 00000000..69faaeae --- /dev/null +++ b/rl_engine/mismatch/pipeline/registry.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plugin registration and factor discovery. + +Adding an operator is adding a directory; adding a factor is adding a file. If +you find yourself editing the planner, a global dict, or another operator's file, +the abstraction is missing something -- raise it rather than patching in place. +""" + +from __future__ import annotations + +import importlib +import pkgutil +from typing import Any, Callable, Mapping, Protocol + +from rl_engine.mismatch.schema import ( + CollectiveContract, + ComparisonRule, + ImplementationResolution, + MismatchFactor, + OperatorContract, + PolicyRole, +) + + +class OperatorChecks(Protocol): + """Everything one operator needs checked. The plugin itself. + + These four are operator-level, not factor-level: reading configuration back + from an engine is the same logic for all of an operator's factors. Factor + files hold declarations only. + """ + + operator: str + + def declare_factors(self) -> tuple[MismatchFactor, ...]: + """Which factors this operator has.""" + ... + + def build_contract( + self, role: PolicyRole, switch_values: Mapping[str, Any] + ) -> OperatorContract: + """Turn this side's switch values into that side's numerical contract.""" + ... + + def read_effective_config(self, role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read the switches' effective values back. A requested value is not + evidence.""" + ... + + def observe_collectives(self, role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """Which collectives actually ran.""" + ... + + def resolve_implementation( + self, factor_id: str, role: PolicyRole, impl_name: str + ) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve the name a variant asks for into a callable. + + Return the trace even when resolution fails: a bare ``None`` leaves a + silent fallback with nothing to investigate. + """ + ... + + +class RegistrationError(ValueError): + """A plugin conflicts with one already registered.""" + + +class PluginRegistry: + """Registered operator plugins, with conflict detection at import time.""" + + def __init__(self) -> None: + self._plugins: dict[str, OperatorChecks] = {} + + def register(self, plugin_cls: type) -> type: + """Instantiate and register a plugin, checking for conflicts.""" + + plugin = plugin_cls() + name = getattr(plugin, "operator", "") + if not name: + raise RegistrationError(f"{plugin_cls.__name__} does not declare an operator name") + if name in self._plugins: + raise RegistrationError(f"operator {name!r} is already registered") + + self._check_factor_conflicts(plugin) + self._plugins[name] = plugin + return plugin_cls + + def _check_factor_conflicts(self, plugin: OperatorChecks) -> None: + """Reject duplicate ids, duplicate switch paths, and one contract field + claimed at two different comparison rules.""" + + known_ids = {factor.id for p in self._plugins.values() for factor in p.declare_factors()} + known_paths = { + factor.switch.path for p in self._plugins.values() for factor in p.declare_factors() + } + rules: dict[str, ComparisonRule] = {} + for existing in self._plugins.values(): + for factor in existing.declare_factors(): + rules.update(factor.comparison_rules) + + local_ids: set[str] = set() + local_paths: set[str] = set() + for factor in plugin.declare_factors(): + if factor.id in known_ids: + raise RegistrationError(f"duplicate factor id {factor.id!r}") + if factor.switch.path in known_paths: + raise RegistrationError(f"duplicate switch path {factor.switch.path!r}") + if factor.id in local_ids: + raise RegistrationError(f"duplicate factor id {factor.id!r}") + if factor.switch.path in local_paths: + raise RegistrationError(f"duplicate switch path {factor.switch.path!r}") + for field_path, rule in factor.comparison_rules.items(): + previous = rules.get(field_path) + if previous is not None and previous is not rule: + raise RegistrationError( + f"contract field {field_path!r} is declared as {previous.value!r} " + f"elsewhere but {rule.value!r} by {factor.id!r}" + ) + rules[field_path] = rule + local_ids.add(factor.id) + local_paths.add(factor.switch.path) + + def operators(self) -> tuple[str, ...]: + return tuple(sorted(self._plugins)) + + def plugin(self, operator: str) -> OperatorChecks: + if operator not in self._plugins: + raise KeyError(f"operator {operator!r} is not registered; known: {self.operators()}") + return self._plugins[operator] + + def factors_for(self, operator: str | None = None) -> tuple[MismatchFactor, ...]: + """All factors, or just one operator's.""" + + if operator is not None: + return tuple(self.plugin(operator).declare_factors()) + collected: list[MismatchFactor] = [] + for name in self.operators(): + collected.extend(self._plugins[name].declare_factors()) + return tuple(collected) + + def clear(self) -> None: + """Test helper: drop every registration.""" + + self._plugins.clear() + + +OPERATOR_CHECKS = PluginRegistry() + + +class FactorDiscoveryError(ValueError): + """A factor module does not follow the discovery convention.""" + + +def discover_factors(package: str) -> tuple[MismatchFactor, ...]: + """Collect the ``FACTOR`` constant from every module under ``.factors``. + + A file's name must equal its factor id with the operator prefix stripped, so + that renaming an id without renaming the file fails at import rather than + silently dropping the factor. + """ + + factors_package = f"{package}.factors" + module = importlib.import_module(factors_package) + collected: list[MismatchFactor] = [] + + for info in pkgutil.iter_modules(module.__path__): + if info.name.startswith("_"): + continue + submodule = importlib.import_module(f"{factors_package}.{info.name}") + factor = getattr(submodule, "FACTOR", None) + if factor is None: + raise FactorDiscoveryError(f"{factors_package}.{info.name} defines no FACTOR constant") + expected_suffix = factor.id.split(".", 1)[-1] + if info.name != expected_suffix: + raise FactorDiscoveryError( + f"{factors_package}.{info.name}.py declares factor id {factor.id!r}; " + f"the file should be named {expected_suffix}.py" + ) + collected.append(factor) + + return tuple(sorted(collected, key=lambda item: item.id)) + + +__all__ = [ + "FactorDiscoveryError", + "OPERATOR_CHECKS", + "OperatorChecks", + "PluginRegistry", + "RegistrationError", + "discover_factors", +] diff --git a/rl_engine/mismatch/pipeline/report.py b/rl_engine/mismatch/pipeline/report.py new file mode 100644 index 00000000..e241adee --- /dev/null +++ b/rl_engine/mismatch/pipeline/report.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""False-positive filtering, root-cause tracing, and the final report.""" + +from __future__ import annotations + +from typing import Sequence + +from rl_engine.mismatch.schema import ( + Diagnosis, + FactorReport, + KnownPitfall, + LibraryPin, + MismatchReport, + ModuleCorrespondence, + NoiseFloor, + PropagationEdge, + RootCauseCategory, + RootCauseHypothesis, +) + +_SIDE_DIAGNOSES = ( + Diagnosis.CAUSED_BY_TRAINING_SIDE, + Diagnosis.CAUSED_BY_ROLLOUT_SIDE, + Diagnosis.CAUSED_BY_BOTH_SIDES, +) + + +def filter_known_equivalences( + correspondences: Sequence[ModuleCorrespondence], + findings: Sequence[str], +) -> tuple[tuple[str, ...], tuple[ModuleCorrespondence, ...]]: + """Drop findings explained by a known structural equivalence. + + An equivalence without ``verified_by`` is not trusted: unproven, "filtering + false positives" quietly becomes "hiding real findings". + """ + + proven = tuple( + item + for item in correspondences + if item.equivalence is not None and item.verified_by is not None + ) + explained = {item.semantic_name for item in proven} + kept = tuple(finding for finding in findings if finding not in explained) + return kept, proven + + +def trace_root_causes( + reports: Sequence[FactorReport], + correspondences: Sequence[ModuleCorrespondence], + edges: Sequence[PropagationEdge], +) -> tuple[RootCauseHypothesis, ...]: + """Walk from the still-aligned anchor down the call chain into ranked + hypotheses.""" + + downstream_of: dict[str, list[str]] = {} + for edge in edges: + downstream_of.setdefault(edge.upstream, []).append(edge.downstream) + + implicated = [report for report in reports if report.diagnosis in _SIDE_DIAGNOSES] + if not implicated: + return () + + module_of = {item.semantic_name: item for item in correspondences} + hypotheses: list[RootCauseHypothesis] = [] + + for rank, report in enumerate(_by_confidence(implicated), start=1): + module = _module_for_factor(report.factor_id, module_of) + anchor = _anchor_for(module, downstream_of) + hypotheses.append( + RootCauseHypothesis( + suspected_module=module, + category=_categorise(report.diagnosis), + anchor_module=anchor, + supporting_factors=(report.factor_id,), + evidence=(report.diagnosis_reason,), + rank=rank, + ) + ) + + return tuple(hypotheses) + + +def _by_confidence(reports: Sequence[FactorReport]) -> list[FactorReport]: + """A one-sided attribution is more actionable than a both-sided one.""" + + order = { + Diagnosis.CAUSED_BY_TRAINING_SIDE: 0, + Diagnosis.CAUSED_BY_ROLLOUT_SIDE: 0, + Diagnosis.CAUSED_BY_BOTH_SIDES: 1, + } + return sorted(reports, key=lambda report: (order[report.diagnosis], report.factor_id)) + + +def _module_for_factor(factor_id: str, module_of: dict[str, ModuleCorrespondence]) -> str: + """Map a factor id onto a semantic module name, falling back to the id.""" + + for name in module_of: + if factor_id.startswith(name.split(".", 1)[0]): + return name + return factor_id + + +def _anchor_for(module: str, downstream_of: dict[str, list[str]]) -> str: + """The last still-aligned position upstream of the suspect.""" + + for upstream, children in downstream_of.items(): + if module in children: + return upstream + return module + + +def _categorise(diagnosis: Diagnosis) -> RootCauseCategory: + if diagnosis is Diagnosis.CAUSED_BY_BOTH_SIDES: + return RootCauseCategory.DIFFERENT_IMPLEMENTATION + return RootCauseCategory.DIFFERENT_IMPLEMENTATION + + +def build_report( + reports: Sequence[FactorReport], + *, + noise_floor: NoiseFloor, + library_pins: Sequence[LibraryPin] = (), + correspondences: Sequence[ModuleCorrespondence] = (), + edges: Sequence[PropagationEdge] = (), + failed_guards: Sequence[KnownPitfall] = (), +) -> MismatchReport: + """Assemble the final report.""" + + _, filtered = filter_known_equivalences( + correspondences, [report.factor_id for report in reports] + ) + hypotheses = trace_root_causes(reports, correspondences, edges) + + return MismatchReport( + noise_floor=noise_floor, + library_pins=tuple(library_pins), + factor_reports=tuple(reports), + hypotheses=hypotheses, + filtered_false_positives=filtered, + failed_guards=tuple(failed_guards), + ) + + +def render_summary(report: MismatchReport) -> str: + """A short human-readable summary, led by the ranked hypotheses.""" + + lines = [ + f"noise floor: {report.noise_floor.value}", + f"factors run: {len(report.factor_reports)}", + ] + + tally: dict[str, int] = {} + for factor_report in report.factor_reports: + tally[factor_report.diagnosis.value] = tally.get(factor_report.diagnosis.value, 0) + 1 + for name in sorted(tally): + lines.append(f" {name}: {tally[name]}") + + if report.hypotheses: + lines.append("root-cause hypotheses (most suspicious first):") + for hypothesis in report.hypotheses: + lines.append( + f" {hypothesis.rank}. {hypothesis.suspected_module} " + f"[{hypothesis.category.value}] anchor={hypothesis.anchor_module}" + ) + else: + lines.append("root-cause hypotheses: none (no factor was attributed to a side)") + + if report.failed_guards: + lines.append(f"failed pitfall guards: {[g.id for g in report.failed_guards]}") + + return "\n".join(lines) + + +__all__ = [ + "build_report", + "filter_known_equivalences", + "render_summary", + "trace_root_causes", +] diff --git a/rl_engine/mismatch/pipeline/runner.py b/rl_engine/mismatch/pipeline/runner.py new file mode 100644 index 00000000..93250f09 --- /dev/null +++ b/rl_engine/mismatch/pipeline/runner.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The execution loop: reuse decisions, variant repeats, metric computation.""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, replace +from typing import Any, Callable, Mapping, Protocol, Sequence + +from rl_engine.mismatch.pipeline.comparison import compare_contracts +from rl_engine.mismatch.pipeline.registry import OperatorChecks +from rl_engine.mismatch.schema import ( + DEFAULT_CLIP_EPS, + ComparisonIdentity, + ExecutionPath, + FactorVariant, + MismatchFactor, + MismatchMetrics, + PolicyRole, + ReuseKey, + SwitchStatus, + VariantResult, + WorstToken, +) + + +class ScoringBackend(Protocol): + """What the runner needs from an engine: logprobs for a fixed sequence. + + Minimal enough that a CPU stub can satisfy it, so the plumbing is testable + without a GPU. + """ + + def score( + self, + role: PolicyRole, + identity: ComparisonIdentity, + switch_values: Mapping[str, Any], + replacement: Callable[..., Any] | None, + ) -> tuple[Sequence[float], Mapping[str, Any]]: + """Return per-token logprobs plus whatever was read back.""" + ... + + def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: ... + + +class ReadOnlyViolation(RuntimeError): + """Computing logprobs modified the model.""" + + +@dataclass(frozen=True) +class RunContext: + """Everything one run needs that is not the factor itself.""" + + identity: ComparisonIdentity + path: ExecutionPath = ExecutionPath.TRAINING_FULL_PREFILL + clip_eps: float = DEFAULT_CLIP_EPS + strict: bool = True + + +def assert_comparison_is_read_only(before: str, after: str) -> None: + """Model tensor fingerprints must match before and after scoring. + + A kernel that mutates weights in place makes ``both_reference`` fail bitwise + at random, and the blame lands on the reference. This assertion is the + self-check gate's own premise. + """ + + if before != after: + raise ReadOnlyViolation( + f"model state changed while computing logprobs: {before} -> {after}" + ) + + +def compute_metrics( + rollout_logprobs: Sequence[float], + training_logprobs: Sequence[float], + active_mask: Sequence[bool], + *, + clip_eps: float = DEFAULT_CLIP_EPS, + token_ids: Sequence[int] | None = None, +) -> MismatchMetrics: + """Per-token metrics over active tokens only. + + ``dlogp = log pi_theta - log pi_old``, so ``rho = exp(dlogp)``, clipped at + ``1 +/- eps``. Past that edge a token's gradient signal is discarded. + """ + + deltas: list[float] = [] + positions: list[int] = [] + for index, (roll, train, active) in enumerate( + # strict: a length disagreement between logprobs and mask is a real bug + zip(rollout_logprobs, training_logprobs, active_mask, strict=True) + ): + if not active: + continue + deltas.append(float(train) - float(roll)) + positions.append(index) + + if not deltas: + return MismatchMetrics( + active_token_count=0, + dlogp_mean=0.0, + dlogp_p99=0.0, + dlogp_max=0.0, + ratio_mean=1.0, + ratio_max=1.0, + clip_fraction=0.0, + approx_kl=0.0, + ) + + magnitudes = [abs(delta) for delta in deltas] + ratios = [math.exp(delta) for delta in deltas] + upper_edge = math.log1p(clip_eps) + lower_edge = -math.log1p(-clip_eps) if clip_eps < 1.0 else float("inf") + clipped = sum(1 for delta in deltas if delta > upper_edge or delta < -abs(lower_edge)) + + # k3 estimator: rho - 1 - ln(rho) + approx_kl = sum(ratio - 1.0 - math.log(ratio) for ratio in ratios) / len(ratios) + + worst_index = max(range(len(deltas)), key=lambda i: magnitudes[i]) + worst = WorstToken( + position=positions[worst_index], + token_id=(token_ids[positions[worst_index]] if token_ids else -1), + dlogp=deltas[worst_index], + ) + + return MismatchMetrics( + active_token_count=len(deltas), + dlogp_mean=sum(magnitudes) / len(magnitudes), + dlogp_p99=_percentile(magnitudes, 0.99), + dlogp_max=max(magnitudes), + ratio_mean=sum(ratios) / len(ratios), + ratio_max=max(ratios, key=lambda r: abs(math.log(r))), + clip_fraction=clipped / len(deltas), + approx_kl=approx_kl, + worst_token=worst, + ) + + +def _percentile(values: Sequence[float], q: float) -> float: + ordered = sorted(values) + if not ordered: + return 0.0 + index = min(len(ordered) - 1, int(math.ceil(q * len(ordered)) - 1)) + return ordered[max(0, index)] + + +def expand_repeats(variant: FactorVariant) -> tuple[Mapping[str, Any], ...]: + """Expand ``repeat_under`` into the cartesian product of environments. + + The only exception to "one variant, one execution". + """ + + if not variant.repeat_under: + return ({},) + keys = sorted(variant.repeat_under) + combos = itertools.product(*(variant.repeat_under[key] for key in keys)) + return tuple(dict(zip(keys, combo, strict=True)) for combo in combos) + + +def assert_order_is_topology_independent(results: Sequence[Sequence[float]]) -> bool: + """Runs of a topology-independent collective must agree bitwise. + + Single-sided reruns, so the cheapest check in the whole set. + """ + + if len(results) < 2: + return True + first = list(results[0]) + return all(list(other) == first for other in results[1:]) + + +def run_variant( + factor: MismatchFactor, + variant: FactorVariant, + checks: OperatorChecks, + backends: Mapping[PolicyRole, ScoringBackend], + context: RunContext, +) -> VariantResult: + """Run one variant on both sides and produce its result.""" + + resolutions = {} + replacements: dict[PolicyRole, Callable[..., Any] | None] = {} + status = SwitchStatus.APPLIED + + for role in (PolicyRole.ROLLOUT, PolicyRole.TRAINING): + wanted = (variant.replace_on or {}).get(role) + if wanted is None: + replacements[role] = None + continue + callable_impl, resolution = checks.resolve_implementation(factor.id, role, wanted) + replacements[role] = callable_impl + resolutions[role] = resolution + if callable_impl is None: + status = SwitchStatus.FELL_BACK + + scores: dict[PolicyRole, Sequence[float]] = {} + readbacks: dict[PolicyRole, Mapping[str, Any]] = {} + repeats: dict[PolicyRole, list[Sequence[float]]] = { + PolicyRole.ROLLOUT: [], + PolicyRole.TRAINING: [], + } + + for environment in expand_repeats(variant): + merged = {**variant.switch_values, **environment} + for role in (PolicyRole.ROLLOUT, PolicyRole.TRAINING): + logprobs, readback = backends[role].score( + role, context.identity, merged, replacements[role] + ) + repeats[role].append(logprobs) + scores[role] = logprobs + readbacks[role] = readback + + if variant.repeat_under: + for role in (PolicyRole.ROLLOUT, PolicyRole.TRAINING): + if not assert_order_is_topology_independent(repeats[role]): + status = SwitchStatus.ERROR + + effective_configs = { + role: checks.read_effective_config(role, readbacks[role]) + for role in (PolicyRole.ROLLOUT, PolicyRole.TRAINING) + } + contracts = {} + for role in (PolicyRole.ROLLOUT, PolicyRole.TRAINING): + contract = checks.build_contract(role, effective_configs[role]) + # Runtime collective traces win over contracts inferred from flags. + # Backends may choose a different path per shape and topology. + observed = checks.observe_collectives(role, effective_configs[role]) + contracts[role] = replace(contract, collectives=observed) + issues = compare_contracts( + contracts[PolicyRole.ROLLOUT], contracts[PolicyRole.TRAINING], (factor,) + ) + + metrics = compute_metrics( + scores[PolicyRole.ROLLOUT], + scores[PolicyRole.TRAINING], + context.identity.active_mask, + clip_eps=context.clip_eps, + token_ids=context.identity.response_token_ids, + ) + + evidence = frozenset( + itertools.chain.from_iterable(readbacks[role].get("evidence", ()) for role in readbacks) + ) + + return VariantResult( + variant=variant, + path=context.path, + status=status, + metrics=metrics, + evidence=evidence, + effective_config={ + f"{role.value}.{key}": value + for role, effective in effective_configs.items() + for key, value in effective.items() + if key != "evidence" + }, + collectives_observed=tuple( + itertools.chain.from_iterable(contracts[role].collectives for role in contracts) + ), + resolution=resolutions.get(PolicyRole.TRAINING) or resolutions.get(PolicyRole.ROLLOUT), + comparison_issues=issues, + ) + + +__all__ = [ + "ReadOnlyViolation", + "RunContext", + "ScoringBackend", + "assert_comparison_is_read_only", + "assert_order_is_topology_independent", + "compute_metrics", + "expand_repeats", + "run_variant", +] diff --git a/rl_engine/mismatch/reference_adapters/__init__.py b/rl_engine/mismatch/reference_adapters/__init__.py new file mode 100644 index 00000000..8c460842 --- /dev/null +++ b/rl_engine/mismatch/reference_adapters/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Wiring implementations in as reference implementations. No operator code here. + +``kernels/`` and the external libraries own whether the arithmetic is right. +This package owns putting them into a deterministic mode and proving the setting +took effect, which is specific to this framework rather than to any operator. +""" + +from rl_engine.mismatch.reference_adapters.settings import ( + SettingDeliveryError, + apply_required_settings, + verify_required_settings, +) + +__all__ = [ + "SettingDeliveryError", + "apply_required_settings", + "verify_required_settings", +] diff --git a/rl_engine/mismatch/reference_adapters/settings.py b/rl_engine/mismatch/reference_adapters/settings.py new file mode 100644 index 00000000..526027c9 --- /dev/null +++ b/rl_engine/mismatch/reference_adapters/settings.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Delivering pinned settings by channel, and reading them back. + +Neither TransformerEngine nor FlashInfer is deterministic by default, so these +have to be pinned explicitly -- and pinning alone is not enough, since a setting +that cannot be read back is only ``UNOBSERVABLE``. +""" + +from __future__ import annotations + +import os +from typing import Any, Mapping, Sequence + +from rl_engine.mismatch.schema import RequiredSetting, SettingChannel, SwitchStatus + + +class SettingDeliveryError(RuntimeError): + """A setting could not be delivered through its declared channel.""" + + +def apply_required_settings( + settings: Sequence[RequiredSetting], + *, + engine_kwargs: dict[str, Any] | None = None, + call_kwargs: dict[str, Any] | None = None, + environ: dict[str, str] | None = None, +) -> dict[SettingChannel, dict[str, Any]]: + """Route each setting to the right place for its channel. + + The channel decides when a setting can take effect and therefore its rebind + cost: an env var needs a process restart, a call argument does not. + """ + + target_env = environ if environ is not None else os.environ + delivered: dict[SettingChannel, dict[str, Any]] = {channel: {} for channel in SettingChannel} + + for setting in settings: + if setting.channel is SettingChannel.ENV_VAR: + target_env[setting.key] = str(setting.value) + elif setting.channel is SettingChannel.TORCH_GLOBAL: + _apply_torch_global(setting) + elif setting.channel is SettingChannel.ENGINE_ARG: + if engine_kwargs is None: + raise SettingDeliveryError( + f"{setting.key!r} is an engine argument but no engine kwargs were given; " + "it can only take effect when the engine is rebuilt" + ) + engine_kwargs[setting.key] = setting.value + elif setting.channel is SettingChannel.CALL_ARG: + if call_kwargs is None: + raise SettingDeliveryError( + f"{setting.key!r} is a call argument but no call kwargs were given" + ) + call_kwargs[setting.key] = setting.value + delivered[setting.channel][setting.key] = setting.value + + return delivered + + +def _apply_torch_global(setting: RequiredSetting) -> None: + """Set a ``torch.backends.*`` flag by dotted path.""" + + try: + import torch + except ImportError as exc: # pragma: no cover - torch is a hard dependency + raise SettingDeliveryError(f"cannot set {setting.key!r} without torch") from exc + + target: Any = torch + parts = setting.key.split(".") + if parts[0] == "torch": + parts = parts[1:] + for part in parts[:-1]: + target = getattr(target, part) + setattr(target, parts[-1], setting.value) + + +def verify_required_settings( + settings: Sequence[RequiredSetting], + readback: Mapping[str, Any], +) -> tuple[SwitchStatus, tuple[str, ...]]: + """Check the values that came back against what was pinned. + + A setting with no readback path is ``UNOBSERVABLE``: delivered, unproven. + """ + + unobservable: list[str] = [] + mismatched: list[str] = [] + + for setting in settings: + if setting.readback is None: + unobservable.append(setting.key) + continue + if setting.key not in readback: + unobservable.append(setting.key) + continue + actual = readback[setting.key] + if isinstance(setting.value, str) and setting.value.startswith(">="): + continue # a constraint rather than an exact value + if actual != setting.value: + mismatched.append(f"{setting.key}: pinned {setting.value!r}, read {actual!r}") + + if mismatched: + return SwitchStatus.FELL_BACK, tuple(mismatched) + if unobservable: + return SwitchStatus.UNOBSERVABLE, tuple(f"{key}: no readback path" for key in unobservable) + return SwitchStatus.APPLIED, () + + +__all__ = [ + "SettingDeliveryError", + "apply_required_settings", + "verify_required_settings", +] diff --git a/rl_engine/mismatch/schema/__init__.py b/rl_engine/mismatch/schema/__init__.py new file mode 100644 index 00000000..b093accf --- /dev/null +++ b/rl_engine/mismatch/schema/__init__.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pure data structures: public fields, frozen, no meaningful methods. + +Behaviour lives in free functions under ``pipeline/``. The framework grows by +adding functions while the set of types stays stable, which is the side +procedural code is good at -- so do not hang methods on these structures. Write +``requires_fixed_order(contract)``, not ``contract.requires_fixed_order()``. + +Chained field access such as ``factor.switch.path`` is fine: Demeter constrains +an object's internals, and plain data is meant to expose its fields. + +Modules are ordered by dependency; each imports only from the ones above it. +""" + +from rl_engine.mismatch.schema.collectives import ( + ALL_REDUCE_AS_SCATTER_GATHER, + ALL_TO_ALL_AS_GATHER_SLICE, + CollectiveContract, + CollectiveOp, + CollectiveRewrite, + DeterminismLevel, + ParallelDim, + ReductionOrder, +) +from rl_engine.mismatch.schema.contracts import ( + ComparisonIssue, + ComparisonIssueCode, + ComparisonRule, + OperatorContract, +) +from rl_engine.mismatch.schema.factors import ( + BATCH_PLACEMENT, + COLLECTIVE_CONTRACT, + LSE_EXPORT, + MODEL_SHAPE, + POSITION_CACHE, + VOCAB_SHARD_MAP, + Evidence, + FactorCategory, + MismatchFactor, + Prerequisites, + ReferenceAuthority, + ReferenceImplementation, + Switch, + declared_collectives, + requires_fixed_order, +) +from rl_engine.mismatch.schema.fingerprints import ( + EnvironmentFingerprint, + ExecutionFingerprint, + ReuseKey, + VariantRecord, + canonical_fingerprint, + reuse_level, +) +from rl_engine.mismatch.schema.metrics import ( + DEFAULT_CLIP_EPS, + FactorReport, + ImplementationResolution, + LogprobShard, + MismatchMetrics, + RejectedCandidate, + VariantResult, + WorstToken, + is_silent_failure, + missing_evidence, +) +from rl_engine.mismatch.schema.pitfalls import FailureMode, KnownPitfall +from rl_engine.mismatch.schema.rollout_context import ( + BatchPlacement, + ComparisonIdentity, + DynamicSamplingDecision, + RolloutGroup, +) +from rl_engine.mismatch.schema.thresholds import ( + ANY_MODEL_FAMILY, + EXPECTED_RANGES, + ThresholdLookupError, + expected_range, + tolerance_floor, +) +from rl_engine.mismatch.schema.tracing import ( + MismatchReport, + ModuleCorrespondence, + PropagationEdge, + RootCauseCategory, + RootCauseHypothesis, +) +from rl_engine.mismatch.schema.values import ( + DowncastPoint, + ExecutionPath, + LibraryPin, + PolicyRole, + Precision, + PrecisionProfile, + RebindCost, + RequiredSetting, + SettingChannel, + choice_parser, + positive_int, + strict_bool, +) +from rl_engine.mismatch.schema.variants import ( + Diagnosis, + ExpectedOutcome, + ExpectedRange, + FactorVariant, + NoiseFloor, + SwitchStatus, + VariantExpansion, +) + +__all__ = [ + "ALL_REDUCE_AS_SCATTER_GATHER", + "ALL_TO_ALL_AS_GATHER_SLICE", + "CollectiveContract", + "CollectiveOp", + "CollectiveRewrite", + "DeterminismLevel", + "ParallelDim", + "ReductionOrder", + "ComparisonIssue", + "ComparisonIssueCode", + "ComparisonRule", + "OperatorContract", + "BATCH_PLACEMENT", + "COLLECTIVE_CONTRACT", + "LSE_EXPORT", + "MODEL_SHAPE", + "POSITION_CACHE", + "VOCAB_SHARD_MAP", + "Evidence", + "FactorCategory", + "MismatchFactor", + "Prerequisites", + "ReferenceAuthority", + "ReferenceImplementation", + "Switch", + "declared_collectives", + "requires_fixed_order", + "EnvironmentFingerprint", + "ExecutionFingerprint", + "ReuseKey", + "VariantRecord", + "canonical_fingerprint", + "reuse_level", + "DEFAULT_CLIP_EPS", + "FactorReport", + "ImplementationResolution", + "LogprobShard", + "MismatchMetrics", + "RejectedCandidate", + "VariantResult", + "WorstToken", + "is_silent_failure", + "missing_evidence", + "FailureMode", + "KnownPitfall", + "BatchPlacement", + "ComparisonIdentity", + "DynamicSamplingDecision", + "RolloutGroup", + "ANY_MODEL_FAMILY", + "EXPECTED_RANGES", + "ThresholdLookupError", + "expected_range", + "tolerance_floor", + "MismatchReport", + "ModuleCorrespondence", + "PropagationEdge", + "RootCauseCategory", + "RootCauseHypothesis", + "DowncastPoint", + "ExecutionPath", + "LibraryPin", + "PolicyRole", + "Precision", + "PrecisionProfile", + "RebindCost", + "RequiredSetting", + "SettingChannel", + "choice_parser", + "positive_int", + "strict_bool", + "Diagnosis", + "ExpectedOutcome", + "ExpectedRange", + "FactorVariant", + "NoiseFloor", + "SwitchStatus", + "VariantExpansion", +] diff --git a/rl_engine/mismatch/schema/collectives.py b/rl_engine/mismatch/schema/collectives.py new file mode 100644 index 00000000..44bbd0fa --- /dev/null +++ b/rl_engine/mismatch/schema/collectives.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Collective communication as a first-class concept. + +Mismatch comes from floating-point addition not being associative, and the +accumulation order is almost entirely decided by collective communication. Six +factors across gemm, attention, logprob and MoE are instances of this one +semantic model; written separately they would drift apart. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +from rl_engine.mismatch.schema.values import DowncastPoint, LibraryPin, Precision + + +class CollectiveOp(str, Enum): + ALL_REDUCE = "all_reduce" + REDUCE_SCATTER = "reduce_scatter" + ALL_GATHER = "all_gather" + ALL_TO_ALL = "all_to_all" + BROADCAST = "broadcast" + POINT_TO_POINT = "point_to_point" + NONE = "none" # single-device path, recorded rather than left blank + + +class ParallelDim(str, Enum): + """Which parallel dimension the communication happens on. + + Not ``ProcessGroupKind``: ``ProcessGroup`` is an existing torch type. + """ + + TENSOR = "tensor" + SEQUENCE = "sequence" + CONTEXT = "context" + EXPERT = "expert" + PIPELINE = "pipeline" + DATA = "data" + + +class ReductionOrder(str, Enum): + """Accumulation order -- the direct root of mismatch, not the collective.""" + + ARRIVAL = "arrival" # control group only + NCCL_ALGORITHM = "nccl_algorithm" # varies with world size and message size + GLOBAL_RANK_INDEX = "global_rank_index" + GLOBAL_BLOCK_INDEX = "global_block_index" # CP / split-K merge + GLOBAL_VOCAB_SHARD_INDEX = "global_vocab_shard_index" + + +class DeterminismLevel(str, Enum): + """How strong a reproducibility guarantee an implementation offers.""" + + NONE = "none" + STABLE_WITHIN_PROCESS = "stable_within_process" + STABLE_ACROSS_RUNS = "stable_across_runs" + STABLE_ACROSS_TOPOLOGY = "stable_across_topology" + + +@dataclass(frozen=True) +class CollectiveContract: + """The full numerical semantics of one collective. + + Claiming ``STABLE_ACROSS_TOPOLOGY`` while reducing with ``NCCL_ALGORITHM`` or + ``ARRIVAL`` produces numbers that mean nothing; + ``reject_contradictory_factors()`` rejects that at planning time. + """ + + op: CollectiveOp + group: ParallelDim + group_size: int + reduction_order: ReductionOrder + accumulate_precision: Precision + downcast_at: DowncastPoint + determinism: DeterminismLevel + backend: str # "nccl" / "vllm_custom_ipc" / "mnnvl" / "transformer_engine" / "rl_kernel" + pinned_libraries: tuple[LibraryPin, ...] = () + + +@dataclass(frozen=True) +class CollectiveRewrite: + """A rewrite that is equal in algebra and unequal in floating point. + + Whether each side applies one is itself a source of mismatch, so it has to be + declarable. + """ + + name: str + original: tuple[CollectiveOp, ...] + rewritten: tuple[CollectiveOp, ...] + preserves_bitwise: bool = False # always False -- that is the whole problem + + +ALL_REDUCE_AS_SCATTER_GATHER = CollectiveRewrite( + name="all_reduce -> reduce_scatter + all_gather", + original=(CollectiveOp.ALL_REDUCE,), + rewritten=(CollectiveOp.REDUCE_SCATTER, CollectiveOp.ALL_GATHER), +) + +ALL_TO_ALL_AS_GATHER_SLICE = CollectiveRewrite( + name="all_to_all -> all_gather + local slice", + original=(CollectiveOp.ALL_TO_ALL,), + rewritten=(CollectiveOp.ALL_GATHER,), +) + + +__all__ = [ + "ALL_REDUCE_AS_SCATTER_GATHER", + "ALL_TO_ALL_AS_GATHER_SLICE", + "CollectiveContract", + "CollectiveOp", + "CollectiveRewrite", + "DeterminismLevel", + "ParallelDim", + "ReductionOrder", +] diff --git a/rl_engine/mismatch/schema/contracts.py b/rl_engine/mismatch/schema/contracts.py new file mode 100644 index 00000000..4da47f83 --- /dev/null +++ b/rl_engine/mismatch/schema/contracts.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Numerical contracts and how the two sides are compared field by field.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Mapping + +from rl_engine.mismatch.schema.collectives import CollectiveContract +from rl_engine.mismatch.schema.values import PolicyRole, PrecisionProfile + + +class ComparisonRule(str, Enum): + """What the framework does when this field differs between the two sides.""" + + MUST_MATCH_BITWISE = "must_match_bitwise" # differ -> this case is void + MUST_MATCH_SEMANTICALLY = "must_match_semantically" # implementations may differ + RECORD_ONLY = "record_only" # recorded, never compared + + +class ComparisonIssueCode(str, Enum): + """Stable reason codes for the two sides disagreeing. + + These strings go into the artifact schema and callers branch on them, so + renaming one requires a schema version bump. + """ + + REQUIRED_FIELD_MISSING = "required_field_missing" + BITWISE_MISMATCH = "bitwise_mismatch" + SEMANTIC_MISMATCH = "semantic_mismatch" + DETERMINISM_INCOMPATIBLE = "determinism_incompatible" + + +@dataclass(frozen=True) +class ComparisonIssue: + """One record of the two sides disagreeing.""" + + code: ComparisonIssueCode + rule: ComparisonRule # which tier this field was declared at + field_path: str # "collectives[0].reduction_order" + values: Mapping[PolicyRole, Any] # each side's actual value + message: str = "" + + +@dataclass(frozen=True) +class OperatorContract: + """One operator's numerical contract on one side. + + Only the three fields common to every operator live here; the rest goes in + ``extra``, keyed by the paths a factor's ``comparison_rules`` declare:: + + "precision.accumulate" + "collectives[0].reduction_order" + "extra.rope_theta" + + The framework indexes both contracts by path, so a plugin only has to put + fields in the right place. Keep ``extra`` flat -- nesting makes paths long and + unreadable. + """ + + operator: str + role: PolicyRole + precision: PrecisionProfile + collectives: tuple[CollectiveContract, ...] = () + extra: Mapping[str, Any] = field(default_factory=dict) + + +__all__ = [ + "ComparisonIssue", + "ComparisonIssueCode", + "ComparisonRule", + "OperatorContract", +] diff --git a/rl_engine/mismatch/schema/factors.py b/rl_engine/mismatch/schema/factors.py new file mode 100644 index 00000000..beca3fd3 --- /dev/null +++ b/rl_engine/mismatch/schema/factors.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The factor model: one suspected cause of training-inference mismatch.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Callable, Mapping + +from rl_engine.mismatch.schema.collectives import CollectiveContract, DeterminismLevel +from rl_engine.mismatch.schema.contracts import ComparisonRule +from rl_engine.mismatch.schema.pitfalls import KnownPitfall +from rl_engine.mismatch.schema.values import ( + ExecutionPath, + LibraryPin, + PolicyRole, + RebindCost, + RequiredSetting, + choice_parser, +) + + +class FactorCategory(str, Enum): + """Which family a factor belongs to, and so where to look when it fires.""" + + INPUT_IDENTITY = "input_identity" # tokens / mask / position_ids / eps + ENVIRONMENT = "environment" # framework versions, NCCL, determinism switches + KERNEL_IMPLEMENTATION = "kernel_implementation" # backend, fusion, inner precision + SHARDING_AND_REDUCTION = "sharding_and_reduction" # TP/CP/SP, reduction order, split-K + OUTPUT_NUMERICS = "output_numerics" # logits / logp / (out, lse) / gradients + + +class Evidence(str, Enum): + """Evidence every factor must have before a verdict is allowed. + + Operator-specific evidence stays out of this enum: putting it here would turn + "add an operator" into "change the framework". Plugins declare their own as + plain strings, like the constants below. + """ + + EFFECTIVE_CONFIG_READBACK = "effective_config_readback" # read back, not requested + MODEL_STATE_FINGERPRINT = "model_state_fingerprint" + LIBRARY_VERSIONS = "library_versions" + + +COLLECTIVE_CONTRACT = "collective_contract" +BATCH_PLACEMENT = "batch_placement" +MODEL_SHAPE = "model_shape" +POSITION_CACHE = "position_cache" +VOCAB_SHARD_MAP = "vocab_shard_map" +LSE_EXPORT = "lse_export" + + +class ReferenceAuthority(str, Enum): + """Where a reference implementation comes from, most authoritative first. + + A decision order, not a description: look for a SHARED_BACKEND first, and + write SELF_WRITTEN only when the first two cannot cover it. + """ + + FP64_ORACLE = "fp64_oracle" # slow, exact, lowest noise floor only + SHARED_BACKEND = "shared_backend" # TransformerEngine / FlashInfer + SELF_WRITTEN = "self_written" + + +@dataclass(frozen=True) +class Switch: + """A switch's one definition: allowed values and parser declared together.""" + + path: str # "gemm.forward_reduce" + rebind_cost: RebindCost + applies_to: tuple[PolicyRole, ...] + allowed_values: tuple[Any, ...] | None = None + parse: Callable[[Any], Any] | None = None + + def __post_init__(self) -> None: + if self.parse is None and self.allowed_values is not None: + object.__setattr__(self, "parse", choice_parser(*self.allowed_values)) + + +@dataclass(frozen=True) +class Prerequisites: + """What a factor needs in order to run. A whitelist, not a blacklist.""" + + required_ops: tuple[str, ...] = () + min_gpu_count: int = 1 + required_packages: tuple[str, ...] = () # "transformer_engine>=2.0" + required_model_traits: tuple[str, ...] = () # "moe" / "linear_attention" + blocked_by: tuple[str, ...] = () # work this factor waits on + + +@dataclass(frozen=True) +class ReferenceImplementation: + """What replaces the native implementation, and which paths it covers. + + ``covers_paths`` defines the shape of the self-check gate: every path this + reference covers must agree bitwise on the same sequence. Covering two paths + puts the gate across the two sides; a reference that also covers + ``ROLLOUT_DECODE`` puts it inside the rollout side, and then no decode stub is + needed on the training side. + + See ``docs/add-a-kernel-factor.md`` for how that shapes attribution. + """ + + name: str + tier: ReferenceAuthority + training_impl: str + rollout_impl: str + covers_paths: tuple[ExecutionPath, ...] + fp64_oracle: str | None = None + required_settings: tuple[RequiredSetting, ...] = () + pinned_libraries: tuple[LibraryPin, ...] = () + + +@dataclass(frozen=True) +class MismatchFactor: + """One suspected cause of training-inference mismatch. + + Not ``DivergenceFactor``: in an RL context, divergence means KL divergence. + + ``reference is None`` makes it a parameter sweep, otherwise an implementation + swap. There is no separate ``kind`` field -- derivable state is state that can + disagree with itself. + """ + + id: str # "gemm.forward_reduce", globally unique + operator: str + category: FactorCategory + question: str # what this factor answers, one line, goes into the docs + switch: Switch + comparison_rules: Mapping[str, ComparisonRule] # contract field path -> rule + prerequisites: Prerequisites + required_evidence: tuple[str, ...] = () # Evidence values or plugin constants + reference: ReferenceImplementation | None = None + call_sites: tuple[str, ...] = () # one factor acting in several places + pitfalls: tuple[KnownPitfall, ...] = () + variants: tuple[Any, ...] = () # empty -> expand the standard set + + +def declared_collectives(factor: MismatchFactor) -> tuple[CollectiveContract, ...]: + """Collectives a factor's reference pins, for the planner's static check.""" + + reference = factor.reference + if reference is None: + return () + return tuple( + setting.value + for setting in reference.required_settings + if isinstance(setting.value, CollectiveContract) + ) + + +def requires_fixed_order(contract: CollectiveContract) -> bool: + """Whether this contract claims its result is independent of topology.""" + + return contract.determinism is DeterminismLevel.STABLE_ACROSS_TOPOLOGY + + +__all__ = [ + "BATCH_PLACEMENT", + "COLLECTIVE_CONTRACT", + "Evidence", + "FactorCategory", + "LSE_EXPORT", + "MODEL_SHAPE", + "MismatchFactor", + "POSITION_CACHE", + "Prerequisites", + "ReferenceAuthority", + "ReferenceImplementation", + "Switch", + "VOCAB_SHARD_MAP", + "declared_collectives", + "requires_fixed_order", +] diff --git a/rl_engine/mismatch/schema/fingerprints.py b/rl_engine/mismatch/schema/fingerprints.py new file mode 100644 index 00000000..442de7e1 --- /dev/null +++ b/rl_engine/mismatch/schema/fingerprints.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Execution lifetime and identity. + +Two sides of one question: identical identity is what makes reuse safe, and a +changed identity is what makes historical results stale. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from typing import Any, Mapping + +from rl_engine.mismatch.schema.metrics import VariantResult +from rl_engine.mismatch.schema.values import LibraryPin, PolicyRole, RebindCost + + +@dataclass(frozen=True) +class ReuseKey: + """Whether an already-built runtime can be reused. + + Four parts matching the four ``RebindCost`` levels, compared coarse to fine. + """ + + process: str # env vars, determinism switches, compile-time flags + process_group: str # world size, TP/CP/PP split, comm backend + engine: str # dtype, backend choice, KV layout, operator implementation + request: str # batch size, sequence + + +@dataclass(frozen=True) +class EnvironmentFingerprint: + """The execution environment. Change this layer and every number is stale.""" + + python_version: str + torch_version: str + torch_build_hash: str # hash of the build config (cuda/hip build, op set) + driver_version: str + device_model: str + libraries: tuple[LibraryPin, ...] + determinism_env: Mapping[str, str] # NVTE_* / CUBLAS_* / NCCL_* / torch backends + source_revision: str # this framework's own version + + +@dataclass(frozen=True) +class ExecutionFingerprint: + """One execution's full identity. Any part differing makes two runs + incomparable. + + What goes in is the value read back, never the value requested: asking for + ``num_splits=1`` and the backend using 1 are two different facts. Thresholds + go in too, so changing one makes every historical pass/fail stale -- which is + why thresholds are code constants, a configurable value cannot be pinned into + an identity. + """ + + identity: str # fingerprint of the ComparisonIdentity + environment: EnvironmentFingerprint + switch_binding: str # effective switch values, read back + implementation: Mapping[PolicyRole, str] # what each side actually instantiated + model_state: Mapping[PolicyRole, str] # each side's weights + collectives: tuple[str, ...] # fingerprints of the collectives that ran + threshold_table: str # fingerprint of EXPECTED_RANGES + + +@dataclass(frozen=True) +class VariantRecord: + """One variant's archived record: a ``VariantResult`` plus its identity. + + ``content_hash`` is its integrity seal. Only a record whose hash verifies may + be reused on resume. + """ + + variant_name: str + fingerprint: ExecutionFingerprint + result: VariantResult + content_hash: str + + +def canonical_fingerprint(payload: Any) -> str: + """Hash a JSON-serialisable payload, normalising key order.""" + + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def reuse_level(previous: ReuseKey, current: ReuseKey) -> RebindCost: + """How far a rebuild has to go between two cases. Earlier is costlier.""" + + if previous.process != current.process: + return RebindCost.PROCESS_RESTART + if previous.process_group != current.process_group: + return RebindCost.PROCESS_GROUP_REBUILD + if previous.engine != current.engine: + return RebindCost.ENGINE_REBUILD + return RebindCost.PER_REQUEST + + +__all__ = [ + "EnvironmentFingerprint", + "ExecutionFingerprint", + "ReuseKey", + "VariantRecord", + "canonical_fingerprint", + "reuse_level", +] diff --git a/rl_engine/mismatch/schema/metrics.py b/rl_engine/mismatch/schema/metrics.py new file mode 100644 index 00000000..1267ee89 --- /dev/null +++ b/rl_engine/mismatch/schema/metrics.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Metrics and per-variant results. + +``dlogp`` is an intermediate quantity; what the objective acts on is +``rho = exp(dlogp)``, clipped at ``1 +/- eps``. Past ``ln(1 + eps) ~= 0.182`` a +token's gradient signal is discarded, which is how mismatch breaks training: not +random samples are dropped, the most mismatched ones are. + +Healthy means run 0.002-0.008 (dense) and 0.01-0.03 (large MoE), all far below +that edge, so judging on the mean alone always concludes everything is fine. The +danger is entirely in the tail. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping + +from rl_engine.mismatch.schema.collectives import CollectiveContract +from rl_engine.mismatch.schema.contracts import ComparisonIssue +from rl_engine.mismatch.schema.values import ExecutionPath +from rl_engine.mismatch.schema.variants import Diagnosis, FactorVariant, NoiseFloor, SwitchStatus + +DEFAULT_CLIP_EPS = 0.2 + + +@dataclass(frozen=True) +class WorstToken: + """The largest-deviation token, which often points straight at a layer.""" + + position: int + token_id: int + dlogp: float + layer_hint: str | None = None + expert_hint: int | None = None + + +@dataclass(frozen=True) +class MismatchMetrics: + """Every metric from one comparison. Active tokens only.""" + + active_token_count: int + dlogp_mean: float + dlogp_p99: float + dlogp_max: float + ratio_mean: float # rho = exp(dlogp) + ratio_max: float + clip_fraction: float # share of active tokens past the clip edge + approx_kl: float # k3 estimator: rho - 1 - ln(rho) + worst_token: WorstToken | None = None + + +@dataclass(frozen=True) +class RejectedCandidate: + """A rejected candidate implementation and why it was rejected.""" + + name: str + reason: str + + +@dataclass(frozen=True) +class ImplementationResolution: + """Which candidates were tried and why each was rejected. + + What you look at when the status is ``FELL_BACK``. A single reason string is + not enough -- a silent fallback leaves nothing to investigate. + """ + + requested: str + resolved: str | None # None = no candidate was usable + rejected: tuple[RejectedCandidate, ...] = () # in the order tried + + +@dataclass(frozen=True) +class LogprobShard: + """The slice of logprobs one rank holds under TP/CP. + + Missing a slice is wrong in a way that does not show: drop one vocab shard + and the LSE denominator loses a chunk, so logp comes out systematically high. + Slices are counted against ``world_size`` before merging. + """ + + rank: int + world_size: int + selected_logprobs: Any # torch.Tensor + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class VariantResult: + """What one variant produced.""" + + variant: FactorVariant + path: ExecutionPath + status: SwitchStatus + metrics: MismatchMetrics | None + evidence: frozenset[str] # Evidence values plus plugin constants + effective_config: Mapping[str, Any] # read back, not requested + collectives_observed: tuple[CollectiveContract, ...] = () + resolution: ImplementationResolution | None = None # required unless APPLIED + comparison_issues: tuple[ComparisonIssue, ...] = () + logprob_shards: tuple[LogprobShard, ...] = () + + +@dataclass(frozen=True) +class FactorReport: + """The conclusion for one factor, from all of its variants together.""" + + factor_id: str + noise_floor: NoiseFloor + variants: tuple[VariantResult, ...] + diagnosis: Diagnosis + diagnosis_reason: str + + +def is_silent_failure(metrics: MismatchMetrics) -> bool: + """Mean within band but the tail already past the clip edge. + + The most common false negative: the headline number looks healthy while + gradient signal is being discarded. + """ + + return metrics.dlogp_mean < 0.01 and metrics.clip_fraction > 0.0 + + +def missing_evidence(collected: frozenset[str], required: tuple[str, ...]) -> frozenset[str]: + """Which required evidence is absent.""" + + return frozenset(required) - collected + + +__all__ = [ + "DEFAULT_CLIP_EPS", + "FactorReport", + "ImplementationResolution", + "LogprobShard", + "MismatchMetrics", + "RejectedCandidate", + "VariantResult", + "WorstToken", + "is_silent_failure", + "missing_evidence", +] diff --git a/rl_engine/mismatch/schema/pitfalls.py b/rl_engine/mismatch/schema/pitfalls.py new file mode 100644 index 00000000..7341f099 --- /dev/null +++ b/rl_engine/mismatch/schema/pitfalls.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Known pitfalls, encoded as data. + +Prose pitfalls get read once and never again. As data, the framework can block +them before a run instead of relying on somebody remembering afterwards. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +from rl_engine.mismatch.schema.variants import NoiseFloor + + +class FailureMode(str, Enum): + """How a pitfall fools you, which decides what tool works against it.""" + + STRUCTURAL_FALSE_POSITIVE = "structural_false_positive" # differs in form, equal in math + SILENT_FALSE_NEGATIVE = "silent_false_negative" # metrics look fine, conclusion is wrong + MISSING_INSTRUMENTATION = "missing_instrumentation" # never captured in the first place + CONFIG_DEFAULT_TRAP = "config_default_trap" # the default is not what you assumed + CONVENTION_MISMATCH = "convention_mismatch" # shift-by-one, log base, ... + RESOURCE_LIMIT = "resource_limit" # this arm simply cannot run + + +@dataclass(frozen=True) +class KnownPitfall: + """A known pitfall together with the assertion that blocks it. + + ``symptom`` and ``actual_cause`` are separate because a pitfall is a pitfall + precisely when its appearance points at the wrong cause. + """ + + id: str + mode: FailureMode + symptom: str + actual_cause: str + guard: str + guard_runs_at: NoiseFloor # lowest floor that can run it -- cheap checks first + + +__all__ = ["FailureMode", "KnownPitfall"] diff --git a/rl_engine/mismatch/schema/rollout_context.py b/rl_engine/mismatch/schema/rollout_context.py new file mode 100644 index 00000000..c8368a8a --- /dev/null +++ b/rl_engine/mismatch/schema/rollout_context.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The online RL context around one comparison. + +All of these belong to ``INPUT_IDENTITY`` rather than to the environment: they +change how samples are grouped for reduction, so two runs that differ here are +not reproducible against each other. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RolloutGroup: + """One GRPO group: K rollouts from the same prompt. + + Advantage is normalised within the group, so one rollout's logp deviating + reaches all K through it. Averaging over tokens alone hides that. Offline + ablation leaves this a placeholder; it means something only when wired into a + real training loop. + """ + + prompt_id: str + rollout_ids: tuple[str, ...] + group_size: int # GRPO's K + + +@dataclass(frozen=True) +class BatchPlacement: + """Where one sample landed in this training step. + + Not ``BatchLayout``: in tensor land, layout means memory ordering. DP and + microbatch splitting decide which samples reduce together, so moving a sample + changes its accumulation order. + """ + + data_parallel_rank: int + microbatch_index: int + position_in_microbatch: int + dropped_by_schedule: bool = False + + +@dataclass(frozen=True) +class DynamicSamplingDecision: + """Record of dropping a group, which changes the batch composition.""" + + kept: bool + reason: str | None = None + + +@dataclass(frozen=True) +class ComparisonIdentity: + """This comparison's input identity. + + If these differ, no numerical comparison means anything. + """ + + prompt_token_ids: tuple[int, ...] + response_token_ids: tuple[int, ...] + active_mask: tuple[bool, ...] # loss mask: which tokens participate + position_ids: tuple[int, ...] + checkpoint_id: str + checkpoint_revision: str + model_shape: str # a trimmed model is a different model + group: RolloutGroup + batch_placement: BatchPlacement + sampling_decision: DynamicSamplingDecision + + +__all__ = [ + "BatchPlacement", + "ComparisonIdentity", + "DynamicSamplingDecision", + "RolloutGroup", +] diff --git a/rl_engine/mismatch/schema/thresholds.py b/rl_engine/mismatch/schema/thresholds.py new file mode 100644 index 00000000..a8f9a8fa --- /dev/null +++ b/rl_engine/mismatch/schema/thresholds.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The threshold table, as code constants. + +Thresholds are not configuration. A tunable threshold means somebody can tune it +until the test passes, so this table lives in code, goes into the execution +fingerprint, and changing it invalidates historical pass/fail results. +""" + +from __future__ import annotations + +from rl_engine.mismatch.schema.variants import ExpectedRange, NoiseFloor + +ANY_MODEL_FAMILY = "*" + + +EXPECTED_RANGES: tuple[ExpectedRange, ...] = ( + # -- production floor, empirical -- + ExpectedRange("dense", NoiseFloor.PRODUCTION, None, (0.002, 0.008), 0.01), + ExpectedRange("moe", NoiseFloor.PRODUCTION, False, (0.007, 0.008), 0.02), + ExpectedRange( + "moe", + NoiseFloor.PRODUCTION, + True, + (0.0, 0.008), + 0.008, + note="routing replay on but no drop means replay never took effect", + ), + ExpectedRange( + "large_moe", + NoiseFloor.PRODUCTION, + False, + (0.01, 0.03), + 0.05, + note="GLM5 / DSv3.2, DSA+MoE, 2k-8k. This band is normal, do not file it as a bug", + ), + ExpectedRange( + "large_moe", + NoiseFloor.PRODUCTION, + True, + (0.0, 0.008), + 0.008, + note="no fall back to e-3 means the routing capture path is broken", + ), + # -- low-noise floors: the expectation is bitwise, not "small" -- + ExpectedRange( + ANY_MODEL_FAMILY, + NoiseFloor.SINGLE_LAYER_ANCHOR, + None, + (0.0, 0.0), + 1e-6, + note="failing here is an operator bug, not training-inference mismatch", + ), + ExpectedRange(ANY_MODEL_FAMILY, NoiseFloor.FULL_MODEL_SINGLE_GPU, None, (0.0, 1e-6), 1e-5), + ExpectedRange(ANY_MODEL_FAMILY, NoiseFloor.SHARDED_SINGLE_NODE, None, (0.0, 1e-4), 1e-3), +) + + +class ThresholdLookupError(LookupError): + """No band declared for this combination.""" + + +def expected_range( + model_family: str, + noise_floor: NoiseFloor, + routing_replay: bool | None = None, +) -> ExpectedRange: + """Look up the normal band for one combination. + + Always keyed by noise floor: judging an anchor-floor run against the + production band would call a definite operator error normal. An exact model + family beats the wildcard. + """ + + exact = [ + candidate + for candidate in EXPECTED_RANGES + if candidate.model_family == model_family + and candidate.noise_floor is noise_floor + and candidate.routing_replay == routing_replay + ] + if exact: + return exact[0] + + wildcard = [ + candidate + for candidate in EXPECTED_RANGES + if candidate.model_family == ANY_MODEL_FAMILY and candidate.noise_floor is noise_floor + ] + if wildcard: + return wildcard[0] + + raise ThresholdLookupError( + f"no expected range declared for model_family={model_family!r} " + f"noise_floor={noise_floor.value!r} routing_replay={routing_replay!r}" + ) + + +def tolerance_floor(model_family: str, noise_floor: NoiseFloor) -> float: + """The floor below which a difference is not treated as a signal.""" + + return expected_range(model_family, noise_floor).suspect_above + + +__all__ = [ + "ANY_MODEL_FAMILY", + "EXPECTED_RANGES", + "ThresholdLookupError", + "expected_range", + "tolerance_floor", +] diff --git a/rl_engine/mismatch/schema/tracing.py b/rl_engine/mismatch/schema/tracing.py new file mode 100644 index 00000000..e9eccb05 --- /dev/null +++ b/rl_engine/mismatch/schema/tracing.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Module correspondence and root-cause tracing.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +from rl_engine.mismatch.schema.metrics import FactorReport +from rl_engine.mismatch.schema.pitfalls import KnownPitfall +from rl_engine.mismatch.schema.values import LibraryPin +from rl_engine.mismatch.schema.variants import NoiseFloor + + +@dataclass(frozen=True) +class ModuleCorrespondence: + """A training-side module paired with its rollout-side counterpart. + + A non-empty ``equivalence`` means "different in form, equal in arithmetic", + which is what lets a false positive be filtered. Without it a difference like + fused QKV turns every weight comparison red and buries the real problem. + """ + + semantic_name: str # "mlp.gate_up" -- framework-independent + training_module: str # "...megatron...linear_fc1" + rollout_module: str # "...vllm...gate_up_proj" + equivalence: str | None = None # "concat_on_dim0" / "transpose" / ... + verified_by: str | None = None # the test proving it. No proof, no claim. + + +@dataclass(frozen=True) +class PropagationEdge: + """One directed edge of the call chain, followed backwards when tracing.""" + + upstream: str # semantic_name + downstream: str + + +class RootCauseCategory(str, Enum): + MISSING_OPERATOR = "missing_operator" # one side does not have it at all + DIFFERENT_IMPLEMENTATION = "different_implementation" + DIFFERENT_PARAMETER = "different_parameter" + UPSTREAM_PROPAGATED = "upstream_propagated" # fine itself, inherited from upstream + + +@dataclass(frozen=True) +class RootCauseHypothesis: + """One hypothesis from walking the call chain after a ``NOT_THIS_FACTOR``. + + The root cause must be downstream of ``anchor_module``, the last position + where the two sides still agree. + """ + + suspected_module: str + category: RootCauseCategory + anchor_module: str + supporting_factors: tuple[str, ...] # MismatchFactor.id + evidence: tuple[str, ...] + rank: int # 1 is the most suspicious + + +@dataclass(frozen=True) +class MismatchReport: + """The final report for one run. + + Thirty factors give thirty diagnoses, and that pile is not the answer. This + combines them with the module correspondence table and the call chain into + ``hypotheses`` -- the few most suspicious modules, ranked. The other fields + are the evidence supporting it. + """ + + noise_floor: NoiseFloor + library_pins: tuple[LibraryPin, ...] + factor_reports: tuple[FactorReport, ...] + hypotheses: tuple[RootCauseHypothesis, ...] # sorted by rank + filtered_false_positives: tuple[ModuleCorrespondence, ...] + failed_guards: tuple[KnownPitfall, ...] + + +__all__ = [ + "MismatchReport", + "ModuleCorrespondence", + "PropagationEdge", + "RootCauseCategory", + "RootCauseHypothesis", +] diff --git a/rl_engine/mismatch/schema/values.py b/rl_engine/mismatch/schema/values.py new file mode 100644 index 00000000..9fc9db43 --- /dev/null +++ b/rl_engine/mismatch/schema/values.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Base value types. Imports nothing from the project. + +Everything here is a plain data structure: public fields, ``frozen=True``, no +meaningful methods. Behaviour lives in free functions under ``pipeline/``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Callable + + +class PolicyRole(str, Enum): + """Which of the two compared policies this side plays. + + Not ``Policy``: in RL that means the policy itself, not the part it plays. + """ + + ROLLOUT = "rollout" # pi_old -- produced the trajectory + TRAINING = "training" # pi_theta -- being updated + + +class ExecutionPath(str, Enum): + """Which physical path produced a set of logprobs. + + Whole sequence at once, in chunks, or token by token: mathematically + equivalent, not equal in floating point. + """ + + TRAINING_FULL_PREFILL = "training_full_prefill" + ROLLOUT_FULL_PREFILL = "rollout_full_prefill" + ROLLOUT_CHUNKED_PREFILL = "rollout_chunked_prefill" + ROLLOUT_DECODE = "rollout_decode" + + +class Precision(str, Enum): + FP64 = "fp64" # oracle only + FP32 = "fp32" + BF16 = "bf16" + FP16 = "fp16" + FP8_E4M3 = "fp8_e4m3" + FP8_E5M2 = "fp8_e5m2" + + +class DowncastPoint(str, Enum): + """When a high-precision accumulator is written back at lower precision. + + Numerical precision reduction, not the C++ sense of casting down a hierarchy. + """ + + NEVER = "never" + PER_BLOCK = "per_block" # largest error + PER_PARTIAL = "per_partial" + FINAL_WRITE = "final_write" # smallest error + + +class RebindCost(str, Enum): + """What it costs to change a switch. Drives case ordering.""" + + PER_REQUEST = "per_request" + ENGINE_REBUILD = "engine_rebuild" + PROCESS_GROUP_REBUILD = "process_group_rebuild" + PROCESS_RESTART = "process_restart" + + +class SettingChannel(str, Enum): + """How a setting reaches the engine, which decides when it takes effect.""" + + ENV_VAR = "env_var" # read once at process start -> PROCESS_RESTART + TORCH_GLOBAL = "torch_global" # torch.backends.* -> PROCESS_RESTART + ENGINE_ARG = "engine_arg" # engine constructor -> ENGINE_REBUILD + CALL_ARG = "call_arg" # passed per call -> PER_REQUEST + + +@dataclass(frozen=True) +class LibraryPin: + """A pinned library version. + + TransformerEngine and FlashInfer change kernel selection across versions, so + the same factor can reach the opposite conclusion on a different one. The pin + goes into the execution fingerprint rather than a footnote. + """ + + package: str + version: str # exact; ranges are not accepted + commit: str | None = None + container_digest: str | None = None # the only truly reproducible anchor + + +@dataclass(frozen=True) +class RequiredSetting: + """A setting that must be pinned, and how to prove it was. + + A setting that cannot be read back can only be recorded ``UNOBSERVABLE``: + delivered but unverifiable is the same as not delivered. + """ + + key: str + value: Any + channel: SettingChannel + readback: str | None = None + guards: str = "" # KnownPitfall.id + + +@dataclass(frozen=True) +class PrecisionProfile: + """Which precision this side uses at each point in the computation. + + Not ``PrecisionPolicy``: in RL, ``Policy`` is pi. + """ + + compute: Precision + accumulate: Precision + downcast_at: DowncastPoint + master_weights: Precision | None = None + lm_head: Precision | None = None # should be FP32 + softmax_accumulate: Precision | None = None + kv_accumulate: Precision | None = None # linear attention KV, should be FP32 + + +def choice_parser(*allowed: Any) -> Callable[[Any], Any]: + """Build a parser accepting only the given values.""" + + permitted = tuple(allowed) + + def parse(value: Any) -> Any: + if value not in permitted: + raise ValueError(f"expected one of {permitted!r}, got {value!r}") + return value + + return parse + + +def positive_int(value: Any) -> int: + parsed = int(value) + if parsed <= 0: + raise ValueError(f"expected a positive integer, got {value!r}") + return parsed + + +def strict_bool(value: Any) -> bool: + if not isinstance(value, bool): + raise ValueError(f"expected a bool, got {value!r}") + return value + + +__all__ = [ + "DowncastPoint", + "ExecutionPath", + "LibraryPin", + "PolicyRole", + "Precision", + "PrecisionProfile", + "RebindCost", + "RequiredSetting", + "SettingChannel", + "choice_parser", + "positive_int", + "strict_bool", +] diff --git a/rl_engine/mismatch/schema/variants.py b/rl_engine/mismatch/schema/variants.py new file mode 100644 index 00000000..2588be6b --- /dev/null +++ b/rl_engine/mismatch/schema/variants.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Variants, diagnosis, and the noise floor ladder.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Mapping + +from rl_engine.mismatch.schema.values import PolicyRole + + +class VariantExpansion(str, Enum): + """Which variants a factor expands into.""" + + STANDARD_FOUR = "standard_four" # swap factors + VALUE_SWEEP = "value_sweep" # sweep factors: one run per allowed value + PAIRWISE = "pairwise" # only after COUPLED_WITH_OTHER_FACTORS is diagnosed + + +class ExpectedOutcome(str, Enum): + BITWISE_IDENTICAL = "bitwise_identical" # failing means the reference is at fault + MEASURE_ONLY = "measure_only" + + +class SwitchStatus(str, Enum): + """Whether the switch actually reached the engine. + + A silent fallback is far more harmful than an error. + """ + + APPLIED = "applied" + FELL_BACK = "fell_back" # requested, silently reverted to native + UNSUPPORTED = "unsupported" + UNOBSERVABLE = "unobservable" # delivered but unreadable -- no evidence + ERROR = "error" + + +class Diagnosis(str, Enum): + """One factor's conclusion after its variants have run. + + The first three mean "cannot judge" and must stay strictly separate from + "judged, nothing here". Answering what causes the mismatch overall needs + cross-factor synthesis, which is ``MismatchReport``. + """ + + VARIANT_DID_NOT_APPLY = "variant_did_not_apply" + INSUFFICIENT_EVIDENCE = "insufficient_evidence" + REFERENCE_ITSELF_IS_BROKEN = "reference_itself_is_broken" + CAUSED_BY_TRAINING_SIDE = "caused_by_training_side" + CAUSED_BY_ROLLOUT_SIDE = "caused_by_rollout_side" + CAUSED_BY_BOTH_SIDES = "caused_by_both_sides" + NOT_THIS_FACTOR = "not_this_factor" + COUPLED_WITH_OTHER_FACTORS = "coupled_with_other_factors" + + +class NoiseFloor(str, Enum): + """How small a difference this run can resolve. Orthogonal to factors. + + Each step down introduces exactly one new noise source, so when a floor + starts failing the suspect set is whatever that floor just added. A floor + that has not passed blocks the next one. + """ + + SINGLE_LAYER_ANCHOR = "single_layer_anchor" + # 1 layer, single device, determinism on, one token. No noise sources at all, + # which is why failing bitwise here is an operator bug rather than mismatch, + # and the other three floors need not run. + + FULL_MODEL_SINGLE_GPU = "full_model_single_gpu" + # All layers, still single device. New: accumulation over depth. Tests + # whether error grows linearly or exponentially with layer count. + + SHARDED_SINGLE_NODE = "sharded_single_node" + # TP + SP on one node. New: reduction-order differences from sharding, so + # this is the first floor with real training-inference mismatch. + + PRODUCTION = "production" + # Target TP/CP/PP, determinism off, decode path. New: everything else. The + # only floor whose numbers may be read against the threshold table. + + +@dataclass(frozen=True) +class FactorVariant: + """One arm of a controlled experiment, as pasteable switch values. + + ``repeat_under`` runs this same arm once per environment and requires bitwise + equality -- the only exception to "one variant, one execution", expanded by + the runner as a cartesian product. It never compares across frameworks, so it + is cheap, and it verifies the premise of the self-check gate: an arm can only + anchor the others if its fixed-order implementation really did fix the order. + """ + + name: str + switch_values: Mapping[str, Any] + replace_on: Mapping[PolicyRole, str] | None = None + expected: ExpectedOutcome = ExpectedOutcome.MEASURE_ONLY + why: str = "" + repeat_under: Mapping[str, tuple[Any, ...]] | None = None + + +@dataclass(frozen=True) +class ExpectedRange: + """The normal band for a metric under one (model family, noise floor, config). + + A code constant, never configuration: a tunable threshold is one somebody + tunes until the test passes. It enters the execution fingerprint, so changing + it invalidates every historical pass/fail. + """ + + model_family: str # "dense" / "moe" / "large_moe" / "*" + noise_floor: NoiseFloor + routing_replay: bool | None # None = not applicable + dlogp_mean: tuple[float, float] # normal band [low, high] + suspect_above: float + note: str = "" + + +__all__ = [ + "Diagnosis", + "ExpectedOutcome", + "ExpectedRange", + "FactorVariant", + "NoiseFloor", + "SwitchStatus", + "VariantExpansion", +] diff --git a/tests/mismatch_cpu_backend.py b/tests/mismatch_cpu_backend.py new file mode 100644 index 00000000..51b8c46b --- /dev/null +++ b/tests/mismatch_cpu_backend.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A CPU scoring backend for exercising the framework without a GPU. + +Plumbing only: it says nothing about real Megatron or vLLM numerics. What it can +do is simulate the failure modes the framework exists to catch -- a per-side +bias, a switch that silently does nothing, and output that changes with the +environment -- so the gates and the matrix are tested against something that +really fails. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, Sequence + +from rl_engine.mismatch.schema import ComparisonIdentity, Evidence, PolicyRole, ReuseKey + + +@dataclass +class CpuScoringBackend: + """Deterministic synthetic logprobs, with injectable deviation. + + Setting ``bias`` on one side only simulates a one-sided root cause. + """ + + role: PolicyRole + bias: float = 0.0 + silently_ignores: frozenset[str] = frozenset() + unstable_under: frozenset[str] = frozenset() + evidence: frozenset[str] = frozenset( + { + Evidence.EFFECTIVE_CONFIG_READBACK.value, + Evidence.MODEL_STATE_FINGERPRINT.value, + Evidence.LIBRARY_VERSIONS.value, + } + ) + applied_calls: list[Mapping[str, Any]] = field(default_factory=list) + + def score( + self, + role: PolicyRole, + identity: ComparisonIdentity, + switch_values: Mapping[str, Any], + replacement: Callable[..., Any] | None, + ) -> tuple[Sequence[float], Mapping[str, Any]]: + """Produce logprobs for the fixed sequence. + + The base signal depends only on the identity, so both sides agree unless + a bias is injected. + """ + + self.applied_calls.append(dict(switch_values)) + base = _base_signal(identity) + + # The reference is defined to have no bias of its own. + bias = 0.0 if replacement is not None else self.bias + + drift = 0.0 + for key in sorted(self.unstable_under): + if key in switch_values: + drift += _env_jitter(key, switch_values[key]) + + logprobs = [value + bias + drift for value in base] + + effective = { + key: ("ignored" if key in self.silently_ignores else value) + for key, value in switch_values.items() + } + effective["evidence"] = tuple(self.evidence) + return logprobs, effective + + def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: + """Group switches by tier, mirroring the four ``RebindCost`` levels.""" + + def digest(prefix: str) -> str: + payload = sorted( + f"{key}={value}" for key, value in switch_values.items() if key.startswith(prefix) + ) + return hashlib.sha256("|".join(payload).encode()).hexdigest()[:12] + + return ReuseKey( + process=digest("env."), + process_group=digest("dist."), + engine=digest("engine."), + request=digest("batch."), + ) + + +def _base_signal(identity: ComparisonIdentity) -> list[float]: + """A stable pseudo-logprob per token, derived only from the identity.""" + + tokens = identity.response_token_ids + signal: list[float] = [] + for index, token in enumerate(tokens): + seed = hashlib.sha256(f"{identity.checkpoint_id}:{index}:{token}".encode()).digest() + magnitude = int.from_bytes(seed[:4], "big") / 0xFFFFFFFF + signal.append(-(0.5 + magnitude)) # logprobs are negative + return signal + + +def _env_jitter(key: str, value: Any) -> float: + """Tiny environment-dependent shift, used to make instability observable.""" + + seed = hashlib.sha256(f"{key}={value}".encode()).digest() + return (int.from_bytes(seed[:2], "big") / 0xFFFF) * 1e-3 + + +__all__ = ["CpuScoringBackend"] diff --git a/tests/test_mismatch_attention_adapter.py b/tests/test_mismatch_attention_adapter.py new file mode 100644 index 00000000..b7ed0976 --- /dev/null +++ b/tests/test_mismatch_attention_adapter.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Attention adapter tests: actual Split-KV/CP evidence, not requested flags.""" + +from __future__ import annotations + +import json +from copy import deepcopy + +import pytest + +from rl_engine.mismatch.operator_checks.attention import AttentionChecks, adapter +from rl_engine.mismatch.operator_checks.attention._common import CP_MERGE_REFERENCE +from rl_engine.mismatch.schema import ( + CollectiveOp, + DowncastPoint, + PolicyRole, + Precision, + ReductionOrder, +) + + +def _plan_set( + *, + cp_world_size: int = 2, + boundaries: dict[int, list[list[int]]] | None = None, + actual_mode: str = "auto", + actual_size: int | None = None, + fallback: bool = False, +) -> dict: + total = 8 + ranges = [(rank * 4, (rank + 1) * 4) for rank in range(cp_world_size)] + if boundaries is None: + boundaries = { + owner: [[start, start + 2], [start + 2, end]] + for owner, (start, end) in enumerate(ranges) + } + entries = [] + for cp_rank in range(cp_world_size): + for owner, expected_range in enumerate(ranges): + entries.append( + { + "batch_index": 0, + "tp_rank": 0, + "cp_rank": cp_rank, + "owner_cp_rank": owner, + "expected_kv_range": list(expected_range), + "requested_split_kv_policy": "auto", + "requested_split_kv_size": None, + "actual_split_kv_policy": actual_mode, + "actual_split_kv_size": actual_size, + "actual_split_boundaries": boundaries[owner], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "fixture", + "split_kv_plan_source": "runtime_trace", + "split_kv_fallback": fallback, + "split_kv_fallback_reason": "shape fallback" if fallback else None, + } + ) + return { + "batch_size": 1, + "tp_world_size": 1, + "cp_world_size": cp_world_size, + "total_kv_tokens": [total], + "entries": entries, + } + + +def _manifest() -> list[dict]: + return [ + { + "global_block_index": 0, + "kv_block_start": 0, + "kv_block_end": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0, + }, + { + "global_block_index": 1, + "kv_block_start": 4, + "kv_block_end": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0, + }, + ] + + +def _collective(**overrides) -> dict: + result = { + "op": "point_to_point", + "group_size": 2, + "reduction_order": "global_block_index", + "accumulate_precision": "fp32", + "downcast_at": "final_write", + "determinism": "stable_across_topology", + "backend": "p2p_nccl_reference", + } + result.update(overrides) + return result + + +def _effective(**overrides) -> dict: + result = { + "attn.compute_dtype": "bf16", + "attn.accumulate_dtype": "fp32", + "attn.downcast_at": "final_write", + "attn.batch_size": 1, + "attn.tp_world_size": 1, + "attn.cp_world_size": 2, + "attn.actual_split_kv_plan_set": _plan_set(), + "attn.cp_block_manifest": _manifest(), + "attn.cp_collective": _collective(), + "attn.lse_domain": "attention", + "attn.export_lse": True, + "attn.merge_state": "out_lse", + "attn.rope_theta": 1_000_000.0, + "attn.position_ids_digest": "positions:abc", + "attn.post_rope_qk_digest": "qk:def", + "attn.q_rope_state": "post_rope", + "attn.k_rope_state": "post_rope", + "attn.k_cache_rope_state": "post_rope", + "attn.fusion_boundary": "unfused_rope_attention", + } + result.update(overrides) + return result + + +def test_contract_maps_actual_attention_state_onto_the_generic_schema(): + contract = adapter.build_contract(PolicyRole.ROLLOUT, _effective()) + + assert contract.precision.compute is Precision.BF16 + assert contract.precision.accumulate is Precision.FP32 + assert contract.precision.softmax_accumulate is Precision.FP32 + assert contract.precision.downcast_at is DowncastPoint.FINAL_WRITE + assert contract.collectives[0].op is CollectiveOp.POINT_TO_POINT + assert contract.collectives[0].reduction_order is ReductionOrder.GLOBAL_BLOCK_INDEX + assert len(contract.extra["split_kv_coordinates"]) == 4 + assert contract.extra["split_kv_boundaries"][0][1] == ((0, 2), (2, 4)) + assert contract.extra["cp_block_manifest"][1][1:3] == (4, 8) + assert contract.extra["lse_domain"] == "attention" + + +def test_missing_actual_plan_is_not_filled_from_requested_policy(): + config = _effective() + del config["attn.actual_split_kv_plan_set"] + config["attn.requested_split_kv_policy"] = "fixed" + config["attn.requested_split_kv_size"] = 2 + + contract = adapter.build_contract(PolicyRole.TRAINING, config) + assert contract.extra["requested_split_kv_policy"] == "fixed" + assert "split_kv_runtime_plan_set" not in contract.extra + assert "split_kv_boundaries" not in contract.extra + + +def test_incomplete_or_rank_variant_plan_set_fails_closed(): + missing = _plan_set() + missing["entries"].pop() + with pytest.raises(adapter.AttentionAdapterError, match="coverage is incomplete"): + adapter.build_contract( + PolicyRole.TRAINING, + _effective(**{"attn.actual_split_kv_plan_set": missing}), + ) + + rank_variant = _plan_set() + rank_variant["entries"][2]["actual_split_boundaries"] = [[0, 4]] + with pytest.raises(adapter.AttentionAdapterError, match="differs across TP/CP consumers"): + adapter.build_contract( + PolicyRole.TRAINING, + _effective(**{"attn.actual_split_kv_plan_set": rank_variant}), + ) + + +def test_reported_split_count_must_match_actual_boundaries(): + plan = _plan_set() + plan["entries"][0]["actual_split_kv_count"] = 3 + with pytest.raises(adapter.AttentionAdapterError, match="must equal"): + adapter.build_contract( + PolicyRole.TRAINING, + _effective(**{"attn.actual_split_kv_plan_set": plan}), + ) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("split_kv_accum_dtype", "bf16", "accumulate in fp32"), + ("split_kv_downcast_at", "per_partial", "only at final_write"), + ("split_kv_merge_order", "arrival", "global_block_index"), + ], +) +def test_invalid_split_kv_numerical_contract_is_rejected(field, value, message): + plan = _plan_set() + plan["entries"][0][field] = value + with pytest.raises(adapter.AttentionAdapterError, match=message): + adapter.build_contract( + PolicyRole.ROLLOUT, + _effective(**{"attn.actual_split_kv_plan_set": plan}), + ) + + +def test_non_fp32_attention_accumulation_is_rejected_before_comparison(): + with pytest.raises(adapter.AttentionAdapterError, match="must accumulate in fp32"): + adapter.build_contract( + PolicyRole.TRAINING, + _effective(**{"attn.accumulate_dtype": "bf16"}), + ) + + +def test_cp_manifest_must_be_complete_and_gap_free(): + manifest = _manifest() + manifest[1]["kv_block_start"] = 5 + with pytest.raises(adapter.AttentionAdapterError, match="gap-free"): + adapter.build_contract( + PolicyRole.ROLLOUT, + _effective(**{"attn.cp_block_manifest": manifest}), + ) + + +def test_reference_cp_switch_builds_the_p2p_contract_without_claiming_a_runtime_plan(): + config = { + "attn.cp_world_size": 2, + "attn.cp_merge": CP_MERGE_REFERENCE.name, + } + contract = adapter.build_contract(PolicyRole.TRAINING, config) + assert contract.collectives[0].backend == "p2p_nccl_reference" + assert contract.collectives[0].reduction_order is ReductionOrder.GLOBAL_BLOCK_INDEX + assert "split_kv_runtime_plan_set" not in contract.extra + + +def test_role_specific_downcast_supports_a_training_only_ablation(): + config = _effective(**{"attn.training_downcast_at": "per_partial"}) + training = adapter.build_contract(PolicyRole.TRAINING, config) + rollout = adapter.build_contract(PolicyRole.ROLLOUT, config) + assert training.precision.downcast_at is DowncastPoint.PER_PARTIAL + assert rollout.precision.downcast_at is DowncastPoint.FINAL_WRITE + + +def test_readback_accepts_mapping_reader_and_attribute_but_rejects_requested_only(): + assert adapter.read_effective_config(PolicyRole.TRAINING, _effective())["attn.batch_size"] == 1 + + class Engine: + role = PolicyRole.ROLLOUT + + def read_effective_config(self): + return {"attn.cp_world_size": 2} + + assert adapter.read_effective_config(PolicyRole.ROLLOUT, Engine()) == { + "attn.cp_world_size": 2 + } + + class Bare: + effective_config = {"attn.compute_dtype": "bf16"} + + assert adapter.read_effective_config(PolicyRole.TRAINING, Bare()) == { + "attn.compute_dtype": "bf16" + } + + with pytest.raises(adapter.AttentionAdapterError, match="requested_config"): + adapter.read_effective_config( + PolicyRole.TRAINING, {"requested_config": {"attn.cp_world_size": 2}} + ) + + +def test_observed_collective_comes_from_effective_runtime_state(): + observed = adapter.observe_collectives(PolicyRole.ROLLOUT, _effective()) + assert observed[0].backend == "p2p_nccl_reference" + + config = _effective() + del config["attn.cp_collective"] + assert adapter.observe_collectives(PolicyRole.ROLLOUT, config) == () + + +def test_implementation_resolution_has_a_trace_for_every_failed_candidate(): + impl, resolution = adapter.resolve_implementation( + "attn.split_kv", PolicyRole.TRAINING, "does.not.exist.Missing" + ) + assert impl is None + assert resolution.resolved is None + assert resolution.rejected + + impl, resolution = adapter.resolve_implementation( + "attn.split_kv", PolicyRole.TRAINING, "math.sqrt" + ) + assert impl(9.0) == 3.0 + assert resolution.resolved == "math.sqrt" + + +def test_plugin_wires_adapter_and_discovers_all_attention_factors(): + assert AttentionChecks.build_contract is adapter.build_contract + ids = [factor.id for factor in AttentionChecks().declare_factors()] + assert ids == [ + "attn.cp_merge", + "attn.precision_downcast", + "attn.rope_fusion", + "attn.split_kv", + ] + + +def test_plan_helper_is_not_mutated_by_contract_building(): + plan = _plan_set() + before = deepcopy(plan) + adapter.build_contract( + PolicyRole.ROLLOUT, + _effective(**{"attn.actual_split_kv_plan_set": plan}), + ) + assert plan == before + + +def test_attention_extra_is_json_serializable_for_report_artifacts(): + contract = adapter.build_contract(PolicyRole.ROLLOUT, _effective()) + json.dumps(contract.extra) diff --git a/tests/test_mismatch_attention_factors.py b/tests/test_mismatch_attention_factors.py new file mode 100644 index 00000000..877ec954 --- /dev/null +++ b/tests/test_mismatch_attention_factors.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Attention factor declarations and contract-comparison behavior.""" + +from __future__ import annotations + +from copy import deepcopy + +from rl_engine.mismatch.operator_checks.attention import adapter +from rl_engine.mismatch.operator_checks.attention.factors.cp_merge import FACTOR as CP_FACTOR +from rl_engine.mismatch.operator_checks.attention.factors.precision_downcast import ( + FACTOR as PRECISION_FACTOR, +) +from rl_engine.mismatch.operator_checks.attention.factors.rope_fusion import FACTOR as ROPE_FACTOR +from rl_engine.mismatch.operator_checks.attention.factors.split_kv import FACTOR as SPLIT_FACTOR +from rl_engine.mismatch.pipeline import ( + build_variants, + compare_contracts, + reject_contradictory_factors, +) +from rl_engine.mismatch.schema import ComparisonIssueCode, PolicyRole +from tests.test_mismatch_attention_adapter import _effective, _manifest, _plan_set + + +def _compare(config_rollout: dict, config_training: dict, factor): + return compare_contracts( + adapter.build_contract(PolicyRole.ROLLOUT, config_rollout), + adapter.build_contract(PolicyRole.TRAINING, config_training), + (factor,), + ) + + +def test_equal_requested_policy_but_different_actual_boundaries_is_a_finding(): + rollout = _effective() + training = _effective() + training["attn.actual_split_kv_plan_set"] = _plan_set( + boundaries={0: [[0, 4]], 1: [[4, 8]]} + ) + + issues = _compare(rollout, training, SPLIT_FACTOR) + paths = {issue.field_path for issue in issues} + assert "extra.split_kv_boundaries" in paths + assert "extra.split_kv_runtime_plan_set" in paths + assert all( + issue.code is ComparisonIssueCode.SEMANTIC_MISMATCH + for issue in issues + if issue.field_path in paths + ) + + +def test_missing_runtime_plan_is_reported_as_missing_not_clean(): + rollout = _effective() + training = _effective() + del training["attn.actual_split_kv_plan_set"] + + issues = _compare(rollout, training, SPLIT_FACTOR) + missing = { + issue.field_path + for issue in issues + if issue.code is ComparisonIssueCode.REQUIRED_FIELD_MISSING + } + assert "extra.split_kv_runtime_plan_set" in missing + assert "extra.split_kv_boundaries" in missing + + +def test_split_kv_backend_and_trace_source_are_provenance_only(): + rollout = _effective() + training = _effective() + plan = training["attn.actual_split_kv_plan_set"] + for entry in plan["entries"]: + entry["split_kv_backend"] = "different-but-equivalent-backend" + entry["split_kv_plan_source"] = "different-runtime-hook" + + assert _compare(rollout, training, SPLIT_FACTOR) == () + + +def test_merge_order_and_collective_path_mismatches_are_visible(): + rollout = _effective() + training = _effective() + training["attn.cp_collective"] = { + **training["attn.cp_collective"], + "op": "all_gather", + "reduction_order": "nccl_algorithm", + "determinism": "none", + "backend": "nccl", + } + + issues = _compare(rollout, training, CP_FACTOR) + paths = {issue.field_path for issue in issues} + assert "collectives[0].op" in paths + assert "collectives[0].reduction_order" in paths + assert "collectives[0].backend" not in paths + + +def test_cp_owner_manifest_mismatch_voids_the_comparison_identity(): + rollout = _effective() + training = _effective() + manifest = deepcopy(_manifest()) + manifest[0]["owner_cp_rank"] = 1 + manifest[1]["owner_cp_rank"] = 0 + training["attn.cp_block_manifest"] = manifest + + issues = _compare(rollout, training, CP_FACTOR) + by_path = {issue.field_path: issue for issue in issues} + assert by_path["extra.cp_block_manifest"].code is ComparisonIssueCode.BITWISE_MISMATCH + + +def test_downcast_and_compute_dtype_mismatch_are_separate_findings(): + rollout = _effective() + training = _effective( + **{ + "attn.compute_dtype": "fp16", + "attn.training_downcast_at": "per_partial", + } + ) + issues = _compare(rollout, training, PRECISION_FACTOR) + by_path = {issue.field_path: issue for issue in issues} + assert by_path["precision.compute"].code is ComparisonIssueCode.BITWISE_MISMATCH + assert by_path["precision.downcast_at"].code is ComparisonIssueCode.SEMANTIC_MISMATCH + + +def test_rope_position_theta_and_post_qk_state_are_compared(): + rollout = _effective() + training = _effective( + **{ + "attn.rope_theta": 10_000.0, + "attn.position_ids_digest": "positions:other", + "attn.post_rope_qk_digest": "qk:other", + } + ) + issues = _compare(rollout, training, ROPE_FACTOR) + assert {issue.field_path for issue in issues} >= { + "extra.rope_theta", + "extra.position_ids_digest", + "extra.post_rope_qk_digest", + } + + +def test_reference_factors_expand_to_four_arms_and_static_checks_pass(): + reject_contradictory_factors((SPLIT_FACTOR, CP_FACTOR, ROPE_FACTOR, PRECISION_FACTOR)) + for factor in (SPLIT_FACTOR, CP_FACTOR, ROPE_FACTOR): + assert [variant.name for variant in build_variants(factor)] == [ + "both_native", + "both_reference", + "training_reference_only", + "rollout_reference_only", + ] + assert [variant.name for variant in build_variants(PRECISION_FACTOR)] == [ + "value_final_write", + "value_per_partial", + "value_per_block", + ] diff --git a/tests/test_mismatch_framework.py b/tests/test_mismatch_framework.py new file mode 100644 index 00000000..58f11936 --- /dev/null +++ b/tests/test_mismatch_framework.py @@ -0,0 +1,883 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Framework tests. No operator plugins involved. + +The operators are written separately, so everything here uses fixture factors: +what is under test is the framework's own logic -- the gates, the matrix, the +ordering, the conflict detection. + +The failure modes these lock down are the ones that make an attribution +framework worse than useless: a silently ignored switch read as "nothing here", +a missing shard read as a clean number, a broken reference quietly steering every +conclusion. +""" + +from __future__ import annotations + +import pytest + +from rl_engine.mismatch.pipeline import ( + ContradictoryFactor, + PluginRegistry, + RegistrationError, + RunContext, + build_report, + build_variants, + compare_contracts, + compute_metrics, + diagnose, + expand_repeats, + missing_prerequisites, + order_cases_by_rebind_cost, + reject_contradictory_factors, + render_summary, + run_variant, +) +from rl_engine.mismatch.schema import ( + BatchPlacement, + CollectiveContract, + CollectiveOp, + ComparisonIdentity, + ComparisonIssueCode, + ComparisonRule, + DeterminismLevel, + Diagnosis, + DowncastPoint, + DynamicSamplingDecision, + Evidence, + ExecutionPath, + ExpectedOutcome, + FactorCategory, + FactorVariant, + FailureMode, + ImplementationResolution, + KnownPitfall, + LogprobShard, + MismatchFactor, + MismatchMetrics, + NoiseFloor, + OperatorContract, + ParallelDim, + PolicyRole, + Precision, + PrecisionProfile, + Prerequisites, + RebindCost, + ReductionOrder, + ReferenceAuthority, + ReferenceImplementation, + RejectedCandidate, + ReuseKey, + RolloutGroup, + Switch, + SwitchStatus, + VariantResult, + expected_range, + is_silent_failure, + reuse_level, + tolerance_floor, +) +from tests.mismatch_cpu_backend import CpuScoringBackend + +# ---------------------------------------------------------------- fixtures -- + + +def make_identity(tokens: int = 8) -> ComparisonIdentity: + return ComparisonIdentity( + prompt_token_ids=tuple(range(tokens // 2)), + response_token_ids=tuple(range(tokens)), + active_mask=tuple([False] * (tokens // 2) + [True] * (tokens - tokens // 2)), + position_ids=tuple(range(tokens)), + checkpoint_id="fixture/model", + checkpoint_revision="deadbeef", + model_shape="L=1,H=8,Hq=2,Hkv=1,D=4", + group=RolloutGroup(prompt_id="p0", rollout_ids=("r0",), group_size=1), + batch_placement=BatchPlacement( + data_parallel_rank=0, microbatch_index=0, position_in_microbatch=0 + ), + sampling_decision=DynamicSamplingDecision(kept=True), + ) + + +def make_reference(name: str = "fixture_ref", **kwargs) -> ReferenceImplementation: + return ReferenceImplementation( + name=name, + tier=ReferenceAuthority.SELF_WRITTEN, + training_impl=f"{name}.training", + rollout_impl=f"{name}.rollout", + covers_paths=( + ExecutionPath.TRAINING_FULL_PREFILL, + ExecutionPath.ROLLOUT_FULL_PREFILL, + ), + **kwargs, + ) + + +def make_factor( + factor_id: str = "fixture.swap", + *, + reference: ReferenceImplementation | None = None, + rules: dict[str, ComparisonRule] | None = None, + prerequisites: Prerequisites | None = None, + required_evidence: tuple[str, ...] = (), + rebind_cost: RebindCost = RebindCost.PER_REQUEST, + allowed_values: tuple = ("native", "fixture_ref"), +) -> MismatchFactor: + return MismatchFactor( + id=factor_id, + operator=factor_id.split(".")[0], + category=FactorCategory.SHARDING_AND_REDUCTION, + question="fixture", + switch=Switch( + path=factor_id, + rebind_cost=rebind_cost, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=allowed_values, + ), + comparison_rules=rules if rules is not None else {}, + prerequisites=prerequisites or Prerequisites(), + required_evidence=required_evidence, + reference=reference if reference is not None else make_reference(), + ) + + +def make_contract(role: PolicyRole, **extra) -> OperatorContract: + return OperatorContract( + operator="fixture", + role=role, + precision=PrecisionProfile( + compute=Precision.BF16, + accumulate=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + ), + collectives=( + CollectiveContract( + op=CollectiveOp.ALL_REDUCE, + group=ParallelDim.TENSOR, + group_size=2, + reduction_order=extra.pop("reduction_order", ReductionOrder.GLOBAL_RANK_INDEX), + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=extra.pop("determinism", DeterminismLevel.STABLE_ACROSS_RUNS), + backend="fixture", + ), + ), + extra=extra, + ) + + +def make_result( + name: str, + *, + dlogp_mean: float = 0.0, + dlogp_max: float = 0.0, + clip_fraction: float = 0.0, + status: SwitchStatus = SwitchStatus.APPLIED, + evidence: frozenset[str] = frozenset(), + expected: ExpectedOutcome = ExpectedOutcome.MEASURE_ONLY, + shards: tuple[LogprobShard, ...] = (), + resolution: ImplementationResolution | None = None, +) -> VariantResult: + return VariantResult( + variant=FactorVariant(name=name, switch_values={}, expected=expected), + path=ExecutionPath.TRAINING_FULL_PREFILL, + status=status, + metrics=MismatchMetrics( + active_token_count=4, + dlogp_mean=dlogp_mean, + dlogp_p99=dlogp_max, + dlogp_max=dlogp_max, + ratio_mean=1.0, + ratio_max=1.0, + clip_fraction=clip_fraction, + approx_kl=0.0, + ), + evidence=evidence, + effective_config={}, + logprob_shards=shards, + resolution=resolution, + ) + + +def four_arms(**overrides) -> list[VariantResult]: + """A standard four-arm set where nothing is wrong.""" + + defaults = { + "both_native": {"dlogp_mean": 0.02, "clip_fraction": 0.2}, + "both_reference": {"dlogp_max": 0.0, "expected": ExpectedOutcome.BITWISE_IDENTICAL}, + "training_reference_only": {"dlogp_mean": 0.02, "clip_fraction": 0.2}, + "rollout_reference_only": {"dlogp_mean": 0.02, "clip_fraction": 0.2}, + } + for name, patch in overrides.items(): + defaults[name] = {**defaults.get(name, {}), **patch} + return [make_result(name, **kwargs) for name, kwargs in defaults.items()] + + +# ------------------------------------------------------------ variant plan -- + + +def test_swap_factor_expands_to_four_arms_not_two(): + """A single swap cannot attribute a side. + + Only a one-sided swap says which side is at fault, and only the two-sided + swap proves the reference itself is sound -- so the standard set is four. + """ + + variants = build_variants(make_factor()) + assert [v.name for v in variants] == [ + "both_native", + "both_reference", + "training_reference_only", + "rollout_reference_only", + ] + both = next(v for v in variants if v.name == "both_reference") + assert both.expected is ExpectedOutcome.BITWISE_IDENTICAL + assert both.replace_on == { + PolicyRole.ROLLOUT: "fixture_ref.rollout", + PolicyRole.TRAINING: "fixture_ref.training", + } + + +def test_factor_without_reference_is_a_value_sweep(): + """No reference implementation means it is a parameter sweep. + + The distinction is derived from ``reference is None`` rather than stored in a + separate field -- derivable state is state that can disagree with itself. + """ + + factor = make_factor(reference=None, allowed_values=(1, 2, 4)) + factor = MismatchFactor(**{**factor.__dict__, "reference": None}) + variants = build_variants(factor) + assert [v.name for v in variants] == ["value_1", "value_2", "value_4"] + + +def test_fp64_oracle_arm_is_added_when_declared(): + factor = make_factor(reference=make_reference(fp64_oracle="fixture.fp64")) + assert "fp64_oracle" in [v.name for v in build_variants(factor)] + + +def test_cases_are_ordered_cheapest_rebuild_first(): + """160 cases in random order restart the process for nearly every one. + + Ordering is what makes the whole run finish, not a nicety. + """ + + cheap = make_factor("a.cheap", rebind_cost=RebindCost.PER_REQUEST) + expensive = make_factor("b.expensive", rebind_cost=RebindCost.PROCESS_RESTART) + middle = make_factor("c.middle", rebind_cost=RebindCost.ENGINE_REBUILD) + + cases = [ + (expensive, build_variants(expensive)[0]), + (cheap, build_variants(cheap)[0]), + (middle, build_variants(middle)[0]), + ] + ordered = order_cases_by_rebind_cost(cases) + assert [factor.id for factor, _ in ordered] == ["a.cheap", "c.middle", "b.expensive"] + + +# --------------------------------------------------------- static rejection -- + + +def test_topology_independence_claim_with_nccl_order_is_rejected_before_running(): + """Claiming topology independence while reducing by NCCL's choice is + self-contradictory -- reject at planning time, not after burning machine time.""" + + from rl_engine.mismatch.schema import RequiredSetting, SettingChannel + + contradictory = CollectiveContract( + op=CollectiveOp.ALL_REDUCE, + group=ParallelDim.TENSOR, + group_size=4, + reduction_order=ReductionOrder.NCCL_ALGORITHM, + accumulate_precision=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + determinism=DeterminismLevel.STABLE_ACROSS_TOPOLOGY, + backend="nccl", + ) + factor = make_factor( + reference=make_reference( + required_settings=( + RequiredSetting( + key="collective", + value=contradictory, + channel=SettingChannel.CALL_ARG, + ), + ) + ) + ) + + with pytest.raises(ContradictoryFactor, match="stable_across_topology"): + reject_contradictory_factors([factor]) + + +def test_prerequisites_report_what_is_missing_not_a_yes_no(): + """A whitelist probe returning the missing items, not an opaque boolean.""" + + factor = make_factor( + prerequisites=Prerequisites( + required_ops=("rl_kernel.reduce_scatter",), + min_gpu_count=2, + blocked_by=("#247",), + ) + ) + unmet = missing_prerequisites(factor, available_ops=frozenset(), gpu_count=0) + reasons = " ".join(item.reason for item in unmet) + assert "reduce_scatter" in reasons + assert "2 devices" in reasons + assert "#247" in reasons + + +# ------------------------------------------------------------- comparison --- + + +def test_record_only_fields_are_never_compared(): + """RECORD_ONLY exists so structural differences do not drown real findings. + + Packed QKV differs in form between the two sides while the arithmetic is + identical; comparing it would turn every run red. + """ + + rollout = make_contract(PolicyRole.ROLLOUT, qkv_layout="packed") + training = make_contract(PolicyRole.TRAINING, qkv_layout="split") + factor = make_factor(rules={"extra.qkv_layout": ComparisonRule.RECORD_ONLY}) + + assert compare_contracts(rollout, training, [factor]) == () + + +def test_semantic_mismatch_is_reported_with_both_sides_values(): + rollout = make_contract(PolicyRole.ROLLOUT, rope_theta=10000.0) + training = make_contract(PolicyRole.TRAINING, rope_theta=1000000.0) + factor = make_factor(rules={"extra.rope_theta": ComparisonRule.MUST_MATCH_SEMANTICALLY}) + + issues = compare_contracts(rollout, training, [factor]) + assert len(issues) == 1 + assert issues[0].code is ComparisonIssueCode.SEMANTIC_MISMATCH + assert issues[0].values[PolicyRole.ROLLOUT] == 10000.0 + assert issues[0].values[PolicyRole.TRAINING] == 1000000.0 + + +def test_missing_field_is_reported_rather_than_raising(): + factor = make_factor(rules={"extra.absent": ComparisonRule.MUST_MATCH_BITWISE}) + issues = compare_contracts( + make_contract(PolicyRole.ROLLOUT), make_contract(PolicyRole.TRAINING), [factor] + ) + assert issues[0].code is ComparisonIssueCode.REQUIRED_FIELD_MISSING + + +def test_indexed_path_reaches_into_collectives(): + rollout = make_contract(PolicyRole.ROLLOUT, reduction_order=ReductionOrder.ARRIVAL) + training = make_contract(PolicyRole.TRAINING, reduction_order=ReductionOrder.GLOBAL_RANK_INDEX) + factor = make_factor( + rules={"collectives[0].reduction_order": ComparisonRule.MUST_MATCH_SEMANTICALLY} + ) + issues = compare_contracts(rollout, training, [factor]) + assert issues[0].field_path == "collectives[0].reduction_order" + + +def test_one_side_promising_more_determinism_is_flagged(): + """Comparing a topology-independent implementation against a + non-reproducible one measures the weaker side's noise, not the gap.""" + + rollout = make_contract(PolicyRole.ROLLOUT, determinism=DeterminismLevel.STABLE_ACROSS_TOPOLOGY) + training = make_contract(PolicyRole.TRAINING, determinism=DeterminismLevel.NONE) + issues = compare_contracts(rollout, training, [make_factor()]) + assert any(i.code is ComparisonIssueCode.DETERMINISM_INCOMPATIBLE for i in issues) + + +# ---------------------------------------------------------------- metrics --- + + +def test_identical_logprobs_give_zero_clip_fraction(): + metrics = compute_metrics([-1.0, -2.0], [-1.0, -2.0], [True, True]) + assert metrics.dlogp_max == 0.0 + assert metrics.clip_fraction == 0.0 + + +def test_clip_fraction_counts_tokens_past_the_grpo_clip_edge(): + """Past ln(1+eps) a token is clipped and its gradient signal is discarded. + + That is the actual mechanism by which mismatch breaks training, which is why + this and not the mean is the headline number. + """ + + # 0.30 > ln(1.2) = 0.182, so one of the two is clipped. + metrics = compute_metrics([-1.0, -1.0], [-0.70, -1.05], [True, True]) + assert metrics.clip_fraction == 0.5 + + +def test_inactive_tokens_are_excluded(): + metrics = compute_metrics([-1.0, -1.0], [-9.0, -1.0], [False, True]) + assert metrics.active_token_count == 1 + assert metrics.dlogp_max == 0.0 + + +def test_mean_within_band_but_tail_past_the_edge_is_a_silent_failure(): + """The most common false negative: the headline looks healthy while gradient + signal is being thrown away.""" + + metrics = MismatchMetrics( + active_token_count=100, + dlogp_mean=0.004, # squarely inside the dense production band + dlogp_p99=0.25, + dlogp_max=0.4, + ratio_mean=1.0, + ratio_max=1.5, + clip_fraction=0.02, + approx_kl=0.001, + ) + assert is_silent_failure(metrics) + + +# ------------------------------------------------------------- thresholds --- + + +def test_low_floor_expects_bitwise_not_the_production_band(): + """Reading the anchor floor against the production table hides real bugs. + + At the anchor floor the expectation is bitwise; judging it by 0.002-0.008 + would call a definite operator error "normal". + """ + + production = expected_range("dense", NoiseFloor.PRODUCTION) + anchor = expected_range("dense", NoiseFloor.SINGLE_LAYER_ANCHOR) + assert production.dlogp_mean == (0.002, 0.008) + assert anchor.dlogp_mean == (0.0, 0.0) + assert tolerance_floor("dense", NoiseFloor.SINGLE_LAYER_ANCHOR) < 1e-5 + + +def test_large_moe_band_is_wider_and_not_a_bug(): + band = expected_range("large_moe", NoiseFloor.PRODUCTION, routing_replay=False) + assert band.dlogp_mean == (0.01, 0.03) + assert "do not file it as a bug" in band.note + + +# ---------------------------------------------------------------- the gates -- + + +def test_a_variant_that_did_not_apply_never_reaches_the_matrix(): + """A silently reverted switch reads as "the deviation did not change", which + reads as NOT_THIS_FACTOR. A false negative that looks like a clean result.""" + + results = four_arms() + results[2] = make_result( + "training_reference_only", + status=SwitchStatus.FELL_BACK, + resolution=ImplementationResolution( + requested="fixture_ref.training", + resolved=None, + rejected=(RejectedCandidate(name="fixture_ref.training", reason="library missing"),), + ), + ) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.VARIANT_DID_NOT_APPLY + assert "library missing" in report.diagnosis_reason + + +def test_missing_evidence_is_not_the_same_as_nothing_found(): + factor = make_factor(required_evidence=(Evidence.MODEL_STATE_FINGERPRINT.value,)) + report = diagnose(factor, four_arms(), noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.INSUFFICIENT_EVIDENCE + + +def test_incomplete_logprob_shards_block_the_verdict(): + """Under TP/CP each rank holds one slice. One slice short and the LSE + denominator loses a chunk, so logp comes out systematically high -- wrong in + a way that does not show.""" + + partial = (LogprobShard(rank=0, world_size=4, selected_logprobs=[0.0]),) + results = four_arms() + results[0] = make_result("both_native", dlogp_mean=0.02, clip_fraction=0.2, shards=partial) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.INSUFFICIENT_EVIDENCE + assert "1 of 4" in report.diagnosis_reason + + +def test_failed_pitfall_guard_blocks_the_verdict(): + guard = KnownPitfall( + id="rope_hook_not_covered", + mode=FailureMode.MISSING_INSTRUMENTATION, + symptom="RoPE looks identical", + actual_cause="the hook never captured it", + guard="dump post-RoPE Q/K on both sides", + guard_runs_at=NoiseFloor.SINGLE_LAYER_ANCHOR, + ) + report = diagnose( + make_factor(), four_arms(), noise_floor=NoiseFloor.PRODUCTION, failed_guards=[guard] + ) + assert report.diagnosis is Diagnosis.INSUFFICIENT_EVIDENCE + assert "rope_hook_not_covered" in report.diagnosis_reason + + +# ------------------------------------------------------- the matrix proper -- + + +def test_broken_reference_voids_the_factor_rather_than_attributing_a_side(): + """Without this gate, one wrong reference quietly steers every attribution -- + worse than having no framework at all.""" + + results = four_arms(both_reference={"dlogp_max": 0.5}) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.REFERENCE_ITSELF_IS_BROKEN + + +def test_only_training_side_converging_attributes_the_training_side(): + results = four_arms(training_reference_only={"dlogp_mean": 0.0, "clip_fraction": 0.0}) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.CAUSED_BY_TRAINING_SIDE + + +def test_only_rollout_side_converging_attributes_the_rollout_side(): + results = four_arms(rollout_reference_only={"dlogp_mean": 0.0, "clip_fraction": 0.0}) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.CAUSED_BY_ROLLOUT_SIDE + + +def test_both_sides_converging_leaves_the_reference_as_the_only_anchor(): + results = four_arms( + training_reference_only={"dlogp_mean": 0.0, "clip_fraction": 0.0}, + rollout_reference_only={"dlogp_mean": 0.0, "clip_fraction": 0.0}, + ) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.CAUSED_BY_BOTH_SIDES + + +def test_neither_side_moving_means_look_upstream(): + report = diagnose(make_factor(), four_arms(), noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.NOT_THIS_FACTOR + + +def test_convergence_is_judged_on_clip_fraction_not_the_mean(): + """At every production floor the mean sits far below the clip edge, so + judging on the mean would mark almost every factor NOT_THIS_FACTOR.""" + + results = four_arms( + both_native={"dlogp_mean": 0.005, "clip_fraction": 0.30}, + training_reference_only={"dlogp_mean": 0.005, "clip_fraction": 0.01}, + ) + report = diagnose(make_factor(), results, noise_floor=NoiseFloor.PRODUCTION) + assert report.diagnosis is Diagnosis.CAUSED_BY_TRAINING_SIDE + + +# ------------------------------------------------------------- repeats ------ + + +def test_repeat_under_expands_to_the_cartesian_product(): + """The one exception to "one variant, one execution". + + It verifies the self-check gate's own premise: both_reference can only anchor + if the fixed-order implementation really did fix the order. + """ + + variant = FactorVariant( + name="both_reference", + switch_values={}, + repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")}, + ) + assert len(expand_repeats(variant)) == 4 + + +def test_variant_without_repeats_runs_once(): + assert expand_repeats(FactorVariant(name="x", switch_values={})) == ({},) + + +# ------------------------------------------------------------- registry ---- + + +def test_duplicate_factor_id_is_rejected_at_registration(): + registry = PluginRegistry() + + class First: + operator = "first" + + def declare_factors(self): + return (make_factor("shared.id"),) + + class Second: + operator = "second" + + def declare_factors(self): + return (make_factor("shared.id"),) + + registry.register(First) + with pytest.raises(RegistrationError, match="duplicate factor id"): + registry.register(Second) + + +def test_same_contract_field_at_two_different_rules_is_rejected(): + """One field cannot be bitwise-required in one operator and record-only in + another; the comparison would depend on which factor happened to run.""" + + registry = PluginRegistry() + + class Strict: + operator = "strict" + + def declare_factors(self): + return ( + make_factor("strict.a", rules={"extra.shared": ComparisonRule.MUST_MATCH_BITWISE}), + ) + + class Loose: + operator = "loose" + + def declare_factors(self): + return (make_factor("loose.b", rules={"extra.shared": ComparisonRule.RECORD_ONLY}),) + + registry.register(Strict) + with pytest.raises(RegistrationError, match="declared as"): + registry.register(Loose) + + +def test_duplicate_factor_inside_one_plugin_is_rejected(): + registry = PluginRegistry() + + class Broken: + operator = "broken" + + def declare_factors(self): + return (make_factor("broken.same"), make_factor("broken.same")) + + with pytest.raises(RegistrationError, match="duplicate factor id"): + registry.register(Broken) + + +def test_duplicate_switch_path_inside_one_plugin_is_rejected(): + registry = PluginRegistry() + first = make_factor("broken.first") + second = make_factor("broken.second") + second = MismatchFactor( + **{ + **second.__dict__, + "switch": Switch( + path=first.switch.path, + rebind_cost=RebindCost.PER_REQUEST, + applies_to=(PolicyRole.ROLLOUT, PolicyRole.TRAINING), + allowed_values=("native", "fixture_ref"), + ), + } + ) + + class Broken: + operator = "broken" + + def declare_factors(self): + return (first, second) + + with pytest.raises(RegistrationError, match="duplicate switch path"): + registry.register(Broken) + + +def test_registry_starts_empty_because_operators_ship_separately(): + assert PluginRegistry().operators() == () + + +# ------------------------------------------------------------ fingerprints -- + + +def test_reuse_level_returns_the_coarsest_thing_that_changed(): + base = ReuseKey(process="p", process_group="g", engine="e", request="r") + assert reuse_level(base, base) is RebindCost.PER_REQUEST + assert reuse_level(base, ReuseKey("p", "g", "e2", "r")) is RebindCost.ENGINE_REBUILD + assert reuse_level(base, ReuseKey("p", "g2", "e", "r")) is RebindCost.PROCESS_GROUP_REBUILD + assert reuse_level(base, ReuseKey("p2", "g", "e", "r")) is RebindCost.PROCESS_RESTART + + +# ---------------------------------------------------------------- runner ---- + + +class _Checks: + """Minimal plugin used to drive the runner. Not a real operator.""" + + operator = "fixture" + + def __init__(self, resolvable: bool = True): + self.resolvable = resolvable + + def declare_factors(self): + return (make_factor(),) + + def build_contract(self, role, switch_values): + return make_contract(role) + + def read_effective_config(self, role, adapter): + return {} + + def observe_collectives(self, role, adapter): + return () + + def resolve_implementation(self, factor_id, role, impl_name): + if not self.resolvable: + return None, ImplementationResolution( + requested=impl_name, + resolved=None, + rejected=(RejectedCandidate(name=impl_name, reason="not built on this host"),), + ) + return (lambda *a, **k: None), ImplementationResolution( + requested=impl_name, resolved=impl_name + ) + + +def test_runner_detects_a_one_sided_injected_deviation(): + """End to end on CPU: bias one side and the metrics must see exactly that.""" + + identity = make_identity() + backends = { + PolicyRole.ROLLOUT: CpuScoringBackend(role=PolicyRole.ROLLOUT), + PolicyRole.TRAINING: CpuScoringBackend(role=PolicyRole.TRAINING, bias=0.25), + } + factor = make_factor() + variant = build_variants(factor)[0] # both_native + + result = run_variant(factor, variant, _Checks(), backends, RunContext(identity=identity)) + assert result.status is SwitchStatus.APPLIED + assert result.metrics.dlogp_mean == pytest.approx(0.25, abs=1e-9) + + +def test_runner_builds_contracts_from_effective_readback_not_requested_values(): + class ReadbackBackend(CpuScoringBackend): + actual: str + + def __init__(self, *, role, actual): + super().__init__(role=role) + self.actual = actual + + def score(self, role, identity, switch_values, replacement): + scores, readback = super().score(role, identity, switch_values, replacement) + return scores, {**readback, "fixture.actual": self.actual} + + class ReadbackChecks(_Checks): + def build_contract(self, role, switch_values): + return make_contract(role, actual=switch_values["fixture.actual"]) + + def read_effective_config(self, role, adapter): + return dict(adapter) + + identity = make_identity() + backends = { + PolicyRole.ROLLOUT: ReadbackBackend(role=PolicyRole.ROLLOUT, actual="runtime-a"), + PolicyRole.TRAINING: ReadbackBackend(role=PolicyRole.TRAINING, actual="runtime-b"), + } + factor = make_factor( + rules={"extra.actual": ComparisonRule.MUST_MATCH_SEMANTICALLY} + ) + variant = build_variants(factor)[0] + + result = run_variant( + factor, + variant, + ReadbackChecks(), + backends, + RunContext(identity=identity), + ) + assert result.comparison_issues[0].field_path == "extra.actual" + assert result.effective_config["rollout.fixture.actual"] == "runtime-a" + + +def test_reference_swap_removes_the_injected_deviation(): + """both_reference puts both sides on one implementation, so an injected + per-side bias must vanish -- this is the self-check gate working.""" + + identity = make_identity() + backends = { + PolicyRole.ROLLOUT: CpuScoringBackend(role=PolicyRole.ROLLOUT), + PolicyRole.TRAINING: CpuScoringBackend(role=PolicyRole.TRAINING, bias=0.25), + } + factor = make_factor() + both_reference = build_variants(factor)[1] + + result = run_variant(factor, both_reference, _Checks(), backends, RunContext(identity=identity)) + assert result.metrics.dlogp_max == 0.0 + + +def test_unresolvable_implementation_is_recorded_as_fell_back_with_a_trace(): + """A single fallback_reason string is not enough: you need to know which + candidates were tried and why each was rejected.""" + + identity = make_identity() + backends = { + PolicyRole.ROLLOUT: CpuScoringBackend(role=PolicyRole.ROLLOUT), + PolicyRole.TRAINING: CpuScoringBackend(role=PolicyRole.TRAINING), + } + factor = make_factor() + both_reference = build_variants(factor)[1] + + result = run_variant( + factor, both_reference, _Checks(resolvable=False), backends, RunContext(identity=identity) + ) + assert result.status is SwitchStatus.FELL_BACK + assert result.resolution.resolved is None + assert result.resolution.rejected[0].reason == "not built on this host" + + +def test_unstable_backend_fails_the_topology_independence_assertion(): + """A backend whose output moves with the environment must not pass as + topology independent.""" + + identity = make_identity() + backends = { + PolicyRole.ROLLOUT: CpuScoringBackend( + role=PolicyRole.ROLLOUT, unstable_under=frozenset({"NCCL_ALGO"}) + ), + PolicyRole.TRAINING: CpuScoringBackend(role=PolicyRole.TRAINING), + } + factor = make_factor() + variant = FactorVariant( + name="both_reference", + switch_values={}, + expected=ExpectedOutcome.BITWISE_IDENTICAL, + repeat_under={"NCCL_ALGO": ("Ring", "Tree")}, + ) + + result = run_variant(factor, variant, _Checks(), backends, RunContext(identity=identity)) + assert result.status is SwitchStatus.ERROR + + +# ---------------------------------------------------------------- report ---- + + +def test_report_ranks_hypotheses_and_summarises(): + from rl_engine.mismatch.model_meta import QWEN3_CORRESPONDENCES, QWEN3_EDGES + + attributed = diagnose( + make_factor("mlp.forward_reduce"), + four_arms(training_reference_only={"dlogp_mean": 0.0, "clip_fraction": 0.0}), + noise_floor=NoiseFloor.SHARDED_SINGLE_NODE, + ) + clean = diagnose( + make_factor("attn.rope_fusion"), + four_arms(), + noise_floor=NoiseFloor.SHARDED_SINGLE_NODE, + ) + + report = build_report( + [attributed, clean], + noise_floor=NoiseFloor.SHARDED_SINGLE_NODE, + correspondences=QWEN3_CORRESPONDENCES, + edges=QWEN3_EDGES, + ) + assert len(report.hypotheses) == 1 + assert report.hypotheses[0].rank == 1 + + summary = render_summary(report) + assert "caused_by_training_side: 1" in summary + assert "not_this_factor: 1" in summary + + +def test_only_proven_equivalences_filter_findings(): + """An equivalence without a test proving it is not trusted -- otherwise + "filtering false positives" quietly becomes "hiding real findings".""" + + from rl_engine.mismatch.pipeline import filter_known_equivalences + from rl_engine.mismatch.schema import ModuleCorrespondence + + unproven = ModuleCorrespondence( + semantic_name="mlp.gate_up", + training_module="t", + rollout_module="r", + equivalence="concat_on_dim0", + verified_by=None, + ) + kept, filtered = filter_known_equivalences([unproven], ["mlp.gate_up"]) + assert kept == ("mlp.gate_up",) + assert filtered == () diff --git a/tests/test_mismatch_logprob_adapter.py b/tests/test_mismatch_logprob_adapter.py new file mode 100644 index 00000000..0c225ad2 --- /dev/null +++ b/tests/test_mismatch_logprob_adapter.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Logprob operator plugin tests: the four adapter methods and the swap factor. + +The framework's own logic is locked down in ``test_mismatch_framework.py`` with +fixture factors; here the subject is the logprob plugin's declarations and how +they map WS2's TP-aware reduction semantics onto the mismatch schema. +""" + +from __future__ import annotations + +import pytest + +from rl_engine.mismatch.operator_checks.logprob import LogprobChecks, adapter +from rl_engine.mismatch.operator_checks.logprob._common import ( + DETERMINISTIC_LSE_REFERENCE, + QWEN3_PADDED_VOCAB, + even_vocab_shard_bounds, +) +from rl_engine.mismatch.operator_checks.logprob.factors.lse_merge_order import ( + FACTOR as LSE_MERGE_FACTOR, +) +from rl_engine.mismatch.pipeline import ( + build_variants, + compare_contracts, + reject_contradictory_factors, +) +from rl_engine.mismatch.schema import ( + CollectiveOp, + ComparisonIssueCode, + PolicyRole, + Precision, + ReductionOrder, +) + +TP2 = {"logp.tp_world_size": 2} + + +# ------------------------------------------------------------ build_contract -- + + +def test_native_contracts_map_each_framework_onto_the_schema(): + training = adapter.build_contract(PolicyRole.TRAINING, TP2) + rollout = adapter.build_contract(PolicyRole.ROLLOUT, TP2) + + assert training.collectives[0].op is CollectiveOp.ALL_REDUCE + assert training.collectives[0].backend == "nccl" + assert rollout.collectives[0].op is CollectiveOp.ALL_GATHER + assert rollout.collectives[0].backend == "vllm_custom_ipc" + for side in (training, rollout): + assert side.collectives[0].group_size == 2 + assert side.collectives[0].accumulate_precision is Precision.FP32 + assert side.collectives[0].reduction_order is ReductionOrder.NCCL_ALGORITHM + assert side.extra["vocab_shard_map"] == even_vocab_shard_bounds(QWEN3_PADDED_VOCAB, 2) + + +def test_tp1_contract_declares_no_collectives(): + contract = adapter.build_contract(PolicyRole.TRAINING, {"logp.tp_world_size": 1}) + assert contract.collectives == () + + +def test_reference_switch_pins_the_shard_order_merge_on_both_sides(): + switches = {**TP2, "logp.lse_merge": DETERMINISTIC_LSE_REFERENCE.name} + training = adapter.build_contract(PolicyRole.TRAINING, switches) + rollout = adapter.build_contract(PolicyRole.ROLLOUT, switches) + + for side in (training, rollout): + assert side.collectives[0].reduction_order is ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + assert side.collectives[0].backend == "rl_kernel" + assert side.extra["lse_export"] is True + + +def test_one_sided_swap_replaces_only_the_named_side(): + switches = {**TP2, "logp.lse_merge": f"{DETERMINISTIC_LSE_REFERENCE.name}@training"} + training = adapter.build_contract(PolicyRole.TRAINING, switches) + rollout = adapter.build_contract(PolicyRole.ROLLOUT, switches) + + assert training.collectives[0].reduction_order is ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + assert rollout.collectives[0].reduction_order is ReductionOrder.NCCL_ALGORITHM + + +def test_only_the_training_side_varies_the_head_dtype(): + switches = {**TP2, "logp.head_dtype": "fp32"} + training = adapter.build_contract(PolicyRole.TRAINING, switches) + rollout = adapter.build_contract(PolicyRole.ROLLOUT, switches) + + assert training.precision.lm_head is Precision.FP32 + assert rollout.precision.lm_head is Precision.BF16 # vLLM: the model dtype + + +def test_unknown_switch_values_fail_loudly(): + with pytest.raises(adapter.LogprobAdapterError, match="head_dtype"): + adapter.build_contract(PolicyRole.TRAINING, {"logp.head_dtype": "fp8"}) + with pytest.raises(adapter.LogprobAdapterError, match="lse_merge"): + adapter.build_contract(PolicyRole.TRAINING, {"logp.lse_merge": "fastest"}) + + +# ------------------------------------------------- comparison with the factor -- + + +def test_native_sides_disagree_semantically_on_the_merge(): + switches = {**TP2, "logp.lse_merge": "native"} + issues = compare_contracts( + adapter.build_contract(PolicyRole.ROLLOUT, switches), + adapter.build_contract(PolicyRole.TRAINING, switches), + (LSE_MERGE_FACTOR,), + ) + + codes = {issue.field_path: issue.code for issue in issues} + assert codes["collectives[0].op"] is ComparisonIssueCode.SEMANTIC_MISMATCH + # Identity fields agree, so the case is a finding rather than void. + assert "collectives[0].group_size" not in codes + assert "extra.vocab_shard_map" not in codes + + +def test_reference_on_both_sides_clears_every_contract_issue(): + switches = {**TP2, "logp.lse_merge": DETERMINISTIC_LSE_REFERENCE.name} + issues = compare_contracts( + adapter.build_contract(PolicyRole.ROLLOUT, switches), + adapter.build_contract(PolicyRole.TRAINING, switches), + (LSE_MERGE_FACTOR,), + ) + assert issues == () + + +def test_lse_merge_factor_expands_to_the_declared_four_arms_and_is_consistent(): + reject_contradictory_factors((LSE_MERGE_FACTOR,)) + names = [variant.name for variant in build_variants(LSE_MERGE_FACTOR)] + assert names == [ + "both_native", + "both_reference", + "training_reference_only", + "rollout_reference_only", + ] + + +# ------------------------------------------- read_effective_config / observe -- + + +def test_read_effective_config_accepts_the_three_adapter_shapes(): + as_mapping = adapter.read_effective_config(PolicyRole.TRAINING, {"logp.head_dtype": "fp32"}) + assert as_mapping == {"logp.head_dtype": "fp32"} + + class Engine: + role = PolicyRole.ROLLOUT + + def read_effective_config(self): + return {"logp.tp_world_size": 2} + + assert adapter.read_effective_config(PolicyRole.ROLLOUT, Engine()) == {"logp.tp_world_size": 2} + + class Bare: + effective_config = {"logp.lse_merge": "native"} + + assert adapter.read_effective_config(PolicyRole.TRAINING, Bare()) == { + "logp.lse_merge": "native" + } + + +def test_read_effective_config_rejects_an_adapter_playing_the_other_role(): + class Engine: + role = PolicyRole.ROLLOUT + effective_config = {} + + with pytest.raises(adapter.LogprobAdapterError, match="plays 'rollout'"): + adapter.read_effective_config(PolicyRole.TRAINING, Engine()) + + +def test_observe_collectives_reflects_the_effective_config_not_the_request(): + observed = adapter.observe_collectives( + PolicyRole.TRAINING, {"logp.tp_world_size": 2, "logp.lse_merge": "native"} + ) + assert len(observed) == 1 + assert observed[0].op is CollectiveOp.ALL_REDUCE + assert observed[0].group_size == 2 + + assert adapter.observe_collectives(PolicyRole.TRAINING, {"logp.tp_world_size": 1}) == () + + +# ------------------------------------------------------ resolve_implementation -- + + +def test_resolution_failure_carries_the_trace_not_a_bare_none(): + impl, resolution = adapter.resolve_implementation( + LSE_MERGE_FACTOR.id, + PolicyRole.TRAINING, + "rl_engine.kernels.ops.pytorch.loss.does_not_exist.MissingOp", + ) + assert impl is None + assert resolution.resolved is None + assert "import failed" in resolution.rejected[0].reason + + impl, resolution = adapter.resolve_implementation( + LSE_MERGE_FACTOR.id, PolicyRole.TRAINING, "math.no_such_attribute" + ) + assert impl is None + assert "no attribute" in resolution.rejected[0].reason + + impl, resolution = adapter.resolve_implementation( + LSE_MERGE_FACTOR.id, PolicyRole.TRAINING, "math.pi" + ) + assert impl is None + assert "not callable" in resolution.rejected[0].reason + + +def test_resolution_success_returns_the_callable_and_a_clean_trace(): + impl, resolution = adapter.resolve_implementation( + LSE_MERGE_FACTOR.id, PolicyRole.TRAINING, "math.sqrt" + ) + assert impl(4.0) == 2.0 + assert resolution.resolved == "math.sqrt" + assert resolution.rejected == () + + +def test_ws2_reference_path_either_resolves_or_leaves_a_trace(): + """On a tree without issue #241 PR3 the WS2 op is absent; either way the + resolution must be investigable, never a silent fallback.""" + + impl, resolution = adapter.resolve_implementation( + LSE_MERGE_FACTOR.id, PolicyRole.TRAINING, DETERMINISTIC_LSE_REFERENCE.training_impl + ) + if impl is None: + assert resolution.resolved is None + assert resolution.rejected + else: + assert callable(impl) + assert resolution.resolved == DETERMINISTIC_LSE_REFERENCE.training_impl + + +# ----------------------------------------------------------------- the plugin -- + + +def test_plugin_wires_the_adapter_methods_and_discovers_both_factors(): + checks = LogprobChecks + assert checks.build_contract is adapter.build_contract + assert checks.read_effective_config is adapter.read_effective_config + assert checks.observe_collectives is adapter.observe_collectives + assert checks.resolve_implementation is adapter.resolve_implementation + + ids = [factor.id for factor in LogprobChecks().declare_factors()] + assert ids == ["logp.lse_merge_order", "logp.precision_downcast"]