From 215d8cc3c0906e2bf9df9b08d5e7ff0cdf58ebf7 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sun, 9 Aug 2026 11:56:37 +0800 Subject: [PATCH 1/7] feat(mismatch): add the training-inference mismatch framework 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. Ships the framework only. Operators are claimed and written separately, so operator_checks/ is empty by design and adding one changes nothing outside its own directory. Layout, in dependency order: schema/ pure data types, no behaviour pipeline/ the seven execution steps, free functions only engines/ the two sides under test reference_adapters/ wiring reference implementations in model_meta/ model shape and module correspondence operator_checks/ plugins, one directory per operator (empty) Three design decisions worth stating: Four variants, 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. That last arm is the self-check gate -- without it one wrong reference quietly steers every attribution, which is worse than having no framework at all. Four gates before the matrix. "Not measured" and "measured and clean" are different, and confusing them is the mistake this kind of framework is most likely to make. A silently reverted switch, missing evidence, an incomplete set of logprob shards, or a failed pitfall guard each block a verdict rather than passing through as a clean result. Convergence is judged on clip_fraction, not dlogp_mean. At every production floor the mean sits far below the GRPO clip edge, so judging on it would mark almost every factor NOT_THIS_FACTOR while gradient signal is being discarded in the tail. Thresholds are code constants keyed by (model family, noise floor), not configuration: a tunable threshold is one somebody can tune until the test passes, and it has to enter the execution fingerprint so changing it invalidates historical results. 39 framework tests, all on CPU via a synthetic scoring backend that can reproduce the failure modes above. Nothing here claims anything about real Megatron or vLLM numerics. --- rl_engine/mismatch/__init__.py | 29 + rl_engine/mismatch/__main__.py | 156 ++++ rl_engine/mismatch/engines/__init__.py | 11 + rl_engine/mismatch/engines/cpu_reference.py | 127 +++ rl_engine/mismatch/model_meta/__init__.py | 12 + rl_engine/mismatch/model_meta/qwen3.py | 73 ++ .../mismatch/operator_checks/__init__.py | 82 ++ rl_engine/mismatch/pipeline/__init__.py | 56 ++ rl_engine/mismatch/pipeline/comparison.py | 179 ++++ rl_engine/mismatch/pipeline/diagnosis.py | 226 +++++ rl_engine/mismatch/pipeline/planner.py | 217 +++++ rl_engine/mismatch/pipeline/registry.py | 189 ++++ rl_engine/mismatch/pipeline/report.py | 188 ++++ rl_engine/mismatch/pipeline/runner.py | 281 ++++++ .../mismatch/reference_adapters/__init__.py | 30 + .../mismatch/reference_adapters/settings.py | 119 +++ rl_engine/mismatch/schema/__init__.py | 129 +++ rl_engine/mismatch/schema/collectives.py | 130 +++ rl_engine/mismatch/schema/contracts.py | 87 ++ rl_engine/mismatch/schema/factors.py | 215 +++++ rl_engine/mismatch/schema/fingerprints.py | 125 +++ rl_engine/mismatch/schema/metrics.py | 172 ++++ rl_engine/mismatch/schema/pitfalls.py | 55 ++ rl_engine/mismatch/schema/rollout_context.py | 103 +++ rl_engine/mismatch/schema/thresholds.py | 116 +++ rl_engine/mismatch/schema/tracing.py | 121 +++ rl_engine/mismatch/schema/values.py | 191 +++++ rl_engine/mismatch/schema/variants.py | 152 ++++ tests/test_mismatch_framework.py | 804 ++++++++++++++++++ 29 files changed, 4375 insertions(+) create mode 100644 rl_engine/mismatch/__init__.py create mode 100644 rl_engine/mismatch/__main__.py create mode 100644 rl_engine/mismatch/engines/__init__.py create mode 100644 rl_engine/mismatch/engines/cpu_reference.py create mode 100644 rl_engine/mismatch/model_meta/__init__.py create mode 100644 rl_engine/mismatch/model_meta/qwen3.py create mode 100644 rl_engine/mismatch/operator_checks/__init__.py create mode 100644 rl_engine/mismatch/pipeline/__init__.py create mode 100644 rl_engine/mismatch/pipeline/comparison.py create mode 100644 rl_engine/mismatch/pipeline/diagnosis.py create mode 100644 rl_engine/mismatch/pipeline/planner.py create mode 100644 rl_engine/mismatch/pipeline/registry.py create mode 100644 rl_engine/mismatch/pipeline/report.py create mode 100644 rl_engine/mismatch/pipeline/runner.py create mode 100644 rl_engine/mismatch/reference_adapters/__init__.py create mode 100644 rl_engine/mismatch/reference_adapters/settings.py create mode 100644 rl_engine/mismatch/schema/__init__.py create mode 100644 rl_engine/mismatch/schema/collectives.py create mode 100644 rl_engine/mismatch/schema/contracts.py create mode 100644 rl_engine/mismatch/schema/factors.py create mode 100644 rl_engine/mismatch/schema/fingerprints.py create mode 100644 rl_engine/mismatch/schema/metrics.py create mode 100644 rl_engine/mismatch/schema/pitfalls.py create mode 100644 rl_engine/mismatch/schema/rollout_context.py create mode 100644 rl_engine/mismatch/schema/thresholds.py create mode 100644 rl_engine/mismatch/schema/tracing.py create mode 100644 rl_engine/mismatch/schema/values.py create mode 100644 rl_engine/mismatch/schema/variants.py create mode 100644 tests/test_mismatch_framework.py diff --git a/rl_engine/mismatch/__init__.py b/rl_engine/mismatch/__init__.py new file mode 100644 index 00000000..218062dc --- /dev/null +++ b/rl_engine/mismatch/__init__.py @@ -0,0 +1,29 @@ +# 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 framework turns "which of the dozens of possible causes +is it" into a set of switches that can be flipped one at a time and attributed. + +Layout:: + + schema/ pure data types, no behaviour + pipeline/ the seven execution steps, free functions only + engines/ the two sides under test + reference_adapters/ wiring reference implementations in + model_meta/ model shape and module correspondence + operator_checks/ plugins, one directory per operator (empty by design) + +Dependency rules: + +1. ``schema/`` never imports ``pipeline/``; within ``schema/`` the imports go one + way only, with ``values.py`` at the top importing nothing from the project. +2. ``pipeline/`` never imports ``operator_checks/`` -- it only sees 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. +""" + +from rl_engine.mismatch import pipeline, schema diff --git a/rl_engine/mismatch/__main__.py b/rl_engine/mismatch/__main__.py new file mode 100644 index 00000000..71964d38 --- /dev/null +++ b/rl_engine/mismatch/__main__.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Command line entry point. **The only place that imports operator plugins.** + +The framework does not know which operators exist -- 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 + +# Operator plugins are listed here and nowhere else. An operator that is not +# imported simply 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: + owner = factor.owner or "unclaimed" + print(f" {factor.id:<40} {factor.category.value:<24} owner={owner}") + 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 + + # Static rejection before anything is expanded: a factor that claims topology + # independence while reducing non-deterministically produces numbers that + # mean nothing, and that is knowable without running. + 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/engines/__init__.py b/rl_engine/mismatch/engines/__init__.py new file mode 100644 index 00000000..2a4f6e4a --- /dev/null +++ b/rl_engine/mismatch/engines/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The two sides under test: process and engine lifetime, configuration readback. + +Shared across operators -- all of attention's factors use one ``vllm.py``. Not to +be confused with ``reference_adapters/``, which wires in *reference* +implementations and contains no operator code. +""" + +from rl_engine.mismatch.engines.cpu_reference import CpuScoringBackend diff --git a/rl_engine/mismatch/engines/cpu_reference.py b/rl_engine/mismatch/engines/cpu_reference.py new file mode 100644 index 00000000..98fb7e69 --- /dev/null +++ b/rl_engine/mismatch/engines/cpu_reference.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A CPU scoring backend used to exercise the framework without a GPU. + +This validates plumbing only. It does not claim anything about real Megatron or +vLLM numerics -- those live in ``megatron.py`` and ``vllm.py`` and need devices. + +It is deliberately able to *simulate* the failure modes the framework is built to +catch, so the gates and the matrix can be tested for real: + +* a configurable per-side bias, so one-sided attribution can be checked; +* a switch that silently does nothing, so ``FELL_BACK`` is reachable; +* an unstable mode whose output changes with the environment, so the + topology-independence assertion can actually fail. +""" + +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. + + ``bias`` is the deviation this side adds on top of the shared base signal. + Setting it on one side only is how a one-sided root cause is simulated. + """ + + 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 only depends on the identity, so both sides agree unless + a bias is injected -- which is what makes an injected deviation the only + thing the metrics can see. + """ + + self.applied_calls.append(dict(switch_values)) + base = _base_signal(identity) + + # A replacement callable means the reference implementation is in play; + # 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 how expensive they are to change. + + Mirrors the four ``RebindCost`` levels so case ordering can be tested. + """ + + 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/rl_engine/mismatch/model_meta/__init__.py b/rl_engine/mismatch/model_meta/__init__.py new file mode 100644 index 00000000..df83c59c --- /dev/null +++ b/rl_engine/mismatch/model_meta/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Model shape and the module correspondence table. + +Filled once per model by whoever brings it up, 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 GLM5 or DSv3.2 means +redoing it -- which is why it lives here and not in any ``operator_checks/``. +""" + +from rl_engine.mismatch.model_meta.qwen3 import 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..0ad09f64 --- /dev/null +++ b/rl_engine/mismatch/model_meta/qwen3.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Qwen3 module correspondence and call chain. + +Every entry with an ``equivalence`` also carries ``verified_by``: **an +equivalence may not be claimed without a test that proves it**, otherwise +"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..341ede21 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/__init__.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Operator plugins, one directory per operator. **Empty by design.** + +The framework ships without operators. Each operator is claimed and written +separately, and adding one changes nothing outside its own directory. + +Importing this package registers nothing on its own -- ``__main__`` imports the +individual operator packages, and that import is what triggers +self-registration. The framework does not know which operators exist; the CLI +brings them in. + +Layout of one operator +---------------------- + +An operator has more than a dozen factors, each with eight or so fields. Putting +them all in one file runs to 800 lines that nobody wants to edit, so it is **one +file per factor**, mirroring "one directory per operator":: + + operator_checks/attention/ + |-- __init__.py ~15 lines: operator name + discover_factors, rest delegated + |-- adapter.py the four operator-level methods + |-- _common.py shared across factors: contract helpers, common + | comparison rules, this operator's reference implementations + `-- factors/ one file per factor, 30-50 lines each + |-- provenance.py attn.provenance + |-- execution_path.py attn.execution_path + `-- ... + +**File name equals the factor id with the operator prefix stripped.** +``discover_factors()`` enforces it, so renaming an id without renaming the file +fails at import rather than silently. + +What a factor file contains +--------------------------- + +One ``FACTOR`` constant and nothing else -- declarations, no behaviour:: + + FACTOR = MismatchFactor( + id="attn.rope_fusion", + operator="attention", + category=FactorCategory.KERNEL_IMPLEMENTATION, + question="...", + switch=Switch(...), + comparison_rules={...}, + prerequisites=Prerequisites(...), + reference=TE_ROPE_REFERENCE, # from _common.py + pitfalls=(KnownPitfall(...),), + required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value, POSITION_CACHE), + owner="", + tracked_by=(), # issues or PRs following this factor + ) + +What ``__init__.py`` contains +----------------------------- + +:: + + @OPERATOR_CHECKS.register + class AttentionChecks: + operator = "attention" + + def declare_factors(self): + return discover_factors(__package__) # scans factors/*.py + + build_contract = adapter.build_contract + read_effective_config = adapter.read_effective_config + observe_collectives = adapter.observe_collectives + resolve_implementation = adapter.resolve_implementation + +Adding a factor means dropping a file into ``factors/``; ``__init__.py`` does not +change. + +Why the four methods stay in ``adapter.py`` +------------------------------------------- + +They are **operator-level**, not factor-level: how to read configuration back +from an engine and how to observe collectives is one piece of logic shared by all +of that operator's factors. Pushing them down would make every factor file copy +them. +""" diff --git a/rl_engine/mismatch/pipeline/__init__.py b/rl_engine/mismatch/pipeline/__init__.py new file mode 100644 index 00000000..b2fd373d --- /dev/null +++ b/rl_engine/mismatch/pipeline/__init__.py @@ -0,0 +1,56 @@ +# 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: + +1. ``registry`` -- plugin registration, factor discovery, import-time conflicts +2--4. ``planner`` -- filter, statically reject, expand variants, order by cost +5. ``runner`` -- execution loop, reuse decisions, variant repeats +6. ``diagnosis`` -- four gates plus the matrix +7. ``report`` -- filter false positives, trace root causes, emit the report + +``comparison`` is used by the runner and holds the declaration-driven contract +comparison. + +Dependency rules: ``schema/`` never imports ``pipeline/``; ``pipeline/`` never +imports ``operator_checks/``. Breaking the second rule turns "add an operator" +into "change the framework". +""" + +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, +) diff --git a/rl_engine/mismatch/pipeline/comparison.py b/rl_engine/mismatch/pipeline/comparison.py new file mode 100644 index 00000000..20e663e7 --- /dev/null +++ b/rl_engine/mismatch/pipeline/comparison.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Contract comparison, driven entirely by declarations. + +The comparison 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``. + + Returns a sentinel when the path is absent so a missing field is reported + rather than raising. + """ + + 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``. + + Returns every disagreement. ``RECORD_ONLY`` fields are never compared -- that + tier exists precisely so structural differences (packed QKV and the like) do + not drown the real problems in false positives. + """ + + 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 a topology-independent implementation against one that is not + reproducible even across runs measures the weaker side's noise, not the gap + between them. + """ + + 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 legitimately declare different numbers of + # collectives; the overlap is what can be compared here. + 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..b376906f --- /dev/null +++ b/rl_engine/mismatch/pipeline/diagnosis.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Four gates plus the diagnosis matrix. Step 6 of the pipeline. + +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 two different things** -- + confusing them is the mistake an attribution framework is most likely to + make, so the gates come first and nothing reaches the matrix without passing + them. + """ + + 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 reference was requested, 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? + + Under TP/CP each rank holds one slice. Missing a slice makes the combined + number wrong in a way that does not show: one vocab shard short and the LSE + denominator loses a chunk, so logp comes out systematically high. + """ + + 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. Before calling it clean, check the tail: the mean stays + # far below the clip edge at every production floor, so judging on the mean + # alone would mark almost everything NOT_THIS_FACTOR. + 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..674af704 --- /dev/null +++ b/rl_engine/mismatch/pipeline/planner.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Planning: filter, statically reject, expand variants, order by rebuild cost. + +Steps 2 to 4 of the pipeline. +""" + +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: + """Static check, run between steps 2 and 3. Nothing executes. + + Claiming topology independence while using a non-deterministic reduction + order produces numbers that mean nothing. Reject at planning time rather + than explaining it away after a full run. + """ + + 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, model traits. + + A whitelist probe, not a hand-maintained blacklist of "unsupported" strings. + """ + + 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. + + A single swap is not enough: **only a one-sided swap tells you which side is + at fault, and only a two-sided swap proves the reference itself is sound.** + Hence four arms rather than on/off. + """ + + if factor.variants: + return tuple(factor.variants) + + path = factor.switch.path + reference = factor.reference + + if reference is None: # no reference implementation means 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. + + **Ordering is required, not a nicety**: run 160 cases in a random order and + the worst case restarts the process for every one of them. + """ + + 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 even show anything at the given floor. + + A factor whose switch is process-level 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..12ed2a7d --- /dev/null +++ b/rl_engine/mismatch/pipeline/registry.py @@ -0,0 +1,189 @@ +# 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. No +existing file changes. If you find yourself having to edit the planner, a global +dict, or another operator's file, the abstraction is missing something -- raise +it and fix the framework 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. + + The four methods below are **operator-level**, not factor-level: how to read + configuration back from an engine and how to observe collectives is the same + logic for all of an operator's factors. Pushing them down would make every + factor file copy them. Factor files hold **declarations only, no behaviour**. + """ + + operator: str + + def declare_factors(self) -> tuple[MismatchFactor, ...]: + """Which factors this operator has. The only thing that must be written + by hand.""" + ... + + 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 from the engine. A requested + value is not evidence.""" + ... + + def observe_collectives(self, role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """Which collectives actually ran. Feeds ``COLLECTIVE_CONTRACT``.""" + ... + + 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** -- which candidates were + tried and why each was rejected. Returning 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: + """Decorator. Instantiates and registers, 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 factor 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) + + 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}") + 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}" + ) + + 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``. + + The convention: **file name equals the factor id with the operator prefix + stripped**. It is enforced here so that renaming an id without renaming the + file fails at import rather than silently. + """ + + 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..89e4893b --- /dev/null +++ b/rl_engine/mismatch/pipeline/report.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""False-positive filtering, root-cause tracing, and the final report. Step 7.""" + +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. + + Without this, a difference like fused QKV turns every weight comparison red + and buries the real problem. An equivalence may only be claimed with a test + that proves it -- one without ``verified_by`` is not trusted here. + """ + + 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. + + Thirty factors give thirty diagnoses, and that pile is not the answer. The + job here is to combine them with the module correspondence table and the + call chain into "the few most suspicious modules, ranked". + """ + + 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. The hypotheses are what you want to read; + everything else is the evidence supporting them.""" + + 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..841c0c31 --- /dev/null +++ b/rl_engine/mismatch/pipeline/runner.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The execution loop: reuse decisions, variant repeats, metric computation. + +Step 5 of the pipeline. +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass +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. + + Deliberately minimal so a CPU stub can satisfy it -- the framework's + 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. + + If a kernel mutates weights in place, the first and second run differ, so + ``both_reference`` fails bitwise at random and gets blamed on the reference + -- attribution then points somewhere entirely wrong. 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)``. The objective + clips ``rho`` at ``1 +/- eps``, which is ``|dlogp| > ln(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, the standard PPO/GRPO diagnostic: 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". It never compares across + frameworks, so it is cheap; what it verifies is the **premise of the + self-check gate** -- ``both_reference`` can only anchor if the fixed-order + implementation really did fix the order. + """ + + 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, no cross-framework comparison -- 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 + + contracts = { + role: checks.build_contract(role, variant.switch_values) + for role in (PolicyRole.ROLLOUT, PolicyRole.TRAINING) + } + 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, readback in readbacks.items() + for key, value in readback.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..f5be7273 --- /dev/null +++ b/rl_engine/mismatch/reference_adapters/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Wiring implementations in as reference implementations. **No operator code here.** + +"Reference implementation" spans two directories. Of the three +``ReferenceAuthority`` levels, two are implemented in ``rl_engine/kernels/``: + +=========================== ===================== ============================== +authority implementation lives what this package does +=========================== ===================== ============================== +``FP64_ORACLE`` ``kernels/`` declare it, call it at the + lowest noise floor +``SHARED_BACKEND`` external library put it in deterministic mode, + read back to verify +``SELF_WRITTEN`` ``kernels/`` wrap it as a + ``ReferenceImplementation`` +=========================== ===================== ============================== + +Split of duties: ``kernels/`` owns "is the arithmetic right"; this package owns +"how do we put it in a deterministic mode, and how do we prove the setting +actually took effect". The latter is specific to this diagnostic framework and +has nothing to do with the operator itself. +""" + +from rl_engine.mismatch.reference_adapters.settings import ( + 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..95a03dda --- /dev/null +++ b/rl_engine/mismatch/reference_adapters/settings.py @@ -0,0 +1,119 @@ +# 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 the +settings have to be pinned explicitly. But pinning alone is not enough: **a +setting that cannot be read back can only be recorded as ``UNOBSERVABLE``** -- +delivered but unverifiable is the same as not delivered. +""" + +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. + + Which channel a setting uses decides when it can take effect, and therefore + its rebind cost -- an env var needs a process restart, a call argument does + not. Flattening them into one dict leaves the framework unable to deliver + them at all. + """ + + 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. + + Returns ``UNOBSERVABLE`` when a setting declares no readback path: it was + delivered but cannot be proven, which is not evidence. + """ + + 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..1dd4c30d --- /dev/null +++ b/rl_engine/mismatch/schema/__init__.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pure data structures: public fields, frozen, no meaningful methods. + +All behaviour lives in free functions under ``pipeline/``. This is a deliberate +trade-off, not an accident of using dataclasses: + +====================== ================== ================== + add a new function add a new type +====================== ================== ================== +data + procedural easy hard +objects + polymorphism hard easy +====================== ================== ================== + +This framework grows by **adding functions** (new diagnosis rules, new report +views, new ordering strategies, new evidence checks) while the set of types +stays stable -- so procedural is the correct side. + +Two consequences: + +* Do not hang methods on these structures. ``spec.requires_fixed_order()`` turns + a data structure into a half-object hybrid, which is the worst of both worlds. + Write ``requires_fixed_order(spec)`` instead. +* Chained field access such as ``factor.switch.path`` is **fine** here and does + not violate the Law of Demeter -- Demeter constrains an object's internals, + and plain data structures are supposed to expose their fields. Do not wrap + them in getters for the sake of it. + +Modules are ordered by dependency; each only imports 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, +) diff --git a/rl_engine/mismatch/schema/collectives.py b/rl_engine/mismatch/schema/collectives.py new file mode 100644 index 00000000..4142c28e --- /dev/null +++ b/rl_engine/mismatch/schema/collectives.py @@ -0,0 +1,130 @@ +# 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. +``gemm.forward_reduce``, ``gemm.dgrad_reduce``, ``attn.cp_split_k_merge_order``, +``logp.reduce_topology``, ``moe.tp_forces_sp_reduce`` and +``gemm.rollout_all_reduce_backend`` are six instances of *one* semantic model -- +writing each separately guarantees they 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 explicitly rather than left blank + + +class ParallelDim(str, Enum): + """Which parallel dimension the communication happens on. + + Not ``ProcessGroupKind`` -- ``ProcessGroup`` is an existing + ``torch.distributed`` type, and this describes a parallel dimension, not the + process group object. + """ + + 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" # first to arrive is first added. Control group only + NCCL_ALGORITHM = "nccl_algorithm" # ring/tree, 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. + + Called a contract rather than a spec: it shares the word with + ``OperatorContract``, and it really is "what both sides must agree on", + not merely a description. + + A self-contradictory combination must be rejected at planning time: claiming + ``determinism == STABLE_ACROSS_TOPOLOGY`` while setting ``reduction_order`` + to ``NCCL_ALGORITHM`` or ``ARRIVAL`` produces numbers that mean nothing. + Static check, no run needed -- see ``reject_contradictory_factors()``. + """ + + 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 mathematically identical rewrite of a collective -- equal in algebra, + unequal in floating point. + + Megatron replacing ``all_reduce`` with ``reduce_scatter + all_gather`` when + sequence parallelism is on is exactly this rule applied. Whether the two + sides apply it 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..2f252264 --- /dev/null +++ b/rl_engine/mismatch/schema/contracts.py @@ -0,0 +1,87 @@ +# 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 -- + renaming one requires a schema version bump.** + + Deliberately excludes "the declaration contradicts itself": that is rejected + by the planner before anything runs (``reject_contradictory_factors()``) and + has nothing to do with comparing the two sides. + """ + + 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 -- an element of what + ``compare_contracts()`` returns.""" + + 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 three things are common to every operator and live here; the rest (RoPE + state, vocab shard map, GQA head mapping, ...) belongs to each plugin and + goes in ``extra``. + + Field path convention: the keys of ``comparison_rules`` are paths counted + from the contract root -- + + "precision.accumulate" + "collectives[0].reduction_order" + "extra.rope_theta" + + The framework indexes both contracts by path and compares entry by entry. So + a plugin only has to put fields in the right place inside + ``build_contract()``; it never needs to write a function that "collects the + comparable fields". Keep ``extra`` flat -- nesting dicts 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..5eedc653 --- /dev/null +++ b/rl_engine/mismatch/schema/factors.py @@ -0,0 +1,215 @@ +# 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 -- decides where to look when it fires. + + Deliberately not ``Layer``: in this domain a layer is a model layer. + """ + + INPUT_IDENTITY = "input_identity" # tokens / mask / position_ids / eps + ENVIRONMENT = "environment" # framework versions, NCCL, determinism switches + KERNEL_IMPLEMENTATION = "kernel_implementation" # backend choice, 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. + + General items only. Operator-specific evidence does not belong in this enum: + putting it here would turn "add an operator" into "change the framework". + Plugins declare their own as plain strings (see the constants below). + """ + + EFFECTIVE_CONFIG_READBACK = "effective_config_readback" # read back, not requested + MODEL_STATE_FINGERPRINT = "model_state_fingerprint" # all variants share one set of weights + LIBRARY_VERSIONS = "library_versions" + + +# Operator-specific evidence: plugins define constants in their own module and +# the framework treats them as plain strings. Adding an operator adds a new set +# of constants; the framework does not change. +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. + + Not ``ReferenceTier``: "tier" says there are levels without saying what + orders them. What orders these is authority. + + This is a decision order, not a description: look for a SHARED_BACKEND + first, and only write SELF_WRITTEN when the first two cannot cover it. + """ + + FP64_ORACLE = "fp64_oracle" # slow, mathematically exact, lowest noise floor only + SHARED_BACKEND = "shared_backend" # TransformerEngine / FlashInfer + SELF_WRITTEN = "self_written" + + +@dataclass(frozen=True) +class Switch: + """The one definition of a switch. Allowed values and parser declared once.""" + + 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 this factor needs in order to run. A whitelist, not a blacklist. + + Capability probing itself is delegated: this only declares what is needed. + """ + + 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 execution paths it covers. + + **``covers_paths`` defines the shape of the self-check gate**, the field most + easily overlooked: the gate is not "both sides use the same implementation", + it is "**every path this reference covers must agree bitwise on the same + sequence**". + + Most factors cover two paths (training full-prefill plus rollout + full-prefill) and the gate holds across the two sides. But an attention + FlashInfer reference also covers ROLLOUT_DECODE -- so the gate holds *inside + the rollout side*: with full-prefill and decode both switched to FlashInfer + and ``num_splits=1`` pinned, ``(out, lse)`` must agree on the same sequence. + A disagreement is the reference's fault, unrelated to the training side. + **No decode stub is needed on the training side.** + + Covering several paths also gives the attribution split directly: + + dlogp(training, real rollout) + = dlogp(TRAINING_FULL_PREFILL, ROLLOUT_FULL_PREFILL) cross-framework + + dlogp(ROLLOUT_FULL_PREFILL, ROLLOUT_DECODE) rollout-side path + + The two terms are fixed in completely different ways: the first by swapping + the kernel backend, the second by changing decode's softmax block size to + match prefill. Without the middle path, attribution stops at the combined + "cross-framework plus cross-path" answer. + """ + + 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, ...] = () # **required**, see LibraryPin + + +@dataclass(frozen=True) +class MismatchFactor: + """One suspected cause of training-inference mismatch. + + Not ``DivergenceFactor``: in an RL context, divergence means KL divergence. + + A factor with ``reference is None`` is a parameter sweep; with a reference it + is 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 + owner: str = "" # "" = unclaimed + tracked_by: tuple[str, ...] = () # issues or PRs following this factor + + +def declared_collectives(factor: MismatchFactor) -> tuple[CollectiveContract, ...]: + """Collectives a factor's reference implementation pins, if any. + + Used by the planner's static check; returns empty when the factor does not + touch collective communication. + """ + + reference = factor.reference + if reference is None: + return () + pinned = tuple( + setting.value + for setting in reference.required_settings + if isinstance(setting.value, CollectiveContract) + ) + return pinned + + +def requires_fixed_order(contract: CollectiveContract) -> bool: + """Whether this contract claims its result is independent of topology. + + Claiming it means it must survive the assertion in ``pipeline/planner.py``. + """ + + 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..33bbca1b --- /dev/null +++ b/rl_engine/mismatch/schema/fingerprints.py @@ -0,0 +1,125 @@ +# 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, coarse to fine. + + Compared coarse to fine: a different ``process`` means restarting the whole + process; equal process but different engine means only rebuilding the engine. + """ + + 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. When this layer changes, every numerical + conclusion has to be re-run.""" + + 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. + + Three rules: + + 1. **What goes in is the value read back, not the value requested.** + Requesting ``num_splits=1`` and the backend actually using 1 are two + different facts. + 2. **Library versions belong here, not in a footnote.** + 3. **Thresholds go in too.** Change a threshold and every historical + pass/fail must go 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 immutable on-disk record. Append-only, never edited. + + Difference from ``VariantResult``: a Result is what was just computed in + memory, a Record is archived -- it additionally carries the execution + identity and a content hash, the latter being its integrity seal. Only a + record whose hash verifies may be reused on resume. + + Not ``VariantArtifact`` -- in an ML context, "artifact" usually means model + weights or checkpoints (MLflow artifacts); this stores an execution record. + """ + + variant_name: str + fingerprint: ExecutionFingerprint + result: VariantResult + content_hash: str + + +def canonical_fingerprint(payload: Any) -> str: + """Stable hash of a JSON-serialisable payload. + + Key order is normalised so the same content always hashes the same, whatever + order the dict was built in. + """ + + 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..4113246f --- /dev/null +++ b/rl_engine/mismatch/schema/metrics.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Metrics and per-variant results. + +``dlogp`` is an intermediate quantity. GRPO's gradient enters the objective +through the importance sampling ratio:: + + rho_{i,t}(theta) = pi_theta(y | x) / pi_old(y | x) = exp(dlogp_{i,t}) + +The objective clips rho, typically at eps = 0.2, i.e. rho in [0.8, 1.2]. Back in +dlogp terms, anything past ``ln(1 + eps) ~= 0.182`` **gets clipped**. + +That is the actual mechanism by which mismatch breaks training: it raises the +clip fraction, and clipped tokens have their gradient signal cut -- **RL discards +samples it should have learned from**, and not random ones, the ones with the +largest mismatch. + +Against the threshold table: dense sits at ``dlogp_mean ~= 0.002-0.008`` and +large MoE at ``0.01-0.03`` -- all far below 0.182. **So the mean never triggers +clipping, and looking only at the mean necessarily 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 -- the entry point for attribution, since it + often points straight at a layer or an expert.""" + + 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 ``fallback_reason: str`` is not enough: what you actually want is + "was the first choice rejected for a missing library, missing devices, or a + version mismatch", plus "which candidate did it land on". Without this trace, + a silent fallback leaves nothing to go on. + """ + + 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. + + With TP/CP on, logprobs are not computed by any single rank -- each holds one + slice. **Missing any slice makes the combined result wrong in a way that does + not show**: drop one vocab shard and the LSE denominator loses a chunk, so + logp is systematically too high. Slices must be collected per rank and + checked against ``world_size`` before merging. + + Not ``RankObservation`` -- in RL, an observation is what the agent sees of + the environment. + """ + + 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. + + A free function rather than a method on a bundle type: making a data + structure reach into a factor's fields would be feature envy. + """ + + 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..c6549769 --- /dev/null +++ b/rl_engine/mismatch/schema/pitfalls.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Known pitfalls, encoded as data. + +Every factor in the docs carries a "pitfall" note. Prose pitfalls get read once +and never again -- they have to be data, so 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 fails -- decides which tool the framework needs against it. + + Borrowed from reliability engineering's failure mode, which is more accurate + than a bare "kind": the useful part is *how it fools you*, not which bucket + it sits in. + """ + + 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 on purpose: a pitfall is a + pitfall precisely because its appearance points at the wrong cause. + + Originally split into ``Pitfall`` and ``Precheck``, but the two referenced + each other's ids -- a two-way reference means they were one thing all along: + a pitfall, and the assertion that stops it. + """ + + id: str + mode: FailureMode + symptom: str # what it looks like + actual_cause: str # what it actually is + guard: str # the assertion that blocks it, one line + 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..2ff29a08 --- /dev/null +++ b/rl_engine/mismatch/schema/rollout_context.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Online RL context. + +The factor model targets **offline pairwise comparison**; real training is an +**online RL loop**. Offline these are placeholders and a factor author never +touches them. They matter only when this framework is wired into a real training +loop for online auditing -- and there, missing them causes two kinds of problem: + +**One, results stop reproducing.** You get ``dlogp = 0.005`` today and ``0.008`` +tomorrow from the same config, and eventually find the batch order changed and +one sample landed in a different microbatch -- **the group it all-reduces with +changed, so the accumulation order changed**. That is what ``BatchPlacement`` +records: not "environment" but part of the **identity**. Changing where a sample +lands is changing the computation. ``DynamicSamplingDecision`` is the same story: +GRPO drops all-correct or all-wrong groups, and which ones were dropped changes +the whole batch composition. + +**Two, metrics get misread.** GRPO normalises advantage within a group: +``A_i = (r_i - mean(r)) / (std(r) + eps)``. So **one rollout's logp drifting +affects all K rollouts in that group through the advantage** -- the error is not +independent per rollout, it amplifies within the group. Averaging over tokens +erases this completely: ``dlogp_mean`` looks fine while a few groups are skewed +wholesale. Hence metrics must be aggregatable by ``RolloutGroup``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RolloutGroup: + """One GRPO group: K rollouts from the same prompt. + + Two modes, depending on who is running: + + * **offline ablation** (the main use): identity is built from a fixed token + fixture and this is a placeholder (1 prompt x 1 rollout); group semantics + are unused; + * **online, wired into real training**: a real group comes from the training + loop, and only then does aggregating by group mean anything. + """ + + 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 (NCHW and + friends); this is about placement. + """ + + 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 (all-correct or all-wrong groups carry no + learning signal). + + Dropping changes the batch composition and therefore the reduction grouping + -- unrecorded, it cannot be reproduced. + """ + + kept: bool + reason: str | None = None + + +@dataclass(frozen=True) +class ComparisonIdentity: + """This comparison's input identity. If these differ, no numerical + comparison means anything. + + Not ``ScoringIdentity`` -- in RL, "score" means reward scoring. + """ + + 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 # "L=1,H=896,Hq=14,Hkv=2,D=64" -- 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..ca235051 --- /dev/null +++ b/rl_engine/mismatch/schema/thresholds.py @@ -0,0 +1,116 @@ +# 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. + + Reading a low floor with the production table hides real bugs: at + SINGLE_LAYER_ANCHOR the expectation is bitwise, and judging it against + 0.002-0.008 would call a definite operator error "normal". So the lookup is + always keyed by noise floor, and 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 absolute floor below which a difference is not treated as a signal. + + Used by the diagnosis matrix as ``tol_floor`` when deciding convergence. + """ + + 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..641a0ee6 --- /dev/null +++ b/rl_engine/mismatch/schema/tracing.py @@ -0,0 +1,121 @@ +# 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, plus any + known structural difference between them. + + Without this table the two sides' tensors cannot even be matched up, let + alone compared entry by entry. + + A non-empty ``equivalence`` means "different in form, equal in arithmetic" -- + the basis for filtering false positives. Without it, a difference like fused + QKV turns every weight comparison red and buries the real problem. + + Filled once per model by whoever brings the model up, 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. Swap in GLM5 or + DSv3.2 and it has to be redone -- so it lives in ``model_meta/``, not in any + ``operator_checks/``. + """ + + 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: mismatch at ``upstream`` propagates + to ``downstream``. + + Tracing walks these backwards: start from the last still-aligned module and + look downstream for the first mismatched one. + """ + + upstream: str # semantic_name + downstream: str + + +class RootCauseCategory(str, Enum): + """How a root cause is classified.""" + + 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 root-cause hypothesis, produced by walking the call chain after a + ``NOT_THIS_FACTOR``. + + ``anchor_module`` is the last position where the two sides still agree -- + the root cause must be downstream of it. + """ + + 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. **Three levels away from a Diagnosis**:: + + one factor's variants -> VariantResult x N + | diagnose() + FactorReport (holds one Diagnosis enum value) + | x 30 factors + trace_root_causes() + v + MismatchReport (ranked hypotheses) + + * ``Diagnosis`` is **one factor's** conclusion, one of eight values; + * ``FactorReport`` is that factor's full record: every variant plus that one + diagnosis plus the reason; + * ``MismatchReport`` is the **whole run**, answering what a Diagnosis cannot. + + Thirty factors give thirty diagnoses -- say three ``CAUSED_BY_TRAINING_SIDE``, + twenty-five ``NOT_THIS_FACTOR``, two ``INSUFFICIENT_EVIDENCE``. That pile is + not the answer. **The report's job is to combine them with the module + correspondence table and the call chain into "the few most suspicious + modules, ranked"** -- the ``hypotheses`` field. The rest is 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, ...] # those with equivalence set + 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..5382f15b --- /dev/null +++ b/rl_engine/mismatch/schema/values.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Base value types. This module 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 ``Side``: that only says "one of the sides" without saying of what. + Not ``Policy`` either -- in RL that means the policy itself (network and + distribution), not the part it plays in this comparison. + """ + + ROLLOUT = "rollout" # pi_old -- the policy that produced the trajectory + TRAINING = "training" # pi_theta -- the policy being updated + + +class ExecutionPath(str, Enum): + """Which physical path produced a set of logprobs. + + The same mathematical quantity (a token's conditional probability) can come + out of several different physical paths: feed the whole sequence at once, + feed it in chunks, or generate 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. + + ``downcast`` here means numerical precision reduction (fp32 -> bf16), not + the C++ sense of casting down a class hierarchy. + """ + + NEVER = "never" + PER_BLOCK = "per_block" # right after each block (largest error) + PER_PARTIAL = "per_partial" # each shard's partial sum + FINAL_WRITE = "final_write" # only on the final write (smallest error) + + +class RebindCost(str, Enum): + """What it costs to change this switch. + + Named after the cost rather than an isolation level because its only job is + to answer "is this one expensive?", which drives case ordering and therefore + the wall-clock of the whole run. + """ + + 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. + + These three are delivered in completely different ways. Flattening them into + one dict leaves the framework unable to deliver them, and unable to derive + the rebind cost that case ordering depends on. + """ + + 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 -- + the same factor can reach the *opposite* conclusion on a different version. + So the version is part of the experiment's identity, not a footnote: it goes + into the fingerprint, and changing it invalidates historical results. + """ + + 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: what to pin, how to deliver it, how to + read it back, and which pitfall it guards against. + + ``readback`` is the main reason this type exists: **a requested value is not + evidence**. A setting that cannot be read back can only be recorded as + ``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, the one word this domain + must not borrow. + + Reserved for the precision factors that get wired in later. Today + ``rollout.dtype`` and ``training.compute_dtype`` are two independent + switches with inconsistent names and no way to express "both sides should + match"; this type converges them. + """ + + 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 that accepts only the given values. + + Used by ``Switch`` so the allowed set and the parser are declared once. + """ + + 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..54216eb4 --- /dev/null +++ b/rl_engine/mismatch/schema/variants.py @@ -0,0 +1,152 @@ +# 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: both_native / both_reference / two one-sided + 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. + + Not ``MaterializationStatus`` -- it describes a **switch**'s state (did it + get delivered), which "materialization" neither reveals nor keeps separate + from attention's execution path. + """ + + 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): + """The conclusion for **one factor** after its variants have run -- an enum + value, not a report. + + The first three are "cannot judge" and must stay strictly separate from + "judged, nothing here". It answers "is this one factor the culprit"; it + cannot answer "what causes the mismatch" -- that needs cross-factor + synthesis, which is what ``MismatchReport`` is for. + """ + + 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): + """The experiment's noise floor: how small a difference this run can resolve. + Orthogonal to factors. A floor that has not passed blocks the next one. + + From signal processing: a signal below the noise floor cannot be measured. + Same here -- at the production floor, an e-6 difference drowns in parallel + reduction noise. Deliberately not named after scale: the levels are graded + by noise, not by size. + + The four are arranged so that **each step down introduces exactly one new + noise source**. When a floor starts failing, the suspect set is decided: it + is whatever that floor just added. + """ + + SINGLE_LAYER_ANCHOR = "single_layer_anchor" + # 1 layer / single device / determinism fully on / batch=1 / one token. + # Noise sources: none, only the operator itself. + # Called an anchor because it is the reference point for the other three: + # **failing bitwise here is an operator bug, not "training-inference + # mismatch"**, and the other three floors need not run at all. + + FULL_MODEL_SINGLE_GPU = "full_model_single_gpu" + # All layers / still single device / determinism still on. + # New: accumulation over depth. Tests whether error grows linearly or + # exponentially with layer count. + + SHARDED_SINGLE_NODE = "sharded_single_node" + # All layers / TP + SP on one node / determinism still on. + # New: reduction-order differences from sharding. + # **This is the first floor with "real" training-inference mismatch.** + + PRODUCTION = "production" + # Target TP/CP/PP / determinism off / target batch / 2k-8k, decode path. + # New: everything else (cross-node comms, non-deterministic kernels, the + # real generation path). The floor reports are written from, and the only + # one whose numbers may be read against the threshold table. + + +@dataclass(frozen=True) +class FactorVariant: + """One arm of a controlled experiment -- the machine-readable form of + "how do I configure this ablation".""" + + 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 + # Run this same variant once under each environment and require bitwise + # equality -- **the only exception to "one variant, one execution"**. The + # runner expands the cartesian product automatically. + # + # repeat_under = {"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")} + # -> four runs, asserted bitwise identical. + # + # It never compares across frameworks, so it is very cheap; and what it + # verifies is the **premise of the self-check gate**: both_reference can + # only serve as an anchor if the "fixed-order implementation" really did fix + # the order. Without this, REFERENCE_ITSELF_IS_BROKEN is itself unreliable. + + +@dataclass(frozen=True) +class ExpectedRange: + """The normal band for a metric under one (model family, noise floor, config). + + **Must be a code constant, never a config file** -- a tunable threshold means + somebody can tune it until the test passes. It enters the execution + fingerprint: changing a threshold requires a code change and review, and + 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/test_mismatch_framework.py b/tests/test_mismatch_framework.py new file mode 100644 index 00000000..79b493e9 --- /dev/null +++ b/tests/test_mismatch_framework.py @@ -0,0 +1,804 @@ +# 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.engines.cpu_reference import CpuScoringBackend +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, +) + +# ---------------------------------------------------------------- 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_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_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 == () From 4f8357158dfcef1779fcbbc509dfacd2d06afed0 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sun, 9 Aug 2026 12:02:12 +0000 Subject: [PATCH 2/7] feat(mismatch): add operator interfaces, engine placeholders, and docs The framework shipped without operators. This adds the three interfaces that show the three shapes a factor can take, so each can be claimed and implemented independently: attention/rope_fusion implementation swap against a SHARED_BACKEND gemm/forward_reduce collective communication, SELF_WRITTEN reference logprob/precision_downcast parameter sweep, no reference Every adapter method raises NotImplementedError; the declaration layer works, so `list` and `plan` verify a factor is wired before anything is implemented. engines/ gets megatron.py and vllm.py placeholders whose docstrings carry the settings that must be pinned and the readback path for each. Since engines/ is for the two sides under test, the CPU scoring harness moves out to tests/mismatch_cpu_backend.py -- satisfying ScoringBackend does not make something a side under test. Drops MismatchFactor.owner and .tracked_by: ownership belongs in the issue tracker, not in every factor declaration. Adds README.md and two tutorials covering how to add a kernel factor and how to add a communication feature. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01QjJLV6aQZTtm4X6pjex8Ri --- rl_engine/mismatch/README.md | 196 +++++++++ rl_engine/mismatch/__main__.py | 9 +- rl_engine/mismatch/docs/README.md | 32 ++ rl_engine/mismatch/docs/add-a-comm-feature.md | 300 +++++++++++++ .../mismatch/docs/add-a-kernel-factor.md | 413 ++++++++++++++++++ rl_engine/mismatch/engines/__init__.py | 23 +- rl_engine/mismatch/engines/megatron.py | 96 ++++ rl_engine/mismatch/engines/vllm.py | 94 ++++ .../mismatch/operator_checks/__init__.py | 2 - .../operator_checks/attention/__init__.py | 20 + .../operator_checks/attention/_common.py | 41 ++ .../operator_checks/attention/adapter.py | 94 ++++ .../attention/factors/__init__.py | 0 .../attention/factors/rope_fusion.py | 61 +++ .../mismatch/operator_checks/gemm/__init__.py | 20 + .../mismatch/operator_checks/gemm/_common.py | 99 +++++ .../mismatch/operator_checks/gemm/adapter.py | 79 ++++ .../operator_checks/gemm/factors/__init__.py | 0 .../gemm/factors/forward_reduce.py | 100 +++++ .../operator_checks/logprob/__init__.py | 20 + .../operator_checks/logprob/_common.py | 28 ++ .../operator_checks/logprob/adapter.py | 78 ++++ .../logprob/factors/__init__.py | 0 .../logprob/factors/precision_downcast.py | 66 +++ rl_engine/mismatch/schema/factors.py | 2 - .../mismatch_cpu_backend.py | 0 tests/test_mismatch_framework.py | 2 +- 27 files changed, 1861 insertions(+), 14 deletions(-) create mode 100644 rl_engine/mismatch/README.md create mode 100644 rl_engine/mismatch/docs/README.md create mode 100644 rl_engine/mismatch/docs/add-a-comm-feature.md create mode 100644 rl_engine/mismatch/docs/add-a-kernel-factor.md create mode 100644 rl_engine/mismatch/engines/megatron.py create mode 100644 rl_engine/mismatch/engines/vllm.py create mode 100644 rl_engine/mismatch/operator_checks/attention/__init__.py create mode 100644 rl_engine/mismatch/operator_checks/attention/_common.py create mode 100644 rl_engine/mismatch/operator_checks/attention/adapter.py create mode 100644 rl_engine/mismatch/operator_checks/attention/factors/__init__.py create mode 100644 rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py create mode 100644 rl_engine/mismatch/operator_checks/gemm/__init__.py create mode 100644 rl_engine/mismatch/operator_checks/gemm/_common.py create mode 100644 rl_engine/mismatch/operator_checks/gemm/adapter.py create mode 100644 rl_engine/mismatch/operator_checks/gemm/factors/__init__.py create mode 100644 rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py create mode 100644 rl_engine/mismatch/operator_checks/logprob/__init__.py create mode 100644 rl_engine/mismatch/operator_checks/logprob/_common.py create mode 100644 rl_engine/mismatch/operator_checks/logprob/adapter.py create mode 100644 rl_engine/mismatch/operator_checks/logprob/factors/__init__.py create mode 100644 rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py rename rl_engine/mismatch/engines/cpu_reference.py => tests/mismatch_cpu_backend.py (100%) diff --git a/rl_engine/mismatch/README.md b/rl_engine/mismatch/README.md new file mode 100644 index 00000000..3edcf698 --- /dev/null +++ b/rl_engine/mismatch/README.md @@ -0,0 +1,196 @@ + + + +# `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 a set of named **factors**, +each of which is a switch that can be flipped one at a time, run, and attributed +to a side. + +The output is not a number. It is a **report**: which factors were measured, +which could not be measured and why, and the few most suspicious modules ranked +by suspicion. + +## Why the mismatch matters + +`dlogp = log π_θ − log π_old` enters the GRPO/PPO objective through the +importance ratio `ρ = exp(dlogp)`, which the objective clips at `1 ± ε`. With the +usual `ε = 0.2`, any token past `|dlogp| > ln(1.2) ≈ 0.182` has its **gradient +signal discarded** — and the discarded tokens are not random, they are the ones +with the largest mismatch. + +The trap: at the production floor a healthy `dlogp_mean` is `0.002–0.008` (dense) +or `0.01–0.03` (large MoE), all far below the clip edge. **Judging on the mean +alone always concludes "everything is fine."** The danger lives in the tail, which +is why `MismatchMetrics` carries `dlogp_p99` / `dlogp_max` / `clip_fraction` / +`approx_kl` / `worst_token`, and why the diagnosis matrix converges on +`clip_fraction`, not on the mean (`pipeline/diagnosis.py::_converged`). + +## What it does, concretely + +1. **Collects factors** from operator plugins at import time, rejecting duplicate + ids, duplicate switch paths, and one contract field claimed at two different + comparison rules. +2. **Rejects self-contradictory declarations statically** — claiming topology + independence while reducing in NCCL's chosen order produces numbers that mean + nothing, and that is knowable before anything runs. +3. **Expands each factor into four arms**, not on/off: + + | arm | rollout | training | what it buys | + |---|---|---|---| + | `both_native` | native | native | the baseline every other arm is measured against | + | `both_reference` | reference | reference | **self-check gate** — must be bitwise identical, or every conclusion from this factor is void | + | `training_reference_only` | native | reference | deviation gone ⇒ the training side is the source | + | `rollout_reference_only` | reference | native | deviation gone ⇒ the rollout side is the source | + + A factor with no reference implementation is a **parameter sweep** instead: + one arm per allowed value. +4. **Orders cases by rebuild cost** so a run reuses engines instead of restarting + the process ~160 times. +5. **Runs four gates before any verdict** — "not measured" and "measured and + clean" are different, and confusing them is the mistake an attribution + framework is most likely to make. +6. **Diagnoses** each factor, then **traces root causes** across all of them into + a ranked `MismatchReport`. + +## The five types that carry the whole design + +| type | what it is | +|---|---| +| `MismatchFactor` | one suspected cause — a switch, its comparison rules, its prerequisites, its pitfalls | +| `FactorVariant` | one ablation arm of that factor, as pasteable switch values | +| `VariantResult` | what one arm produced: status, metrics, evidence, effective config | +| `Diagnosis` | the verdict for one factor: training side, rollout side, both, not this factor, or *cannot tell* | +| `OperatorChecks` | the plugin protocol — one operator's factors plus how to read that operator back from an engine | + +One factor expands into arms, each arm produces a result, the results become one +diagnosis, and all diagnoses become one report. Everything else is detail. + +## The pipeline, in order + +``` +1 registry collect plugins and factors pipeline/registry.py +2 planner filter by prerequisites pipeline/planner.py +2.5 planner reject contradictory declarations +3 planner expand factors into variants +4 planner order cases by rebind cost +5 runner run each arm on both sides pipeline/runner.py + → contracts compared field by field pipeline/comparison.py + → dlogp / ρ / clip_fraction +6 diagnosis four gates, then the matrix pipeline/diagnosis.py +7 report filter false positives, rank causes pipeline/report.py +``` + +The four gates in step 6, in order — none of them may be skipped: + +| gate | fails when | verdict | +|---|---|---| +| 1 · did it apply | any arm is not `APPLIED` | `VARIANT_DID_NOT_APPLY` (with the resolution trace: what was tried, why rejected) | +| 2 · evidence | `required_evidence` is incomplete | `INSUFFICIENT_EVIDENCE` | +| 3 · shards | fewer logprob shards than `world_size` | `INSUFFICIENT_EVIDENCE` — a missing shard is wrong in a way that does not show | +| 4 · guards | a pitfall guard failed | `INSUFFICIENT_EVIDENCE` | + +`SwitchStatus.FELL_BACK` is the dangerous state: the reference was requested, the +engine silently reverted to native, and "the deviation did not change" then reads +as a clean `NOT_THIS_FACTOR`. Gate 1 exists for exactly this. + +## Layout + +``` +mismatch/ +├── schema/ pure data types, frozen dataclasses and enums, no behaviour +│ ├── values.py PolicyRole, ExecutionPath, Precision, RebindCost, RequiredSetting +│ ├── collectives.py CollectiveContract, ReductionOrder, DeterminismLevel, rewrites +│ ├── contracts.py OperatorContract, ComparisonRule, ComparisonIssue +│ ├── factors.py MismatchFactor, Switch, Prerequisites, ReferenceImplementation +│ ├── variants.py FactorVariant, SwitchStatus, Diagnosis, NoiseFloor +│ ├── thresholds.py EXPECTED_RANGES — thresholds as code, not configuration +│ ├── rollout_context.py ComparisonIdentity, RolloutGroup, BatchPlacement +│ ├── metrics.py MismatchMetrics, VariantResult, FactorReport, LogprobShard +│ ├── fingerprints.py ReuseKey, ExecutionFingerprint, VariantRecord +│ ├── pitfalls.py KnownPitfall, FailureMode +│ └── tracing.py ModuleCorrespondence, PropagationEdge, RootCauseHypothesis +├── pipeline/ the seven steps, free functions only, no state +├── engines/ the two sides under test — megatron.py and vllm.py, nothing else +├── reference_adapters/ delivering pinned settings by channel and reading them back +├── model_meta/ per-model module correspondence and call chain (qwen3.py) +├── operator_checks/ plugins, one directory per operator — **empty by design** +└── __main__.py CLI, and the only module that imports operator plugins +``` + +### What belongs in `engines/`, and what does not + +`engines/` holds **exactly two modules: `megatron.py` (training side) and +`vllm.py` (rollout side)** — the two policies as they really run. Each one owns +how its engine is constructed, how a switch is delivered to it, and how the +*effective* value is read back off it. They are shared across operators: all of +attention's factors use the one `vllm.py`, and **adding an operator never adds a +file here**. Both are placeholders today; their docstrings list the settings that +must be pinned and the readback path for each. + +Anything that merely satisfies the `ScoringBackend` protocol is **not** an +engine. Test harnesses live under `tests/`: + +| module | what it is | +|---|---| +| `tests/mismatch_cpu_backend.py` | CPU stub for exercising the gates and the matrix with no GPU, with an injectable one-sided bias | + +The line is role, not protocol: an engine is *a side under test*; a harness is +what lets you run the framework when that side is not available. + +Three dependency rules hold the plugin seam open: + +1. `schema/` never imports `pipeline/`; inside `schema/` the imports go one way, + with `values.py` at the top importing nothing from the project. +2. `pipeline/` never imports `operator_checks/` — it only sees what the registry + hands it. Break this and "add an operator" becomes "change the framework". +3. Only `__main__` imports `operator_checks/`, which is what triggers plugin + self-registration. + +## Running it + +```bash +python -m rl_engine.mismatch list # registered operators and their factors +python -m rl_engine.mismatch plan --gpu-count 2 # expand into cases, cheapest rebuild first +python -m rl_engine.mismatch plan --json # same, machine readable +``` + +`list` currently prints "no operator plugins registered", and that is the +intended state: **the framework ships without operators.** Each operator is +claimed and written separately, and adding one changes nothing outside its own +directory plus one line in `__main__._OPERATOR_PACKAGES`. + +The plumbing is testable without a GPU: `tests/mismatch_cpu_backend.py` is a +scoring backend that can simulate a one-sided bias, a switch that silently does +nothing, and an implementation that is unstable across environments — so the +gates and the matrix are exercised for real in `tests/test_mismatch_framework.py`. + +## Noise floors: run them in order + +A factor's result only means something at a floor that can resolve it. The four +floors are arranged so **each step down adds exactly one new noise source**, which +is what makes a failure at one floor point at a known suspect set. + +| floor | configuration | new noise source | +|---|---|---| +| `SINGLE_LAYER_ANCHOR` | 1 layer, single device, determinism on, one token | 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 differences — the first floor with *real* mismatch | +| `PRODUCTION` | target TP/CP/PP, determinism off, decode path | everything else; the only floor whose numbers may be read against `EXPECTED_RANGES` | + +A floor that has not passed blocks the next one. + +## Adding to this package + +| you want to | read | +|---|---| +| add a kernel's mismatch factor (new operator, or a factor on an existing one) | [`docs/add-a-kernel-factor.md`](docs/add-a-kernel-factor.md) | +| add a communication feature (a collective, a reduction order, a rewrite) | [`docs/add-a-comm-feature.md`](docs/add-a-comm-feature.md) | + +Both are the same shape of work: **adding an operator is adding a directory, +adding a factor is adding a file**, and no existing file changes. If you find +yourself editing the planner, a global dict, or another operator's file, the +abstraction is missing something — say so and fix the framework rather than +patching in place. diff --git a/rl_engine/mismatch/__main__.py b/rl_engine/mismatch/__main__.py index 71964d38..a6bf2c39 100644 --- a/rl_engine/mismatch/__main__.py +++ b/rl_engine/mismatch/__main__.py @@ -28,9 +28,9 @@ # Operator plugins are listed here and nowhere else. An operator that is not # imported simply 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", + "rl_engine.mismatch.operator_checks.gemm", + "rl_engine.mismatch.operator_checks.attention", + "rl_engine.mismatch.operator_checks.logprob", ) @@ -83,8 +83,7 @@ def command_list(operator: str | None) -> int: factors = OPERATOR_CHECKS.factors_for(name) print(f"{name}: {len(factors)} factors") for factor in factors: - owner = factor.owner or "unclaimed" - print(f" {factor.id:<40} {factor.category.value:<24} owner={owner}") + print(f" {factor.id:<40} {factor.category.value}") return 0 diff --git a/rl_engine/mismatch/docs/README.md b/rl_engine/mismatch/docs/README.md new file mode 100644 index 00000000..99ee4c57 --- /dev/null +++ b/rl_engine/mismatch/docs/README.md @@ -0,0 +1,32 @@ + + + +# `rl_engine/mismatch` tutorials + +Start with the package overview in [`../README.md`](../README.md) — what a factor +is, the four arms, the four gates, and the seven pipeline steps. These tutorials +assume it. + +| tutorial | when you need it | +|---|---| +| [add-a-kernel-factor.md](add-a-kernel-factor.md) | a kernel (attention, GEMM, RoPE, logprob, SwiGLU, …) is suspected of computing something different on the two sides, and you want it measured and attributed | +| [add-a-comm-feature.md](add-a-comm-feature.md) | the suspect is a **collective**: a reduction order, an `all_reduce` → `reduce_scatter + all_gather` rewrite, a CP/split-K merge, a communication backend | + +Both end with the same two commands, which are how you check your work without a +GPU: + +```bash +python -m rl_engine.mismatch list +python -m rl_engine.mismatch plan --gpu-count 2 +``` + +## The rule both tutorials follow + +**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` +when the operator is new. + +If your change needs an edit inside `pipeline/`, a global dict, or another +operator's directory, stop: the framework is missing an abstraction. Raise it and +fix the framework, rather than patching around it — the seam is the only thing +keeping forty-odd factors from turning into forty-odd special cases. 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..837ef901 --- /dev/null +++ b/rl_engine/mismatch/docs/add-a-comm-feature.md @@ -0,0 +1,300 @@ + + + +# Tutorial: add a communication feature + +Mismatch exists because floating-point addition is not associative, and **the +accumulation order is almost entirely decided by collective communication**. So +collectives are not an implementation detail here — they are a first-class +declared object, and `gemm.forward_reduce`, `gemm.dgrad_reduce`, +`attn.cp_split_k_merge_order`, `logp.reduce_topology`, `moe.tp_forces_sp_reduce` +and `gemm.rollout_all_reduce_backend` are six instances of *one* semantic model. +Written separately they drift apart; written against `CollectiveContract` they +cannot. + +This tutorial adds one: `gemm.forward_reduce` — RowParallel forward reduction, +where the training side does `all_reduce` without sequence parallelism and +`reduce_scatter + all_gather` with it, and the two paths accumulate in different +orders. + +Read [add-a-kernel-factor.md](add-a-kernel-factor.md) first. The file layout, the +four arms, and the gates are identical; this tutorial only covers what is +different when the suspect is a collective. + +## Step 1 · Describe the collective, do not just name it + +`CollectiveContract` is the full numerical semantics of one collective. Every +field is load-bearing: + +```python +from rl_engine.mismatch.schema import ( + CollectiveContract, + CollectiveOp, + DeterminismLevel, + DowncastPoint, + ParallelDim, + Precision, + ReductionOrder, +) + +ORDERED_REDUCE_SCATTER = CollectiveContract( + op=CollectiveOp.REDUCE_SCATTER, + group=ParallelDim.TENSOR, + group_size=2, + 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 explicitly rather than leaving it blank | +| `group` | which parallel dimension: `TENSOR` / `SEQUENCE` / `CONTEXT` / `EXPERT` / `PIPELINE` / `DATA` | +| `reduction_order` | **the direct root of mismatch.** `ARRIVAL` and `NCCL_ALGORITHM` are order-unstable; `GLOBAL_RANK_INDEX`, `GLOBAL_BLOCK_INDEX` and `GLOBAL_VOCAB_SHARD_INDEX` are the fixed orders | +| `accumulate_precision` + `downcast_at` | how much error the accumulation keeps. `PER_BLOCK` is the largest, `FINAL_WRITE` the smallest | +| `determinism` | how strong a reproducibility guarantee this implementation offers | +| `backend` | `"nccl"` / `"vllm_custom_ipc"` / `"mnnvl"` / `"transformer_engine"` / `"rl_kernel"` | + +**One combination is rejected before anything runs**: claiming +`determinism=STABLE_ACROSS_TOPOLOGY` while reducing with `NCCL_ALGORITHM` or +`ARRIVAL` produces numbers that mean nothing. `reject_contradictory_factors()` +raises `ContradictoryFactor` at planning time — a static check, no GPU, no wasted +run. Do not weaken the claim to get past it; fix whichever half is wrong. + +## Step 2 · Pin the contract so the planner can see it + +The planner finds your contract through `declared_collectives(factor)`, which +collects `CollectiveContract` values out of the reference's `required_settings`. +So the contract is pinned like any other setting — with the channel that actually +delivers it: + +```python +# operator_checks/gemm/_common.py +from rl_engine.mismatch.schema import ( + ExecutionPath, + LibraryPin, + ReferenceAuthority, + ReferenceImplementation, + RequiredSetting, + SettingChannel, +) + +DETERMINISTIC_REDUCE_REFERENCE = ReferenceImplementation( + name="rl_kernel", + tier=ReferenceAuthority.SELF_WRITTEN, # justify this in the PR body + 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, # the contract itself, pinned as a value + 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", container_digest="sha256:..."),), +) +``` + +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, so a deterministic `all_reduce` / `reduce_scatter + all_gather` has to +be written. Say that in the PR rather than leaving the tier unexplained. + +The channel is not cosmetic — it decides when a setting can take effect and +therefore the rebind cost: `ENV_VAR` and `TORCH_GLOBAL` need a process restart, +`ENGINE_ARG` an engine rebuild, `CALL_ARG` nothing. + +## Step 3 · The factor, with the collective in the comparison rules + +```python +# operator_checks/gemm/factors/forward_reduce.py +from rl_engine.mismatch.operator_checks.gemm._common import DETERMINISTIC_REDUCE_REFERENCE +from rl_engine.mismatch.schema import ( + COLLECTIVE_CONTRACT, + ComparisonRule, + Evidence, + FactorCategory, + FailureMode, + KnownPitfall, + MismatchFactor, + NoiseFloor, + PolicyRole, + Prerequisites, + RebindCost, + Switch, +) + +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, + "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=DETERMINISTIC_REDUCE_REFERENCE, + call_sites=("attention.o_linear", "mlp.down_linear", "moe.output"), + 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, + ), + ), +) +``` + +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*, not the library. Declaring `backend` as `MUST_MATCH_*` + buries the real finding under a difference you already knew about. +- **`min_gpu_count=2` plus `PROCESS_GROUP_REBUILD`.** The factor is identically + zero on one device, and `suggested_floor_is_lowest()` will tell you it belongs + at `SHARDED_SINGLE_NODE` or above. Running it at the anchor floor wastes time. +- **`call_sites`** records that one factor acts in several physical places — + attention's O-linear, the MLP's down-linear and the MoE output are all row + parallel linears eating the same accumulation-order problem. One factor, three + sites, not three factors. + +`compare_contracts()` adds one check you do not declare: if the two sides' +collectives promise different `DeterminismLevel`s, it emits +`DETERMINISM_INCOMPATIBLE`. Comparing a topology-independent implementation +against one that is not even reproducible across runs measures the weaker side's +noise, not the gap between the sides. + +## Step 4 · The cheapest strong check: rerun under different NCCL settings + +An implementation claiming `STABLE_ACROSS_TOPOLOGY` must produce **bitwise +identical** results when NCCL is told to use a different algorithm. This needs no +cross-framework comparison at all — one side, rerun — which makes it the highest +value-per-minute check in the whole framework. + +Declare it on the arm with `repeat_under`, the single exception to "one variant, +one execution": + +```python +from rl_engine.mismatch.schema import ExpectedOutcome, FactorVariant, PolicyRole + +TOPOLOGY_INVARIANCE = FactorVariant( + name="both_reference", + switch_values={"gemm.forward_reduce": "rl_kernel"}, + replace_on={ + PolicyRole.ROLLOUT: "rl_engine.kernels.collectives.ordered_reduce_scatter", + PolicyRole.TRAINING: "rl_engine.kernels.collectives.ordered_reduce_scatter", + }, + expected=ExpectedOutcome.BITWISE_IDENTICAL, + repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")}, + why="a fixed order must survive NCCL choosing a different algorithm", +) +``` + +The runner expands the cartesian product (four runs here) and asserts they agree +bitwise; a disagreement sets `SwitchStatus.ERROR`, which gate 1 turns into +`VARIANT_DID_NOT_APPLY` rather than a numeric verdict. + +What this protects is **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 arm explicitly through `MismatchFactor.variants` +(a non-empty `variants` tuple is returned as-is by `build_variants`), so declare +all four arms there when you need it on one of them. + +## Step 5 · Observe what actually ran + +```python +# operator_checks/gemm/adapter.py +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """The collectives this operator really performed, this run.""" + return tuple(adapter.collective_trace()) # not what the config asked for +``` + +This feeds the `COLLECTIVE_CONTRACT` evidence item that the factor declares as +required, and gate 2 refuses a verdict without it. The distinction that matters: +`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 the case where the two disagree, and only the second one is +evidence. + +## Declaring a rewrite + +"Mathematically identical, unequal in floating point" is common enough to be its +own type. Two are already declared in `schema/collectives.py`: + +```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 rewrite when sequence parallelism is on, +while the rollout side does not, *is* `gemm.forward_reduce`. If your factor is +"one side rewrites this collective and the other does not", add a +`CollectiveRewrite` constant next to those two and reference it from the factor's +`question`, rather than describing the rewrite in prose. + +## When you genuinely need a new enum value + +Adding a `CollectiveOp`, `ParallelDim` or `ReductionOrder` value **is** a +framework change, and the plugin seam exists so that this is rare. It is +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. It is *not* justified for a new backend (that is the +`backend` string) or a new library version (that is `LibraryPin`). + +If you do add one: + +1. Put it in `schema/collectives.py` with a comment saying what orders it. +2. Extend `_NON_DETERMINISTIC_ORDERS` in `pipeline/planner.py` if the new order is + not stable, so the contradiction check keeps working. +3. Add a case to `tests/test_mismatch_framework.py`. +4. Say in the PR why an existing value could not express it. + +## Checklist before the PR + +- [ ] Every field of `CollectiveContract` is filled from what the code does, not from what the config requests. +- [ ] `determinism` and `reduction_order` do not contradict each other (the planner will tell you, but know why). +- [ ] `backend` is `RECORD_ONLY`; `group_size` is `MUST_MATCH_BITWISE`. +- [ ] `min_gpu_count` ≥ 2 and the suggested floor is `SHARDED_SINGLE_NODE` or above. +- [ ] `NCCL_ALGO` / `NCCL_PROTO` are pinned as `RequiredSetting`s **with readback**. +- [ ] Any topology-independence claim is backed by a `repeat_under` arm. +- [ ] `observe_collectives()` returns the trace, and `COLLECTIVE_CONTRACT` is in `required_evidence`. +- [ ] `call_sites` lists every physical place this one factor acts. +- [ ] A `SELF_WRITTEN` reference is justified against `SHARED_BACKEND` in the PR body. 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..a35ecd27 --- /dev/null +++ b/rl_engine/mismatch/docs/add-a-kernel-factor.md @@ -0,0 +1,413 @@ + + + +# Tutorial: add a kernel's mismatch factor + +You suspect a kernel — attention, GEMM, RoPE, the logprob path, SwiGLU — of +computing something different on the training side than on the rollout side, and +you want that suspicion **measured and attributed to a side** instead of argued +about. + +This tutorial walks the whole path with one worked example: `attn.rope_fusion`, +"is the difference caused by fused RoPE versus small-operator RoPE, or by +`position_ids` / `theta` / the cast boundary?" + +By the end you will have written **declarations only** — no execution logic. The +framework expands the arms, runs them, gates them, and draws the conclusion. + +## Before you write anything: three decisions + +**1 · Is this a parameter sweep or an implementation swap?** + +| | you declare | the framework expands into | +|---|---|---| +| **sweep** — nothing is replaced, you scan a setting | `Switch.allowed_values`, `reference=None` | one arm per allowed value | +| **swap** — a reference implementation replaces the native one | `Switch` + `ReferenceImplementation` | the four arms (plus an `fp64_oracle` arm if declared) | + +There is deliberately no `kind` field: `reference is None` *is* the distinction, +because derivable state is state that can disagree with itself. RoPE is a swap; +`logp.precision_downcast` (in `operator_checks/logprob/`) is the worked sweep. + +**Know what a sweep gives up.** `diagnose()` runs the four-arm matrix, which +needs a `both_native` baseline and the two one-sided arms. A sweep has neither, +so it returns `INSUFFICIENT_EVIDENCE` — *by design*: with nothing swapped there is +no side to attribute to. A sweep measures and records; it does not conclude. Read +its numbers against `EXPECTED_RANGES` yourself, or declare explicit `variants` +once a reference implementation exists. + +**2 · Which reference implementation, and are you allowed to write one?** + +`ReferenceAuthority` is a decision order, not a description: + +``` +FP64_ORACLE slow, mathematically exact — gold standard at the lowest floor + ↑ validates +SHARED_BACKEND TransformerEngine (training) / FlashInfer (rollout) — 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 (`pytorch/attention/rope.py`) and by +`flashinfer.rope`, so it is `SHARED_BACKEND`. + +**3 · Which noise floor can even show it?** + +Pick the *lowest* floor at which the factor is not identically zero. RoPE shows up +with one layer on one device, so `SINGLE_LAYER_ANCHOR` — cheap, and a bitwise +failure there is an operator bug rather than mismatch. A factor whose switch is +`PROCESS_GROUP_REBUILD` is identical on a single device; running it at the anchor +floor only burns machine time (`planner.suggested_floor_is_lowest`). + +## Step 1 · Create the operator directory + +Skip to step 3 if the operator already exists — then you are only adding one file. + +``` +rl_engine/mismatch/operator_checks/attention/ +├── __init__.py ~15 lines: operator name + discover_factors, everything else delegated +├── adapter.py the four operator-level methods +├── _common.py shared across factors: reference implementations, contract helpers +└── factors/ + ├── __init__.py empty + └── rope_fusion.py one FACTOR constant, 30–50 lines +``` + +Why one file per factor: an operator has a dozen-plus factors of eight-odd fields +each. In one file that is 800 lines nobody wants to edit. + +**The file name is the factor id with the operator prefix stripped.** +`discover_factors()` enforces it, so renaming an id without renaming the file +fails at import instead of silently dropping the factor. + +## Step 2 · `_common.py` — what the whole operator shares + +```python +# operator_checks/attention/_common.py +from rl_engine.mismatch.schema import ( + ExecutionPath, + LibraryPin, + ReferenceAuthority, + ReferenceImplementation, + RequiredSetting, + SettingChannel, +) + +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.3.0", container_digest="sha256:..."), + ), +) +``` + +Three fields decide more than they look like they do: + +- **`covers_paths` defines the shape of the self-check gate.** The gate is not + "both sides use the same implementation", it is "*every path this reference + covers must agree bitwise on the same sequence*". Two paths ⇒ the gate holds + across the two sides. A reference that also covers `ROLLOUT_DECODE` (FlashInfer + attention does) makes the gate hold *inside the rollout side*, and then you do + not need a decode stub on the training side at all. +- **`required_settings` are not documentation.** Each is delivered by its channel + (`reference_adapters/settings.py::apply_required_settings`) and verified by + `verify_required_settings`. A setting with `readback=None` can only ever be + recorded `UNOBSERVABLE` — delivered but unprovable is the same as not + delivered. +- **`pinned_libraries` is required.** TE and FlashInfer change kernel selection + across versions; the same factor can reach the *opposite* conclusion on a + different version. The pin goes into the execution fingerprint, so bumping it + invalidates historical results instead of quietly making them incomparable. + +## Step 3 · The factor file — declarations, no behaviour + +```python +# operator_checks/attention/factors/rope_fusion.py +from rl_engine.mismatch.operator_checks.attention._common import 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, + "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "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), + 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, + ), + ), +) +``` + +### Filling the fields well + +**`comparison_rules` — the tier matters more than the field list.** +Keys are dotted paths from the contract root: `precision.accumulate`, +`collectives[0].reduction_order`, `extra.rope_theta`. The framework indexes both +contracts by path and compares entry by entry, so you never write a "collect the +comparable fields" function — you only put fields in the right place in +`build_contract()`. + +| tier | use it when | consequence | +|---|---|---| +| `MUST_MATCH_BITWISE` | identity: shapes, dtypes, TP size, vocab size | differing ⇒ the case is void, not a finding | +| `MUST_MATCH_SEMANTICALLY` | the two sides may implement it differently but must mean the same thing | differing ⇒ a `SEMANTIC_MISMATCH` issue | +| `RECORD_ONLY` | representation that differs by construction — packed QKV layout, page tables, backend names | recorded, **never compared** | + +`RECORD_ONLY` exists precisely so structural differences do not drown the real +problems. Declare a packed-QKV-style field `MUST_MATCH_*` and your factor +disappears under false positives. + +One more constraint the registry enforces: **two factors may not declare the same +contract field at two different rules.** If that fires, one of the two +declarations is wrong — do not "fix" it by renaming the path. + +**`required_evidence` — what must exist before a verdict is allowed.** +Three items apply to every factor (`Evidence.EFFECTIVE_CONFIG_READBACK`, +`MODEL_STATE_FINGERPRINT`, `LIBRARY_VERSIONS`). Operator-specific evidence is a +plain string constant (`POSITION_CACHE`, `LSE_EXPORT`, `VOCAB_SHARD_MAP`, +`COLLECTIVE_CONTRACT`, …) so that adding an operator never means editing an enum +in the framework. Missing evidence is gate 2: `INSUFFICIENT_EVIDENCE`, which is +**not** the same verdict as "measured and clean". + +**`prerequisites` — a whitelist, declared, not probed by hand.** +`required_ops`, `min_gpu_count`, `required_packages`, `required_model_traits` +(`"moe"`, `"linear_attention"`), `blocked_by` (issues this waits on). The planner +turns each unmet item into a reason string, so `plan` prints *why* a factor was +skipped instead of silently omitting it. + +**`pitfalls` — prose gets read once; data gets enforced.** +`symptom` and `actual_cause` are separate fields on purpose: a pitfall is a +pitfall because its appearance points at the wrong cause. `guard_runs_at` should +be the lowest floor that can run the check — cheap guards first. + +## Step 4 · `adapter.py` — the four operator-level methods + +These are operator-level, not factor-level: how to read config back from an +engine is one piece of logic shared by all of the operator's factors. + +```python +# operator_checks/attention/adapter.py +from typing import Any, Callable, Mapping + +from rl_engine.mismatch.schema import ( + DowncastPoint, + ImplementationResolution, + OperatorContract, + PolicyRole, + Precision, + PrecisionProfile, + RejectedCandidate, +) + + +def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: + """This side's switch values -> this side's numerical contract.""" + return OperatorContract( + operator="attention", + role=role, + precision=PrecisionProfile( + compute=Precision.BF16, + accumulate=Precision.FP32, + downcast_at=DowncastPoint.FINAL_WRITE, + softmax_accumulate=Precision.FP32, + ), + collectives=(), # see the comm tutorial + extra={ # keep flat: paths stay short and readable + "rope_theta": 1_000_000.0, + "position_ids_digest": "...", + "post_rope_qk_digest": "...", + "fusion_boundary": switch_values.get("attn.rope_fusion", "native"), + }, + ) + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read switches back **from the engine**. A requested value is not evidence.""" + ... + + +def observe_collectives(role: PolicyRole, adapter: Any): + """Which collectives actually ran. RoPE has none.""" + return () + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve the name an arm asks for into a callable. + + Return the trace **even when resolution fails**: which candidates were tried + and why each was rejected. + """ + rejected: list[RejectedCandidate] = [] + for candidate in _candidates_for(impl_name): + try: + return _import(candidate), ImplementationResolution(impl_name, candidate) + except ImportError as exc: + rejected.append(RejectedCandidate(candidate, str(exc))) + return None, ImplementationResolution(impl_name, None, tuple(rejected)) +``` + +`resolve_implementation` is where a factor most often dies quietly. Returning a +bare `None` produces `SwitchStatus.FELL_BACK` with nothing to investigate, and a +fallen-back arm whose deviation "did not change" reads exactly like a clean +`NOT_THIS_FACTOR`. Gate 1 catches the status; only your trace explains it. + +## Step 5 · `__init__.py` — register, and nothing else + +```python +# operator_checks/attention/__init__.py +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__) # scans factors/*.py + + 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) +``` + +Adding the next factor drops a file into `factors/`; this file does not change. + +## Step 6 · Make the operator exist + +The framework does not know which operators exist. `__main__` is the only module +that imports plugins, which is what keeps `pipeline/` from ever reaching into +`operator_checks/`: + +```python +# rl_engine/mismatch/__main__.py +_OPERATOR_PACKAGES: tuple[str, ...] = ( + "rl_engine.mismatch.operator_checks.attention", +) +``` + +An operator that is not listed here simply does not exist as far as the framework +is concerned. **This one line is the only edit outside your own directory.** + +## Step 7 · Check it, without a GPU + +```bash +python -m rl_engine.mismatch list +# attention: 1 factors +# attn.rope_fusion kernel_implementation +``` + +`plan` on a workstation without TransformerEngine reports the factor as skipped, +naming every unmet prerequisite rather than silently omitting it: + +```bash +python -m rl_engine.mismatch plan --gpu-count 1 +# noise floor: single_layer_anchor +# runnable factors: 0 cases: 0 +# +# skipped (prerequisites not met): +# attn.rope_fusion: operator 'rope' is not dispatchable +# attn.rope_fusion: package 'transformer_engine>=2.0' is not installed +``` + +That is a statement about this machine, not a defect in the factor. Where the +prerequisites are met, the same declaration expands into the four arms, ordered +cheapest-rebuild-first: + +``` +# runnable factors: 1 cases: 4 +# cases in execution order (cheapest rebuild first): +# [engine_rebuild ] attn.rope_fusion :: both_native +# [engine_rebuild ] attn.rope_fusion :: both_reference +# [engine_rebuild ] attn.rope_fusion :: training_reference_only +# [engine_rebuild ] attn.rope_fusion :: rollout_reference_only +``` + +`plan --json` gives the same thing machine readably. + +Then add a test beside `tests/test_mismatch_framework.py`. The CPU backend in +`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 the framework does so you do not have to + +| you might reach for | it already exists | +|---|---| +| writing the four arms out by hand | `build_variants()` | +| a "compare these fields" helper | `compare_contracts()`, driven by `comparison_rules` | +| deciding whether the numbers converged | `diagnose()` — four gates, then the matrix, judged on `clip_fraction` | +| checking prerequisites and printing why one was skipped | `missing_prerequisites()` | +| ordering runs so engines get reused | `order_cases_by_rebind_cost()` | +| re-running under several `NCCL_ALGO` values and asserting bitwise equality | `FactorVariant.repeat_under` + `assert_order_is_topology_independent()` | + +## Checklist before the PR + +- [ ] File name equals the factor id minus the operator prefix. +- [ ] `question` is one line and says what the factor *answers*. +- [ ] Every representation-only field is `RECORD_ONLY`, not `MUST_MATCH_*`. +- [ ] A `SELF_WRITTEN` reference is justified against `SHARED_BACKEND` in the PR body. +- [ ] Every `RequiredSetting` has a `readback`, or you have accepted `UNOBSERVABLE`. +- [ ] `pinned_libraries` names an exact version, ideally with a container digest. +- [ ] Known pitfalls are encoded as `KnownPitfall`, not written in a comment. +- [ ] `python -m rl_engine.mismatch list` and `plan` both show what you expect. diff --git a/rl_engine/mismatch/engines/__init__.py b/rl_engine/mismatch/engines/__init__.py index 2a4f6e4a..0ed5ebf7 100644 --- a/rl_engine/mismatch/engines/__init__.py +++ b/rl_engine/mismatch/engines/__init__.py @@ -3,9 +3,24 @@ """The two sides under test: process and engine lifetime, configuration readback. -Shared across operators -- all of attention's factors use one ``vllm.py``. Not to -be confused with ``reference_adapters/``, which wires in *reference* -implementations and contains no operator code. +**This package holds exactly two things: ``megatron.py`` and ``vllm.py``.** They +are the training side and the rollout side as they really run -- how the engine +is constructed, how a switch is delivered to it, and how its *effective* value is +read back out. + +Shared across operators: all of attention's factors use one ``vllm.py``, and +adding an operator never adds a file here. + +Not to be confused with: + +* ``reference_adapters/`` -- wires *reference* implementations in and contains no + operator code; +* ``tests/mismatch_cpu_backend.py`` -- a harness that satisfies the same + ``ScoringBackend`` protocol so the framework can be exercised without a real + engine. It is not a side under test, so it does not live here. """ -from rl_engine.mismatch.engines.cpu_reference import CpuScoringBackend +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..6a2aab58 --- /dev/null +++ b/rl_engine/mismatch/engines/megatron.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The training side: Megatron-LM. **Placeholder -- not wired up yet.** + +What has to be built here: + +1. **Construction.** Model parallel state (TP/CP/PP) is fixed when the process + group is built, so anything touching it costs + ``RebindCost.PROCESS_GROUP_REBUILD``; determinism flags are read once at + process start, so they cost ``PROCESS_RESTART``. +2. **Scoring.** ``score()`` returns per-token logprobs for a fixed sequence via a + single forward over ``prompt + response`` -- ``ExecutionPath.TRAINING_FULL_PREFILL``, + the training side's only path. It must be **read-only**: fingerprint the model + tensors before and after and assert they match + (``pipeline/runner.py::assert_comparison_is_read_only``). A kernel that + mutates weights in place makes ``both_reference`` fail bitwise at random, and + the blame lands on the reference. +3. **Readback.** Off the live config objects, never from what was requested. + +Settings to pin and read back, with the pitfall each guards: + + AttnBackend -> explicit, never "auto" guards requested_not_actual + Megatron's auto picks a different kernel per shape, so requested and actual + must be recorded separately (attn.provenance). + + NVTE_ALLOW_NONDETERMINISTIC_ALGO -> "0" os.environ + CUBLAS_WORKSPACE_CONFIG -> ":4096:8" os.environ + NCCL_ALGO / NCCL_PROTO -> pinned guards nccl_algo_unpinned + + torch.backends.cuda.matmul.allow_tf32 -> False + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction -> False + torch.use_deterministic_algorithms(True) + + sequence_parallel -> record + With SP on, Megatron rewrites all_reduce into reduce_scatter + all_gather. + Whether each side applies that rewrite is gemm.forward_reduce. + + lm_head dtype -> fp32 + A BF16 head that is never upcast costs real logprob accuracy + (logp.precision_downcast). + +Under TP/CP no single rank holds the whole logprob vector: collect one +``LogprobShard`` per rank and let gate 3 check the count against ``world_size``. +A missing shard is wrong in a way that does not show. +""" + +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: + """Scores tokens on a live Megatron model. **Every method is a stub.**""" + + 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( + "Megatron scoring is not wired up: run one forward over " + "prompt + response, gather the selected logprobs, and assert the " + "model state fingerprint is unchanged across the call." + ) + + def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: + raise NotImplementedError( + "Group the 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 back off the live config. Requested is not " + "actual, and Megatron's AttnBackend=auto is exactly where that bites." + ) + + def collect_logprob_shards(self) -> tuple[LogprobShard, ...]: + raise NotImplementedError( + "Under TP/CP each rank holds one slice; return one shard per rank so " + "gate 3 can check the count against world_size." + ) + + +__all__ = ["MegatronBackend"] diff --git a/rl_engine/mismatch/engines/vllm.py b/rl_engine/mismatch/engines/vllm.py new file mode 100644 index 00000000..d1c328d8 --- /dev/null +++ b/rl_engine/mismatch/engines/vllm.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The rollout side: vLLM. **Placeholder -- not wired up yet.** + +What has to be built here, and why each part is not optional: + +1. **Construction.** vLLM decides several things at engine-build time, so a + switch at ``ENGINE_ARG`` level cannot be changed afterwards -- it costs a + rebuild, which is what ``RebindCost.ENGINE_REBUILD`` prices in. +2. **Scoring.** ``score()`` returns the per-token logprobs for a fixed sequence. + For the full-prefill path that means feeding ``prompt + response`` in as one + prompt and reading ``prompt_logprobs``; for the decode path it means a real + generation. Which one ran is ``ExecutionPath``, and the two are not equal in + floating point. +3. **Readback.** Every value below must come back off the live engine object. + A requested value is not evidence. + +The settings that must be pinned and read back, with the pitfall each guards +(see ``schema/pitfalls.py`` and the reference's ``required_settings``): + + enable_chunked_prefill -> False guards chunked_prefill_default_on + llm_engine.scheduler_config.chunked_prefill_enabled + Defaults to True. Leave it and you believe you measured a full-sequence + prefill while the engine chunked it. + + 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 + Non-zero splits long prefills by another route. + + enable_prefix_caching -> False guards prompt_logprobs_skips_prefix_cache + llm_engine.cache_config.enable_prefix_caching + prompt_logprobs already skips the prefix cache, so the full-prefill path + has no cache while the decode path does -- unless this is off, the two + paths are not comparable. + + all_reduce backend -> record actual, not requested + vLLM switches between custom IPC, MNNVL and NCCL by world size and + topology. Only what it actually chose is evidence (gemm.rollout_all_reduce_backend). + +Also to record: ``prompt_logprobs`` position 0 is ``None``, which is where the +shift-by-one convention has to be asserted against the training side rather than +assumed (pitfall ``shift_by_one``). +""" + +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: + """Scores tokens on a live vLLM engine. **Every method is a stub.**""" + + role: PolicyRole = PolicyRole.ROLLOUT + engine: Any = None # vllm.LLM once constructed + 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( + "vLLM scoring is not wired up: build the engine with the settings " + "pinned in this module's docstring, run the requested ExecutionPath, " + "and return (logprobs, readback) with the effective values read off " + "the engine." + ) + + def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: + raise NotImplementedError( + "Group the switches by tier so order_cases_by_rebind_cost() can reuse " + "engines: env/determinism -> 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 back off the live engine. A setting with no " + "readback path can only be recorded UNOBSERVABLE." + ) + + +__all__ = ["VllmBackend"] diff --git a/rl_engine/mismatch/operator_checks/__init__.py b/rl_engine/mismatch/operator_checks/__init__.py index 341ede21..b638834d 100644 --- a/rl_engine/mismatch/operator_checks/__init__.py +++ b/rl_engine/mismatch/operator_checks/__init__.py @@ -48,8 +48,6 @@ reference=TE_ROPE_REFERENCE, # from _common.py pitfalls=(KnownPitfall(...),), required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value, POSITION_CACHE), - owner="", - tracked_by=(), # issues or PRs following this factor ) What ``__init__.py`` contains 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..6b584268 --- /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 + +"""Attention operator plugin. Adding a factor drops a file into ``factors/``.""" + +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__) # scans factors/*.py + + 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..71941f35 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/_common.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared across the attention factors: reference implementations, contract helpers.""" + +from __future__ import annotations + +from rl_engine.mismatch.schema import ( + ExecutionPath, + LibraryPin, + ReferenceAuthority, + ReferenceImplementation, + RequiredSetting, + SettingChannel, +) + +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", # zhangj1an/TransformerEngine + ), + ), +) 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..d4270d1a --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/adapter.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The four operator-level methods for attention. **Interface only.** + +They are operator-level, not factor-level: how to read configuration back from an +engine and how to observe collectives is one piece of logic shared by all of +attention's factors. Pushing them down would make every factor file copy them. + +Whoever claims attention implements these four; the factor files under +``factors/`` stay declarations and do not change. +""" + +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: + """This side's switch values -> this side's numerical contract. + + Put each field at the path its factor declared in ``comparison_rules``: + ``precision.*`` for the precision profile, ``collectives[i].*`` for + communication, ``extra.*`` for everything attention-specific (rope_theta, + post-RoPE Q/K digests, GQA head mapping, page layout). Keep ``extra`` flat -- + nested dicts make the paths long and unreadable. + + No "collect the comparable fields" helper is needed: the framework indexes + both contracts by path and compares entry by entry. + """ + + raise NotImplementedError( + "attention.build_contract: return the OperatorContract for this side, " + "with fields at the paths declared in each factor's comparison_rules." + ) + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read the switches back **off the engine**. + + A requested value is not evidence. Megatron's ``AttnBackend=auto`` and vLLM's + backend selection both pick per shape, so requested and actual must be + recorded separately (pitfall ``requested_not_actual``). + """ + + raise NotImplementedError( + "attention.read_effective_config: read actual values off the live engine " + "and return them; anything with no readback path is UNOBSERVABLE." + ) + + +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """Which collectives this operator really performed, this run. + + Feeds the ``COLLECTIVE_CONTRACT`` evidence item. RoPE performs none; CP + attention performs several. Return what ran, not what was configured. + """ + + raise NotImplementedError( + "attention.observe_collectives: return the collectives that actually ran " + "(empty tuple for single-device paths)." + ) + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve the name an arm asks for into a callable. + + **Return the trace even when resolution fails** -- which candidates were + tried and why each was rejected. A bare ``None`` produces + ``SwitchStatus.FELL_BACK`` with nothing to investigate, and a fallen-back arm + whose deviation "did not change" reads exactly like a clean + ``NOT_THIS_FACTOR``. + + Candidates worth trying in order, for the RoPE reference: + ``transformer_engine.pytorch.attention.rope:apply_rotary_pos_emb`` on the + training side; ``flashinfer.rope:apply_rope`` then vLLM's + ``rotary_embedding:get_rope`` on the rollout side -- vLLM forwards to + FlashInfer itself when it is available, so which one answered is evidence. + """ + + raise NotImplementedError( + "attention.resolve_implementation: import the candidates in order and " + "return (callable, ImplementationResolution) -- including the rejection " + "trace on failure." + ) 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/rope_fusion.py b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py new file mode 100644 index 00000000..a55a27f1 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py @@ -0,0 +1,61 @@ +# 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 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, + "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "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), + 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/gemm/__init__.py b/rl_engine/mismatch/operator_checks/gemm/__init__.py new file mode 100644 index 00000000..e67e1f80 --- /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 + +"""GEMM operator plugin. Adding a factor drops a file into ``factors/``.""" + +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__) # scans factors/*.py + + 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..aeae4a3e --- /dev/null +++ b/rl_engine/mismatch/operator_checks/gemm/_common.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared across the gemm factors: collective contracts and the ordered reference.""" + +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 + +# What the reference pins: a reduce_scatter whose accumulation order is keyed on +# global rank, which is what makes the result independent of topology. +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 is rewritten into +# reduce_scatter + all_gather, and NCCL picks the order. +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, and the backend is +# chosen by world size and topology at runtime. +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", +) + +DETERMINISTIC_REDUCE_REFERENCE = ReferenceImplementation( + name="rl_kernel", + # SELF_WRITTEN because neither TE nor FlashInfer exposes a reduction whose + # order is fixed across topologies -- see the PR body. + 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, # the contract itself, pinned as a value + 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..7b88462f --- /dev/null +++ b/rl_engine/mismatch/operator_checks/gemm/adapter.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The four operator-level methods for gemm. **Interface only.** + +Whoever claims gemm implements these four; the factor files under ``factors/`` +stay declarations and do not change. +""" + +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: + """This side's switch values -> this side's numerical contract. + + For a reduction factor the load-bearing part is ``collectives``: the factors + index into it by path (``collectives[0].reduction_order`` and friends), so + the tuple order has to be stable and documented. ``_common.py`` already + declares the three contracts this operator switches between. + + Remember 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``. That difference is the factor, not a bug in + the contract. + """ + + raise NotImplementedError( + "gemm.build_contract: return the OperatorContract for this side, with the " + "collective this switch value selects at collectives[0]." + ) + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read the switches back **off the engine**. Requested is not actual.""" + + raise NotImplementedError( + "gemm.read_effective_config: read actual values off the live engine, " + "including which all-reduce backend vLLM really chose." + ) + + +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """The collectives this operator really performed, this run. + + Not what the config asked for: vLLM switches between custom IPC, MNNVL and + NCCL by world size and topology, and only what actually ran is evidence. + Feeds ``COLLECTIVE_CONTRACT``, which gate 2 requires before any verdict. + """ + + raise NotImplementedError( + "gemm.observe_collectives: return the collective trace for this run." + ) + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve the name an arm asks for into a callable, with the rejection trace. + + ``gemm.forward_reduce`` names ``rl_engine.kernels.collectives.ordered_reduce_scatter``, + the deterministic reduction that does not exist yet. Until it does, this must + fail *loudly* -- return the trace so gate 1 reports ``VARIANT_DID_NOT_APPLY`` + rather than letting a silent fallback read as a clean result. + """ + + raise NotImplementedError( + "gemm.resolve_implementation: import the requested implementation and " + "return (callable, ImplementationResolution) -- including the rejection " + "trace on failure." + ) 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..c21b0114 --- /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, declared explicitly 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, + "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, + # One factor, three physical sites: all row parallel linears, all eating the + # same accumulation-order problem. + 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..35c7834a --- /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 + +"""Logprob operator plugin. Adding a factor drops a file into ``factors/``.""" + +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__) # scans factors/*.py + + 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..1f03b8e6 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/_common.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared across the logprob factors. + +No ``ReferenceImplementation`` here yet: the rollout side has no deterministic +vocab-parallel reduction to swap in, so today's logprob factors are parameter +sweeps. When the self-written reference (kernel (4) of the design doc) lands, +it goes here and the factors that use it become four-arm swaps without any +other file changing. +""" + +from __future__ import annotations + +from rl_engine.mismatch.schema import DowncastPoint, Precision + +# vLLM computes logits at the model dtype and only upcasts for the softmax; +# Megatron can be told to keep the head in fp32. That difference is the switch. +HEAD_DTYPES: dict[str, Precision] = { + "bf16": Precision.BF16, + "fp32": Precision.FP32, +} + +# Where the fp32 accumulator is written back to a lower precision. +DOWNCAST_POINTS: dict[str, DowncastPoint] = { + "final_write": DowncastPoint.FINAL_WRITE, + "per_partial": DowncastPoint.PER_PARTIAL, +} 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..a4bad561 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/adapter.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The four operator-level methods for logprob. **Interface only.** + +Whoever claims logprob implements these four; the factor files under +``factors/`` stay declarations and do not change. +""" + +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: + """This side's switch values -> this side's numerical contract. + + Two asymmetries to encode rather than smooth over: + + * only the training side can vary the head dtype -- vLLM computes logits at + the model dtype, which is why ``logp.precision_downcast`` declares + ``applies_to=(TRAINING,)``; + * under TP the vocabulary is sharded, and MCore's ``VocabUtility`` and vLLM's + ``_get_indices`` do not agree on the boundaries for Qwen3. Record each + rank's ``(vocab_start, vocab_end)`` in ``extra`` -- comparing partial + results across different shard maps is meaningless. + """ + + raise NotImplementedError( + "logprob.build_contract: return the OperatorContract for this side, with " + "precision.lm_head set from the switch and the vocab shard map in extra." + ) + + +def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: + """Read the switches back **off the engine**. Requested is not actual.""" + + raise NotImplementedError( + "logprob.read_effective_config: read the head dtype, the transform chain " + "and the vocab shard boundaries off the live engine." + ) + + +def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: + """The collectives this operator really performed, this run. + + Empty at TP=1. Under vocab parallelism this is the partial-LSE reduction that + ``logp.reduce_topology`` and ``logp.merge_order`` are about. + """ + + raise NotImplementedError( + "logprob.observe_collectives: return the collective trace for this run " + "(empty tuple at TP=1)." + ) + + +def resolve_implementation( + factor_id: str, role: PolicyRole, impl_name: str +) -> tuple[Callable[..., Any] | None, ImplementationResolution]: + """Resolve the name an arm asks for into a callable, with the rejection trace. + + A parameter sweep never asks for one, so this stays unused until a logprob + factor gains a reference -- the deterministic vocab-parallel reduction, which + is not written yet. + """ + + raise NotImplementedError( + "logprob.resolve_implementation: import the requested implementation and " + "return (callable, ImplementationResolution) -- including the rejection " + "trace on failure." + ) 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/precision_downcast.py b/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py new file mode 100644 index 00000000..b39d8c42 --- /dev/null +++ b/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""logp.precision_downcast -- lm_head dtype and where fp32 is written back.""" + +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?" + ), + # A parameter sweep, not an implementation swap: nothing is replaced, the + # head dtype is scanned. ``reference=None`` is what says so -- there is no + # separate "kind" field. + switch=Switch( + path="logp.head_dtype", + # Cheapest tier: the dtype is a call argument, so arms of this factor + # reuse one engine and sort ahead of every rebuild-level case. + rebind_cost=RebindCost.PER_REQUEST, + applies_to=(PolicyRole.TRAINING,), + allowed_values=tuple(HEAD_DTYPES), + ), + comparison_rules={ + # The head GEMM's output dtype is identity, not a matter of taste. + "precision.lm_head": ComparisonRule.MUST_MATCH_BITWISE, + "precision.accumulate": ComparisonRule.MUST_MATCH_SEMANTICALLY, + "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, + # Representation: both sides may name their mode differently. + "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/schema/factors.py b/rl_engine/mismatch/schema/factors.py index 5eedc653..961d6002 100644 --- a/rl_engine/mismatch/schema/factors.py +++ b/rl_engine/mismatch/schema/factors.py @@ -165,8 +165,6 @@ class MismatchFactor: call_sites: tuple[str, ...] = () # one factor acting in several places pitfalls: tuple[KnownPitfall, ...] = () variants: tuple[Any, ...] = () # empty -> expand the standard set - owner: str = "" # "" = unclaimed - tracked_by: tuple[str, ...] = () # issues or PRs following this factor def declared_collectives(factor: MismatchFactor) -> tuple[CollectiveContract, ...]: diff --git a/rl_engine/mismatch/engines/cpu_reference.py b/tests/mismatch_cpu_backend.py similarity index 100% rename from rl_engine/mismatch/engines/cpu_reference.py rename to tests/mismatch_cpu_backend.py diff --git a/tests/test_mismatch_framework.py b/tests/test_mismatch_framework.py index 79b493e9..c9d2cb09 100644 --- a/tests/test_mismatch_framework.py +++ b/tests/test_mismatch_framework.py @@ -17,7 +17,7 @@ import pytest -from rl_engine.mismatch.engines.cpu_reference import CpuScoringBackend +from tests.mismatch_cpu_backend import CpuScoringBackend from rl_engine.mismatch.pipeline import ( ContradictoryFactor, PluginRegistry, From 9213e7196e90b361d32fa65721212c42a4f054c5 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sun, 9 Aug 2026 12:19:22 +0000 Subject: [PATCH 3/7] refactor(mismatch): cut comments back to what the code cannot say Applied Clean Code chapter 4 across the package. The docstrings had grown into design documents: rationale essays on types, restatements of the signature, and system-wide explanation attached to one local declaration -- the "too much information" and "nonlocal information" smells. That material belongs in README.md and docs/, where it already is. Removed roughly 530 lines. What survives is what the code cannot say for itself: why a field is RECORD_ONLY, why a silent fallback is more dangerous than an error, why thresholds are constants rather than configuration. One executable change: declared_collectives() drops an intermediate variable that only repeated the function name. Everything else is comments -- verified by comparing every module's AST with docstrings stripped. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01QjJLV6aQZTtm4X6pjex8Ri --- rl_engine/mismatch/__init__.py | 23 ++--- rl_engine/mismatch/__main__.py | 15 ++- rl_engine/mismatch/engines/__init__.py | 21 ++--- rl_engine/mismatch/engines/megatron.py | 80 +++++----------- rl_engine/mismatch/engines/vllm.py | 78 +++++----------- rl_engine/mismatch/model_meta/__init__.py | 8 +- rl_engine/mismatch/model_meta/qwen3.py | 5 +- .../mismatch/operator_checks/__init__.py | 82 +++-------------- .../operator_checks/attention/__init__.py | 4 +- .../operator_checks/attention/_common.py | 8 +- .../operator_checks/attention/adapter.py | 68 ++++---------- .../attention/factors/rope_fusion.py | 2 + .../mismatch/operator_checks/gemm/__init__.py | 4 +- .../mismatch/operator_checks/gemm/_common.py | 18 ++-- .../mismatch/operator_checks/gemm/adapter.py | 59 ++++-------- .../gemm/factors/forward_reduce.py | 10 +- .../operator_checks/logprob/__init__.py | 4 +- .../operator_checks/logprob/_common.py | 13 +-- .../operator_checks/logprob/adapter.py | 56 +++-------- .../logprob/factors/precision_downcast.py | 14 ++- rl_engine/mismatch/pipeline/__init__.py | 19 ++-- rl_engine/mismatch/pipeline/comparison.py | 20 ++-- rl_engine/mismatch/pipeline/diagnosis.py | 25 +++-- rl_engine/mismatch/pipeline/planner.py | 34 +++---- rl_engine/mismatch/pipeline/registry.py | 38 ++++---- rl_engine/mismatch/pipeline/report.py | 18 ++-- rl_engine/mismatch/pipeline/runner.py | 29 ++---- .../mismatch/reference_adapters/__init__.py | 23 +---- .../mismatch/reference_adapters/settings.py | 16 ++-- rl_engine/mismatch/schema/__init__.py | 30 ++---- rl_engine/mismatch/schema/collectives.py | 37 +++----- rl_engine/mismatch/schema/contracts.py | 27 ++---- rl_engine/mismatch/schema/factors.py | 92 ++++++------------- rl_engine/mismatch/schema/fingerprints.py | 40 +++----- rl_engine/mismatch/schema/metrics.py | 54 ++++------- rl_engine/mismatch/schema/pitfalls.py | 26 ++---- rl_engine/mismatch/schema/rollout_context.py | 57 ++++-------- rl_engine/mismatch/schema/thresholds.py | 12 +-- rl_engine/mismatch/schema/tracing.py | 63 +++---------- rl_engine/mismatch/schema/values.py | 67 ++++---------- rl_engine/mismatch/schema/variants.py | 85 +++++++---------- tests/mismatch_cpu_backend.py | 33 +++---- 42 files changed, 439 insertions(+), 978 deletions(-) diff --git a/rl_engine/mismatch/__init__.py b/rl_engine/mismatch/__init__.py index 218062dc..a8284b3c 100644 --- a/rl_engine/mismatch/__init__.py +++ b/rl_engine/mismatch/__init__.py @@ -4,26 +4,19 @@ """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 framework turns "which of the dozens of possible causes -is it" into a set of switches that can be flipped one at a time and attributed. +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. -Layout:: +Three dependency rules keep the plugin seam open: - schema/ pure data types, no behaviour - pipeline/ the seven execution steps, free functions only - engines/ the two sides under test - reference_adapters/ wiring reference implementations in - model_meta/ model shape and module correspondence - operator_checks/ plugins, one directory per operator (empty by design) - -Dependency rules: - -1. ``schema/`` never imports ``pipeline/``; within ``schema/`` the imports go one - way only, with ``values.py`` at the top importing nothing from the project. -2. ``pipeline/`` never imports ``operator_checks/`` -- it only sees what the +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 diff --git a/rl_engine/mismatch/__main__.py b/rl_engine/mismatch/__main__.py index a6bf2c39..8fd3c566 100644 --- a/rl_engine/mismatch/__main__.py +++ b/rl_engine/mismatch/__main__.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Command line entry point. **The only place that imports operator plugins.** +"""Command line entry point, and the only place that imports operator plugins. -The framework does not know which operators exist -- importing them here is what -triggers self-registration, and it keeps the dependency arrow pointing one way: -``pipeline/`` never reaches into ``operator_checks/``. +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 @@ -25,8 +25,8 @@ ) from rl_engine.mismatch.schema import NoiseFloor -# Operator plugins are listed here and nowhere else. An operator that is not -# imported simply does not exist as far as the framework is concerned. +# 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", @@ -94,9 +94,6 @@ def command_plan(operator: str | None, noise_floor: str, gpu_count: int, as_json print("nothing to plan: no operator plugins are registered.") return 0 - # Static rejection before anything is expanded: a factor that claims topology - # independence while reducing non-deterministically produces numbers that - # mean nothing, and that is knowable without running. reject_contradictory_factors(factors) runnable = [] diff --git a/rl_engine/mismatch/engines/__init__.py b/rl_engine/mismatch/engines/__init__.py index 0ed5ebf7..33a835b2 100644 --- a/rl_engine/mismatch/engines/__init__.py +++ b/rl_engine/mismatch/engines/__init__.py @@ -1,23 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The two sides under test: process and engine lifetime, configuration readback. +"""The two sides under test: engine lifetime and configuration readback. -**This package holds exactly two things: ``megatron.py`` and ``vllm.py``.** They -are the training side and the rollout side as they really run -- how the engine -is constructed, how a switch is delivered to it, and how its *effective* value is -read back out. +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``. -Shared across operators: all of attention's factors use one ``vllm.py``, and -adding an operator never adds a file here. - -Not to be confused with: - -* ``reference_adapters/`` -- wires *reference* implementations in and contains no - operator code; -* ``tests/mismatch_cpu_backend.py`` -- a harness that satisfies the same - ``ScoringBackend`` protocol so the framework can be exercised without a real - engine. It is not a side under test, so it does not live here. +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 diff --git a/rl_engine/mismatch/engines/megatron.py b/rl_engine/mismatch/engines/megatron.py index 6a2aab58..f2f788e0 100644 --- a/rl_engine/mismatch/engines/megatron.py +++ b/rl_engine/mismatch/engines/megatron.py @@ -1,48 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The training side: Megatron-LM. **Placeholder -- not wired up yet.** - -What has to be built here: - -1. **Construction.** Model parallel state (TP/CP/PP) is fixed when the process - group is built, so anything touching it costs - ``RebindCost.PROCESS_GROUP_REBUILD``; determinism flags are read once at - process start, so they cost ``PROCESS_RESTART``. -2. **Scoring.** ``score()`` returns per-token logprobs for a fixed sequence via a - single forward over ``prompt + response`` -- ``ExecutionPath.TRAINING_FULL_PREFILL``, - the training side's only path. It must be **read-only**: fingerprint the model - tensors before and after and assert they match - (``pipeline/runner.py::assert_comparison_is_read_only``). A kernel that - mutates weights in place makes ``both_reference`` fail bitwise at random, and - the blame lands on the reference. -3. **Readback.** Off the live config objects, never from what was requested. - -Settings to pin and read back, with the pitfall each guards: - - AttnBackend -> explicit, never "auto" guards requested_not_actual - Megatron's auto picks a different kernel per shape, so requested and actual - must be recorded separately (attn.provenance). - - NVTE_ALLOW_NONDETERMINISTIC_ALGO -> "0" os.environ - CUBLAS_WORKSPACE_CONFIG -> ":4096:8" os.environ - NCCL_ALGO / NCCL_PROTO -> pinned guards nccl_algo_unpinned - - torch.backends.cuda.matmul.allow_tf32 -> False - torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction -> False - torch.use_deterministic_algorithms(True) - - sequence_parallel -> record - With SP on, Megatron rewrites all_reduce into reduce_scatter + all_gather. - Whether each side applies that rewrite is gemm.forward_reduce. - - lm_head dtype -> fp32 - A BF16 head that is never upcast costs real logprob accuracy - (logp.precision_downcast). - -Under TP/CP no single rank holds the whole logprob vector: collect one -``LogprobShard`` per rank and let gate 3 check the count against ``world_size``. -A missing shard is wrong in a way that does not show. +"""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 @@ -55,8 +32,6 @@ @dataclass class MegatronBackend: - """Scores tokens on a live Megatron model. **Every method is a stub.**""" - role: PolicyRole = PolicyRole.TRAINING model: Any = None effective_config: dict[str, Any] = field(default_factory=dict) @@ -69,28 +44,21 @@ def score( replacement: Callable[..., Any] | None, ) -> tuple[Sequence[float], Mapping[str, Any]]: raise NotImplementedError( - "Megatron scoring is not wired up: run one forward over " - "prompt + response, gather the selected logprobs, and assert the " - "model state fingerprint is unchanged across the call." + "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 the switches by tier: determinism env -> process, TP/CP/PP -> " + "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 back off the live config. Requested is not " - "actual, and Megatron's AttnBackend=auto is exactly where that bites." - ) + raise NotImplementedError("Read each pinned setting off the live config.") def collect_logprob_shards(self) -> tuple[LogprobShard, ...]: - raise NotImplementedError( - "Under TP/CP each rank holds one slice; return one shard per rank so " - "gate 3 can check the count against world_size." - ) + 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 index d1c328d8..08678388 100644 --- a/rl_engine/mismatch/engines/vllm.py +++ b/rl_engine/mismatch/engines/vllm.py @@ -1,49 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The rollout side: vLLM. **Placeholder -- not wired up yet.** - -What has to be built here, and why each part is not optional: - -1. **Construction.** vLLM decides several things at engine-build time, so a - switch at ``ENGINE_ARG`` level cannot be changed afterwards -- it costs a - rebuild, which is what ``RebindCost.ENGINE_REBUILD`` prices in. -2. **Scoring.** ``score()`` returns the per-token logprobs for a fixed sequence. - For the full-prefill path that means feeding ``prompt + response`` in as one - prompt and reading ``prompt_logprobs``; for the decode path it means a real - generation. Which one ran is ``ExecutionPath``, and the two are not equal in - floating point. -3. **Readback.** Every value below must come back off the live engine object. - A requested value is not evidence. - -The settings that must be pinned and read back, with the pitfall each guards -(see ``schema/pitfalls.py`` and the reference's ``required_settings``): - - enable_chunked_prefill -> False guards chunked_prefill_default_on - llm_engine.scheduler_config.chunked_prefill_enabled - Defaults to True. Leave it and you believe you measured a full-sequence - prefill while the engine chunked it. - - 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 - Non-zero splits long prefills by another route. - - enable_prefix_caching -> False guards prompt_logprobs_skips_prefix_cache - llm_engine.cache_config.enable_prefix_caching - prompt_logprobs already skips the prefix cache, so the full-prefill path - has no cache while the decode path does -- unless this is off, the two - paths are not comparable. - - all_reduce backend -> record actual, not requested - vLLM switches between custom IPC, MNNVL and NCCL by world size and - topology. Only what it actually chose is evidence (gemm.rollout_all_reduce_backend). - -Also to record: ``prompt_logprobs`` position 0 is ``None``, which is where the -shift-by-one convention has to be asserted against the training side rather than -assumed (pitfall ``shift_by_one``). +"""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 @@ -56,10 +29,8 @@ @dataclass class VllmBackend: - """Scores tokens on a live vLLM engine. **Every method is a stub.**""" - role: PolicyRole = PolicyRole.ROLLOUT - engine: Any = None # vllm.LLM once constructed + engine: Any = None effective_config: dict[str, Any] = field(default_factory=dict) def score( @@ -70,23 +41,20 @@ def score( replacement: Callable[..., Any] | None, ) -> tuple[Sequence[float], Mapping[str, Any]]: raise NotImplementedError( - "vLLM scoring is not wired up: build the engine with the settings " - "pinned in this module's docstring, run the requested ExecutionPath, " - "and return (logprobs, readback) with the effective values read off " - "the engine." + "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 the switches by tier so order_cases_by_rebind_cost() can reuse " - "engines: env/determinism -> process, world size and TP/CP -> " - "process_group, dtype and backend and KV layout -> engine, batch and " - "sequence -> request." + "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 back off the live engine. A setting with no " + "Read each pinned setting off the live engine. A setting with no " "readback path can only be recorded UNOBSERVABLE." ) diff --git a/rl_engine/mismatch/model_meta/__init__.py b/rl_engine/mismatch/model_meta/__init__.py index df83c59c..cb66664b 100644 --- a/rl_engine/mismatch/model_meta/__init__.py +++ b/rl_engine/mismatch/model_meta/__init__.py @@ -3,10 +3,10 @@ """Model shape and the module correspondence table. -Filled once per model by whoever brings it up, 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 GLM5 or DSv3.2 means -redoing it -- which is why it lives here and not in any ``operator_checks/``. +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 diff --git a/rl_engine/mismatch/model_meta/qwen3.py b/rl_engine/mismatch/model_meta/qwen3.py index 0ad09f64..a7673c47 100644 --- a/rl_engine/mismatch/model_meta/qwen3.py +++ b/rl_engine/mismatch/model_meta/qwen3.py @@ -3,9 +3,8 @@ """Qwen3 module correspondence and call chain. -Every entry with an ``equivalence`` also carries ``verified_by``: **an -equivalence may not be claimed without a test that proves it**, otherwise -"filtering false positives" quietly becomes "hiding real findings". +Every ``equivalence`` carries a ``verified_by``: unproven, "filtering false +positives" quietly becomes "hiding real findings". """ from __future__ import annotations diff --git a/rl_engine/mismatch/operator_checks/__init__.py b/rl_engine/mismatch/operator_checks/__init__.py index b638834d..9964e360 100644 --- a/rl_engine/mismatch/operator_checks/__init__.py +++ b/rl_engine/mismatch/operator_checks/__init__.py @@ -1,80 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Operator plugins, one directory per operator. **Empty by design.** +"""Operator plugins, one directory per operator. -The framework ships without operators. Each operator is claimed and written -separately, and adding one changes nothing outside its own directory. +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. -Importing this package registers nothing on its own -- ``__main__`` imports the -individual operator packages, and that import is what triggers -self-registration. The framework does not know which operators exist; the CLI -brings them in. - -Layout of one operator ----------------------- - -An operator has more than a dozen factors, each with eight or so fields. Putting -them all in one file runs to 800 lines that nobody wants to edit, so it is **one -file per factor**, mirroring "one directory per operator":: +One operator's layout, and the conventions that are enforced:: operator_checks/attention/ - |-- __init__.py ~15 lines: operator name + discover_factors, rest delegated + |-- __init__.py operator name + discover_factors |-- adapter.py the four operator-level methods - |-- _common.py shared across factors: contract helpers, common - | comparison rules, this operator's reference implementations - `-- factors/ one file per factor, 30-50 lines each - |-- provenance.py attn.provenance - |-- execution_path.py attn.execution_path - `-- ... - -**File name equals the factor id with the operator prefix stripped.** -``discover_factors()`` enforces it, so renaming an id without renaming the file -fails at import rather than silently. - -What a factor file contains ---------------------------- - -One ``FACTOR`` constant and nothing else -- declarations, no behaviour:: - - FACTOR = MismatchFactor( - id="attn.rope_fusion", - operator="attention", - category=FactorCategory.KERNEL_IMPLEMENTATION, - question="...", - switch=Switch(...), - comparison_rules={...}, - prerequisites=Prerequisites(...), - reference=TE_ROPE_REFERENCE, # from _common.py - pitfalls=(KnownPitfall(...),), - required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value, POSITION_CACHE), - ) - -What ``__init__.py`` contains ------------------------------ - -:: - - @OPERATOR_CHECKS.register - class AttentionChecks: - operator = "attention" - - def declare_factors(self): - return discover_factors(__package__) # scans factors/*.py - - build_contract = adapter.build_contract - read_effective_config = adapter.read_effective_config - observe_collectives = adapter.observe_collectives - resolve_implementation = adapter.resolve_implementation - -Adding a factor means dropping a file into ``factors/``; ``__init__.py`` does not -change. + |-- _common.py shared reference implementations and contract helpers + `-- factors/ one file per factor, holding one FACTOR constant -Why the four methods stay in ``adapter.py`` -------------------------------------------- +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. -They are **operator-level**, not factor-level: how to read configuration back -from an engine and how to observe collectives is one piece of logic shared by all -of that operator's factors. Pushing them down would make every factor file copy -them. +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 index 6b584268..cac2bf21 100644 --- a/rl_engine/mismatch/operator_checks/attention/__init__.py +++ b/rl_engine/mismatch/operator_checks/attention/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Attention operator plugin. Adding a factor drops a file into ``factors/``.""" +"""The attention operator plugin.""" from rl_engine.mismatch.operator_checks.attention import adapter from rl_engine.mismatch.pipeline import OPERATOR_CHECKS, discover_factors @@ -12,7 +12,7 @@ class AttentionChecks: operator = "attention" def declare_factors(self): - return discover_factors(__package__) # scans factors/*.py + return discover_factors(__package__) build_contract = staticmethod(adapter.build_contract) read_effective_config = staticmethod(adapter.read_effective_config) diff --git a/rl_engine/mismatch/operator_checks/attention/_common.py b/rl_engine/mismatch/operator_checks/attention/_common.py index 71941f35..8754b8cb 100644 --- a/rl_engine/mismatch/operator_checks/attention/_common.py +++ b/rl_engine/mismatch/operator_checks/attention/_common.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Shared across the attention factors: reference implementations, contract helpers.""" +"""Reference implementations shared by attention's factors.""" from __future__ import annotations @@ -32,10 +32,6 @@ ), ), pinned_libraries=( - LibraryPin( - "transformer_engine", - "2.9.0.dev0", - commit="8260f49", # zhangj1an/TransformerEngine - ), + LibraryPin("transformer_engine", "2.9.0.dev0", commit="8260f49"), ), ) diff --git a/rl_engine/mismatch/operator_checks/attention/adapter.py b/rl_engine/mismatch/operator_checks/attention/adapter.py index d4270d1a..9caf36f9 100644 --- a/rl_engine/mismatch/operator_checks/attention/adapter.py +++ b/rl_engine/mismatch/operator_checks/attention/adapter.py @@ -1,14 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The four operator-level methods for attention. **Interface only.** +"""Attention's four operator-level methods. Not implemented. -They are operator-level, not factor-level: how to read configuration back from an -engine and how to observe collectives is one piece of logic shared by all of -attention's factors. Pushing them down would make every factor file copy them. - -Whoever claims attention implements these four; the factor files under -``factors/`` stay declarations and do not change. +They are operator-level rather than factor-level because reading configuration +back from an engine is the same logic for all of attention's factors. """ from __future__ import annotations @@ -24,71 +20,41 @@ def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: - """This side's switch values -> this side's numerical contract. - - Put each field at the path its factor declared in ``comparison_rules``: - ``precision.*`` for the precision profile, ``collectives[i].*`` for - communication, ``extra.*`` for everything attention-specific (rope_theta, - post-RoPE Q/K digests, GQA head mapping, page layout). Keep ``extra`` flat -- - nested dicts make the paths long and unreadable. - - No "collect the comparable fields" helper is needed: the framework indexes - both contracts by path and compares entry by entry. - """ - raise NotImplementedError( - "attention.build_contract: return the OperatorContract for this side, " - "with fields at the paths declared in each factor's comparison_rules." + "Return this side's contract with each field at the path its factor " + "declared in comparison_rules: precision.*, collectives[i].*, and " + "attention's own state under a flat extra.*" ) def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: - """Read the switches back **off the engine**. - - A requested value is not evidence. Megatron's ``AttnBackend=auto`` and vLLM's - backend selection both pick per shape, so requested and actual must be - recorded separately (pitfall ``requested_not_actual``). - """ - raise NotImplementedError( - "attention.read_effective_config: read actual values off the live engine " - "and return them; anything with no readback path is UNOBSERVABLE." + "Read actual values off the live engine. Megatron's AttnBackend=auto and " + "vLLM's backend selection both pick per shape, so a requested value is " + "not evidence." ) def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: - """Which collectives this operator really performed, this run. - - Feeds the ``COLLECTIVE_CONTRACT`` evidence item. RoPE performs none; CP - attention performs several. Return what ran, not what was configured. - """ - - raise NotImplementedError( - "attention.observe_collectives: return the collectives that actually ran " - "(empty tuple for single-device paths)." - ) + raise NotImplementedError("Return the collectives that actually ran.") def resolve_implementation( factor_id: str, role: PolicyRole, impl_name: str ) -> tuple[Callable[..., Any] | None, ImplementationResolution]: - """Resolve the name an arm asks for into a callable. + """Resolve an arm's implementation name, returning the trace even on failure. - **Return the trace even when resolution fails** -- which candidates were - tried and why each was rejected. A bare ``None`` produces - ``SwitchStatus.FELL_BACK`` with nothing to investigate, and a fallen-back arm - whose deviation "did not change" reads exactly like a clean + 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``. - Candidates worth trying in order, for the RoPE reference: + Candidates for the RoPE reference, in order: ``transformer_engine.pytorch.attention.rope:apply_rotary_pos_emb`` on the training side; ``flashinfer.rope:apply_rope`` then vLLM's - ``rotary_embedding:get_rope`` on the rollout side -- vLLM forwards to - FlashInfer itself when it is available, so which one answered is evidence. + ``rotary_embedding:get_rope`` on the rollout side. vLLM forwards to + FlashInfer when it is available, so which one answered is itself evidence. """ raise NotImplementedError( - "attention.resolve_implementation: import the candidates in order and " - "return (callable, ImplementationResolution) -- including the rejection " - "trace on failure." + "Import the candidates in order and return (callable, resolution)." ) diff --git a/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py index a55a27f1..5050b9b5 100644 --- a/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py +++ b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py @@ -40,6 +40,8 @@ "extra.position_ids_digest": ComparisonRule.MUST_MATCH_SEMANTICALLY, "extra.post_rope_qk_digest": 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( diff --git a/rl_engine/mismatch/operator_checks/gemm/__init__.py b/rl_engine/mismatch/operator_checks/gemm/__init__.py index e67e1f80..41d6ba31 100644 --- a/rl_engine/mismatch/operator_checks/gemm/__init__.py +++ b/rl_engine/mismatch/operator_checks/gemm/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""GEMM operator plugin. Adding a factor drops a file into ``factors/``.""" +"""The gemm operator plugin.""" from rl_engine.mismatch.operator_checks.gemm import adapter from rl_engine.mismatch.pipeline import OPERATOR_CHECKS, discover_factors @@ -12,7 +12,7 @@ class GemmChecks: operator = "gemm" def declare_factors(self): - return discover_factors(__package__) # scans factors/*.py + return discover_factors(__package__) build_contract = staticmethod(adapter.build_contract) read_effective_config = staticmethod(adapter.read_effective_config) diff --git a/rl_engine/mismatch/operator_checks/gemm/_common.py b/rl_engine/mismatch/operator_checks/gemm/_common.py index aeae4a3e..754c4710 100644 --- a/rl_engine/mismatch/operator_checks/gemm/_common.py +++ b/rl_engine/mismatch/operator_checks/gemm/_common.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Shared across the gemm factors: collective contracts and the ordered reference.""" +"""Collective contracts and the ordered reference, shared by gemm's factors.""" from __future__ import annotations @@ -23,8 +23,6 @@ TP_SIZE = 2 -# What the reference pins: a reduce_scatter whose accumulation order is keyed on -# global rank, which is what makes the result independent of topology. ORDERED_REDUCE_SCATTER = CollectiveContract( op=CollectiveOp.REDUCE_SCATTER, group=ParallelDim.TENSOR, @@ -36,8 +34,8 @@ backend="rl_kernel", ) -# Megatron with sequence parallelism on: all_reduce is rewritten into -# reduce_scatter + all_gather, and NCCL picks the order. +# 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, @@ -49,8 +47,8 @@ backend="nccl", ) -# vLLM without sequence parallelism: a plain all_reduce, and the backend is -# chosen by world size and topology at runtime. +# 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, @@ -62,10 +60,10 @@ 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", - # SELF_WRITTEN because neither TE nor FlashInfer exposes a reduction whose - # order is fixed across topologies -- see the PR body. tier=ReferenceAuthority.SELF_WRITTEN, training_impl="rl_engine.kernels.collectives.ordered_reduce_scatter", rollout_impl="rl_engine.kernels.collectives.ordered_reduce_scatter", @@ -76,7 +74,7 @@ required_settings=( RequiredSetting( "forward_reduce_contract", - ORDERED_REDUCE_SCATTER, # the contract itself, pinned as a value + ORDERED_REDUCE_SCATTER, SettingChannel.CALL_ARG, readback="module.last_collective_contract", ), diff --git a/rl_engine/mismatch/operator_checks/gemm/adapter.py b/rl_engine/mismatch/operator_checks/gemm/adapter.py index 7b88462f..81b0a5c1 100644 --- a/rl_engine/mismatch/operator_checks/gemm/adapter.py +++ b/rl_engine/mismatch/operator_checks/gemm/adapter.py @@ -1,11 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The four operator-level methods for gemm. **Interface only.** - -Whoever claims gemm implements these four; the factor files under ``factors/`` -stay declarations and do not change. -""" +"""GEMM's four operator-level methods. Not implemented.""" from __future__ import annotations @@ -20,60 +16,43 @@ def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: - """This side's switch values -> this side's numerical contract. - - For a reduction factor the load-bearing part is ``collectives``: the factors - index into it by path (``collectives[0].reduction_order`` and friends), so - the tuple order has to be stable and documented. ``_common.py`` already - declares the three contracts this operator switches between. + """Return this side's contract, with the collective the switch selects first. - Remember 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``. That difference is the factor, not a bug in - the contract. + 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( - "gemm.build_contract: return the OperatorContract for this side, with the " - "collective this switch value selects at collectives[0]." - ) + raise NotImplementedError def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: - """Read the switches back **off the engine**. Requested is not actual.""" - raise NotImplementedError( - "gemm.read_effective_config: read actual values off the live engine, " - "including which all-reduce backend vLLM really chose." + "Read actual values off the live engine, including which all-reduce " + "backend vLLM really chose." ) def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: - """The collectives this operator really performed, this run. + """Return the collective trace for this run. - Not what the config asked for: vLLM switches between custom IPC, MNNVL and - NCCL by world size and topology, and only what actually ran is evidence. - Feeds ``COLLECTIVE_CONTRACT``, which gate 2 requires before any verdict. + vLLM switches between custom IPC, MNNVL and NCCL by world size and topology, + so only what ran is evidence. """ - raise NotImplementedError( - "gemm.observe_collectives: return the collective trace for this run." - ) + raise NotImplementedError def resolve_implementation( factor_id: str, role: PolicyRole, impl_name: str ) -> tuple[Callable[..., Any] | None, ImplementationResolution]: - """Resolve the name an arm asks for into a callable, with the rejection trace. + """Resolve an arm's implementation name, returning the trace even on failure. - ``gemm.forward_reduce`` names ``rl_engine.kernels.collectives.ordered_reduce_scatter``, - the deterministic reduction that does not exist yet. Until it does, this must - fail *loudly* -- return the trace so gate 1 reports ``VARIANT_DID_NOT_APPLY`` - rather than letting a silent fallback read as a clean result. + ``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( - "gemm.resolve_implementation: import the requested implementation and " - "return (callable, ImplementationResolution) -- including the rejection " - "trace on failure." - ) + raise NotImplementedError diff --git a/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py b/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py index c21b0114..3722080a 100644 --- a/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py +++ b/rl_engine/mismatch/operator_checks/gemm/factors/forward_reduce.py @@ -25,9 +25,9 @@ _REF = DETERMINISTIC_REDUCE_REFERENCE -# The standard four arms, declared explicitly so the self-check gate can also -# carry ``repeat_under``: an implementation claiming topology independence must -# survive NCCL choosing a different algorithm. +# 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", @@ -78,13 +78,13 @@ "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, - # One factor, three physical sites: all row parallel linears, all eating the - # same accumulation-order problem. + # All three are row parallel linears eating the same accumulation order. call_sites=("attention.o_linear", "mlp.down_linear", "moe.output"), variants=_VARIANTS, pitfalls=( diff --git a/rl_engine/mismatch/operator_checks/logprob/__init__.py b/rl_engine/mismatch/operator_checks/logprob/__init__.py index 35c7834a..1019ceb4 100644 --- a/rl_engine/mismatch/operator_checks/logprob/__init__.py +++ b/rl_engine/mismatch/operator_checks/logprob/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Logprob operator plugin. Adding a factor drops a file into ``factors/``.""" +"""The logprob operator plugin.""" from rl_engine.mismatch.operator_checks.logprob import adapter from rl_engine.mismatch.pipeline import OPERATOR_CHECKS, discover_factors @@ -12,7 +12,7 @@ class LogprobChecks: operator = "logprob" def declare_factors(self): - return discover_factors(__package__) # scans factors/*.py + return discover_factors(__package__) build_contract = staticmethod(adapter.build_contract) read_effective_config = staticmethod(adapter.read_effective_config) diff --git a/rl_engine/mismatch/operator_checks/logprob/_common.py b/rl_engine/mismatch/operator_checks/logprob/_common.py index 1f03b8e6..f2475bbc 100644 --- a/rl_engine/mismatch/operator_checks/logprob/_common.py +++ b/rl_engine/mismatch/operator_checks/logprob/_common.py @@ -1,27 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Shared across the logprob factors. +"""Values shared by logprob's factors. -No ``ReferenceImplementation`` here yet: the rollout side has no deterministic -vocab-parallel reduction to swap in, so today's logprob factors are parameter -sweeps. When the self-written reference (kernel (4) of the design doc) lands, -it goes here and the factors that use it become four-arm swaps without any -other file changing. +No reference implementation yet: the rollout side has no deterministic +vocab-parallel reduction to swap in, so these factors are parameter sweeps. """ from __future__ import annotations from rl_engine.mismatch.schema import DowncastPoint, Precision -# vLLM computes logits at the model dtype and only upcasts for the softmax; -# Megatron can be told to keep the head in fp32. That difference is the switch. +# 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, } -# Where the fp32 accumulator is written back to a lower precision. DOWNCAST_POINTS: dict[str, DowncastPoint] = { "final_write": DowncastPoint.FINAL_WRITE, "per_partial": DowncastPoint.PER_PARTIAL, diff --git a/rl_engine/mismatch/operator_checks/logprob/adapter.py b/rl_engine/mismatch/operator_checks/logprob/adapter.py index a4bad561..142d6de8 100644 --- a/rl_engine/mismatch/operator_checks/logprob/adapter.py +++ b/rl_engine/mismatch/operator_checks/logprob/adapter.py @@ -1,11 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The four operator-level methods for logprob. **Interface only.** - -Whoever claims logprob implements these four; the factor files under -``factors/`` stay declarations and do not change. -""" +"""Logprob's four operator-level methods. Not implemented.""" from __future__ import annotations @@ -20,59 +16,33 @@ def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: - """This side's switch values -> this side's numerical contract. - - Two asymmetries to encode rather than smooth over: + """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, which is why ``logp.precision_downcast`` declares - ``applies_to=(TRAINING,)``; - * under TP the vocabulary is sharded, and MCore's ``VocabUtility`` and vLLM's - ``_get_indices`` do not agree on the boundaries for Qwen3. Record each - rank's ``(vocab_start, vocab_end)`` in ``extra`` -- comparing partial - results across different shard maps is meaningless. + 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. """ - raise NotImplementedError( - "logprob.build_contract: return the OperatorContract for this side, with " - "precision.lm_head set from the switch and the vocab shard map in extra." - ) + raise NotImplementedError def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: - """Read the switches back **off the engine**. Requested is not actual.""" - raise NotImplementedError( - "logprob.read_effective_config: read the head dtype, the transform chain " - "and the vocab shard boundaries off the live engine." + "Read the head dtype, the transform chain and the vocab shard " + "boundaries off the live engine." ) def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: - """The collectives this operator really performed, this run. + """Return the collective trace: empty at TP=1, the partial-LSE reduction otherwise.""" - Empty at TP=1. Under vocab parallelism this is the partial-LSE reduction that - ``logp.reduce_topology`` and ``logp.merge_order`` are about. - """ - - raise NotImplementedError( - "logprob.observe_collectives: return the collective trace for this run " - "(empty tuple at TP=1)." - ) + raise NotImplementedError def resolve_implementation( factor_id: str, role: PolicyRole, impl_name: str ) -> tuple[Callable[..., Any] | None, ImplementationResolution]: - """Resolve the name an arm asks for into a callable, with the rejection trace. - - A parameter sweep never asks for one, so this stays unused until a logprob - factor gains a reference -- the deterministic vocab-parallel reduction, which - is not written yet. - """ + """Unused until a logprob factor gains a reference implementation.""" - raise NotImplementedError( - "logprob.resolve_implementation: import the requested implementation and " - "return (callable, ImplementationResolution) -- including the rejection " - "trace on failure." - ) + raise NotImplementedError diff --git a/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py b/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py index b39d8c42..2c2abfa1 100644 --- a/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py +++ b/rl_engine/mismatch/operator_checks/logprob/factors/precision_downcast.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""logp.precision_downcast -- lm_head dtype and where fp32 is written back.""" +"""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 @@ -29,23 +33,17 @@ "Is the deviation the lm_head GEMM running at the model dtype, or where " "the fp32 accumulator is written back?" ), - # A parameter sweep, not an implementation swap: nothing is replaced, the - # head dtype is scanned. ``reference=None`` is what says so -- there is no - # separate "kind" field. switch=Switch( path="logp.head_dtype", - # Cheapest tier: the dtype is a call argument, so arms of this factor - # reuse one engine and sort ahead of every rebuild-level case. 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={ - # The head GEMM's output dtype is identity, not a matter of taste. "precision.lm_head": ComparisonRule.MUST_MATCH_BITWISE, "precision.accumulate": ComparisonRule.MUST_MATCH_SEMANTICALLY, "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, - # Representation: both sides may name their mode differently. "extra.logprobs_mode": ComparisonRule.RECORD_ONLY, }, prerequisites=Prerequisites(required_ops=("lm_head",)), diff --git a/rl_engine/mismatch/pipeline/__init__.py b/rl_engine/mismatch/pipeline/__init__.py index b2fd373d..9881660f 100644 --- a/rl_engine/mismatch/pipeline/__init__.py +++ b/rl_engine/mismatch/pipeline/__init__.py @@ -3,20 +3,15 @@ """The execution pipeline: free functions only, no state. -One file per step, in the order they run: +One file per step, in the order they run:: -1. ``registry`` -- plugin registration, factor discovery, import-time conflicts -2--4. ``planner`` -- filter, statically reject, expand variants, order by cost -5. ``runner`` -- execution loop, reuse decisions, variant repeats -6. ``diagnosis`` -- four gates plus the matrix -7. ``report`` -- filter false positives, trace root causes, emit the report + 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`` is used by the runner and holds the declaration-driven contract -comparison. - -Dependency rules: ``schema/`` never imports ``pipeline/``; ``pipeline/`` never -imports ``operator_checks/``. Breaking the second rule turns "add an operator" -into "change the framework". +``comparison`` holds the declaration-driven contract comparison the runner uses. """ from rl_engine.mismatch.pipeline.comparison import compare_contracts, resolve_field_path diff --git a/rl_engine/mismatch/pipeline/comparison.py b/rl_engine/mismatch/pipeline/comparison.py index 20e663e7..6779fa39 100644 --- a/rl_engine/mismatch/pipeline/comparison.py +++ b/rl_engine/mismatch/pipeline/comparison.py @@ -3,8 +3,8 @@ """Contract comparison, driven entirely by declarations. -The comparison loop is generic; the per-operator part is putting fields in the -right place inside ``build_contract()``. There are no operator branches here. +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 @@ -31,8 +31,7 @@ def resolve_field_path(contract: OperatorContract, path: str) -> Any: """Index a contract by a dotted path such as ``collectives[0].reduction_order``. - Returns a sentinel when the path is absent so a missing field is reported - rather than raising. + Absent paths return a sentinel, so a missing field is reported not raised. """ current: Any = contract @@ -77,9 +76,8 @@ def compare_contracts( ) -> tuple[ComparisonIssue, ...]: """Compare the two contracts field by field, per ``comparison_rules``. - Returns every disagreement. ``RECORD_ONLY`` fields are never compared -- that - tier exists precisely so structural differences (packed QKV and the like) do - not drown the real problems in false positives. + ``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] = {} @@ -140,9 +138,8 @@ def _determinism_issues( ) -> tuple[ComparisonIssue, ...]: """One side claiming a stronger reproducibility guarantee than the other. - Comparing a topology-independent implementation against one that is not - reproducible even across runs measures the weaker side's noise, not the gap - between them. + Comparing against an implementation that is not reproducible across runs + measures the weaker side's noise, not the gap between the two. """ strength = { @@ -152,8 +149,7 @@ def _determinism_issues( DeterminismLevel.STABLE_ACROSS_TOPOLOGY: 3, } issues: list[ComparisonIssue] = [] - # strict=False: the two sides may legitimately declare different numbers of - # collectives; the overlap is what can be compared here. + # 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]: diff --git a/rl_engine/mismatch/pipeline/diagnosis.py b/rl_engine/mismatch/pipeline/diagnosis.py index b376906f..9650ae47 100644 --- a/rl_engine/mismatch/pipeline/diagnosis.py +++ b/rl_engine/mismatch/pipeline/diagnosis.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Four gates plus the diagnosis matrix. Step 6 of the pipeline. +"""Four gates plus the diagnosis matrix. The framework draws the conclusion; nobody reads the numbers by hand. """ @@ -45,10 +45,9 @@ def diagnose( ) -> FactorReport: """Run four gates, then the matrix. - **"Not measured" and "measured and clean" are two different things** -- - confusing them is the mistake an attribution framework is most likely to - make, so the gates come first and nothing reaches the matrix without passing - them. + "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 = ( @@ -71,9 +70,9 @@ def diagnose( def _gate_variants_applied(variants: Sequence[VariantResult]) -> _Outcome | None: """Gate 1: did every variant actually take effect? - ``FELL_BACK`` is the dangerous one -- the reference was requested, 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. + ``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: @@ -110,9 +109,8 @@ def _gate_evidence_complete( def _gate_shards_complete(variants: Sequence[VariantResult]) -> _Outcome | None: """Gate 3: were all logprob shards collected? - Under TP/CP each rank holds one slice. Missing a slice makes the combined - number wrong in a way that does not show: one vocab shard short and the LSE - denominator loses a chunk, so logp comes out systematically high. + 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: @@ -185,9 +183,8 @@ def _run_matrix( "both one-sided swaps converged; the reference is the only anchor", ) - # Neither converged. Before calling it clean, check the tail: the mean stays - # far below the clip edge at every production floor, so judging on the mean - # alone would mark almost everything NOT_THIS_FACTOR. + # 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( diff --git a/rl_engine/mismatch/pipeline/planner.py b/rl_engine/mismatch/pipeline/planner.py index 674af704..3cf26297 100644 --- a/rl_engine/mismatch/pipeline/planner.py +++ b/rl_engine/mismatch/pipeline/planner.py @@ -1,10 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Planning: filter, statically reject, expand variants, order by rebuild cost. - -Steps 2 to 4 of the pipeline. -""" +"""Planning: filter, statically reject, expand variants, order by rebuild cost.""" from __future__ import annotations @@ -48,11 +45,10 @@ class UnmetPrerequisite: def reject_contradictory_factors(factors: Sequence[MismatchFactor]) -> None: - """Static check, run between steps 2 and 3. Nothing executes. + """Reject a factor whose declaration contradicts itself. Nothing executes. - Claiming topology independence while using a non-deterministic reduction - order produces numbers that mean nothing. Reject at planning time rather - than explaining it away after a full run. + Claiming topology independence while reducing non-deterministically produces + numbers that mean nothing, and that is knowable before running. """ for factor in factors: @@ -73,10 +69,7 @@ def missing_prerequisites( gpu_count: int = 0, model_traits: frozenset[str] = frozenset(), ) -> tuple[UnmetPrerequisite, ...]: - """What this factor is still missing: operators, devices, packages, model traits. - - A whitelist probe, not a hand-maintained blacklist of "unsupported" strings. - """ + """What this factor is still missing: operators, devices, packages, traits.""" needs = factor.prerequisites unmet: list[UnmetPrerequisite] = [] @@ -110,9 +103,8 @@ def missing_prerequisites( def build_variants(factor: MismatchFactor) -> tuple[FactorVariant, ...]: """Expand one factor into its variants. - A single swap is not enough: **only a one-sided swap tells you which side is - at fault, and only a two-sided swap proves the reference itself is sound.** - Hence four arms rather than on/off. + 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: @@ -121,7 +113,7 @@ def build_variants(factor: MismatchFactor) -> tuple[FactorVariant, ...]: path = factor.switch.path reference = factor.reference - if reference is None: # no reference implementation means a parameter sweep + if reference is None: # a parameter sweep allowed = factor.switch.allowed_values or () return tuple( FactorVariant( @@ -187,18 +179,18 @@ def order_cases_by_rebind_cost( ) -> tuple[tuple[MismatchFactor, FactorVariant], ...]: """Order cases so rebuild cost never decreases, maximising reuse. - **Ordering is required, not a nicety**: run 160 cases in a random order and - the worst case restarts the process for every one of them. + 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 even show anything at the given floor. + """Whether this factor can show anything at the given floor. - A factor whose switch is process-level is identical on a single device, so - running it at the anchor floor wastes machine time. + 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: diff --git a/rl_engine/mismatch/pipeline/registry.py b/rl_engine/mismatch/pipeline/registry.py index 12ed2a7d..e9a39cbe 100644 --- a/rl_engine/mismatch/pipeline/registry.py +++ b/rl_engine/mismatch/pipeline/registry.py @@ -3,10 +3,9 @@ """Plugin registration and factor discovery. -Adding an operator is adding a directory; adding a factor is adding a file. No -existing file changes. If you find yourself having to edit the planner, a global -dict, or another operator's file, the abstraction is missing something -- raise -it and fix the framework rather than patching in place. +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 @@ -28,17 +27,15 @@ class OperatorChecks(Protocol): """Everything one operator needs checked. The plugin itself. - The four methods below are **operator-level**, not factor-level: how to read - configuration back from an engine and how to observe collectives is the same - logic for all of an operator's factors. Pushing them down would make every - factor file copy them. Factor files hold **declarations only, no behaviour**. + 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. The only thing that must be written - by hand.""" + """Which factors this operator has.""" ... def build_contract( @@ -48,12 +45,12 @@ def build_contract( ... def read_effective_config(self, role: PolicyRole, adapter: Any) -> Mapping[str, Any]: - """Read the switches' effective values back from the engine. A requested - value is not evidence.""" + """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. Feeds ``COLLECTIVE_CONTRACT``.""" + """Which collectives actually ran.""" ... def resolve_implementation( @@ -61,8 +58,7 @@ def resolve_implementation( ) -> tuple[Callable[..., Any] | None, ImplementationResolution]: """Resolve the name a variant asks for into a callable. - **Return the trace even when resolution fails** -- which candidates were - tried and why each was rejected. Returning a bare ``None`` leaves a + Return the trace even when resolution fails: a bare ``None`` leaves a silent fallback with nothing to investigate. """ ... @@ -79,7 +75,7 @@ def __init__(self) -> None: self._plugins: dict[str, OperatorChecks] = {} def register(self, plugin_cls: type) -> type: - """Decorator. Instantiates and registers, checking for conflicts.""" + """Instantiate and register a plugin, checking for conflicts.""" plugin = plugin_cls() name = getattr(plugin, "operator", "") @@ -93,8 +89,8 @@ def register(self, plugin_cls: type) -> type: return plugin_cls def _check_factor_conflicts(self, plugin: OperatorChecks) -> None: - """Reject duplicate factor ids, duplicate switch paths, and one contract - field claimed at two different comparison rules.""" + """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 = { @@ -152,9 +148,9 @@ class FactorDiscoveryError(ValueError): def discover_factors(package: str) -> tuple[MismatchFactor, ...]: """Collect the ``FACTOR`` constant from every module under ``.factors``. - The convention: **file name equals the factor id with the operator prefix - stripped**. It is enforced here so that renaming an id without renaming the - file fails at import rather than silently. + 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" diff --git a/rl_engine/mismatch/pipeline/report.py b/rl_engine/mismatch/pipeline/report.py index 89e4893b..e241adee 100644 --- a/rl_engine/mismatch/pipeline/report.py +++ b/rl_engine/mismatch/pipeline/report.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""False-positive filtering, root-cause tracing, and the final report. Step 7.""" +"""False-positive filtering, root-cause tracing, and the final report.""" from __future__ import annotations @@ -33,9 +33,8 @@ def filter_known_equivalences( ) -> tuple[tuple[str, ...], tuple[ModuleCorrespondence, ...]]: """Drop findings explained by a known structural equivalence. - Without this, a difference like fused QKV turns every weight comparison red - and buries the real problem. An equivalence may only be claimed with a test - that proves it -- one without ``verified_by`` is not trusted here. + An equivalence without ``verified_by`` is not trusted: unproven, "filtering + false positives" quietly becomes "hiding real findings". """ proven = tuple( @@ -53,12 +52,8 @@ def trace_root_causes( correspondences: Sequence[ModuleCorrespondence], edges: Sequence[PropagationEdge], ) -> tuple[RootCauseHypothesis, ...]: - """Walk from the still-aligned anchor down the call chain into ranked hypotheses. - - Thirty factors give thirty diagnoses, and that pile is not the answer. The - job here is to combine them with the module correspondence table and the - call chain into "the few most suspicious modules, ranked". - """ + """Walk from the still-aligned anchor down the call chain into ranked + hypotheses.""" downstream_of: dict[str, list[str]] = {} for edge in edges: @@ -150,8 +145,7 @@ def build_report( def render_summary(report: MismatchReport) -> str: - """A short human-readable summary. The hypotheses are what you want to read; - everything else is the evidence supporting them.""" + """A short human-readable summary, led by the ranked hypotheses.""" lines = [ f"noise floor: {report.noise_floor.value}", diff --git a/rl_engine/mismatch/pipeline/runner.py b/rl_engine/mismatch/pipeline/runner.py index 841c0c31..f66c1481 100644 --- a/rl_engine/mismatch/pipeline/runner.py +++ b/rl_engine/mismatch/pipeline/runner.py @@ -1,10 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""The execution loop: reuse decisions, variant repeats, metric computation. - -Step 5 of the pipeline. -""" +"""The execution loop: reuse decisions, variant repeats, metric computation.""" from __future__ import annotations @@ -33,8 +30,8 @@ class ScoringBackend(Protocol): """What the runner needs from an engine: logprobs for a fixed sequence. - Deliberately minimal so a CPU stub can satisfy it -- the framework's - plumbing is testable without a GPU. + Minimal enough that a CPU stub can satisfy it, so the plumbing is testable + without a GPU. """ def score( @@ -67,9 +64,8 @@ class RunContext: def assert_comparison_is_read_only(before: str, after: str) -> None: """Model tensor fingerprints must match before and after scoring. - If a kernel mutates weights in place, the first and second run differ, so - ``both_reference`` fails bitwise at random and gets blamed on the reference - -- attribution then points somewhere entirely wrong. This assertion is the + 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. """ @@ -89,9 +85,8 @@ def compute_metrics( ) -> MismatchMetrics: """Per-token metrics over active tokens only. - ``dlogp = log pi_theta - log pi_old``, so ``rho = exp(dlogp)``. The objective - clips ``rho`` at ``1 +/- eps``, which is ``|dlogp| > ln(1 + eps)`` -- past - that edge a token's gradient signal is discarded. + ``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] = [] @@ -123,7 +118,7 @@ def compute_metrics( 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, the standard PPO/GRPO diagnostic: rho - 1 - ln(rho). + # 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]) @@ -157,10 +152,7 @@ def _percentile(values: Sequence[float], q: float) -> float: 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". It never compares across - frameworks, so it is cheap; what it verifies is the **premise of the - self-check gate** -- ``both_reference`` can only anchor if the fixed-order - implementation really did fix the order. + The only exception to "one variant, one execution". """ if not variant.repeat_under: @@ -173,8 +165,7 @@ def expand_repeats(variant: FactorVariant) -> tuple[Mapping[str, Any], ...]: def assert_order_is_topology_independent(results: Sequence[Sequence[float]]) -> bool: """Runs of a topology-independent collective must agree bitwise. - Single-sided reruns, no cross-framework comparison -- the cheapest check in - the whole set. + Single-sided reruns, so the cheapest check in the whole set. """ if len(results) < 2: diff --git a/rl_engine/mismatch/reference_adapters/__init__.py b/rl_engine/mismatch/reference_adapters/__init__.py index f5be7273..f03fce17 100644 --- a/rl_engine/mismatch/reference_adapters/__init__.py +++ b/rl_engine/mismatch/reference_adapters/__init__.py @@ -1,26 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Wiring implementations in as reference implementations. **No operator code here.** +"""Wiring implementations in as reference implementations. No operator code here. -"Reference implementation" spans two directories. Of the three -``ReferenceAuthority`` levels, two are implemented in ``rl_engine/kernels/``: - -=========================== ===================== ============================== -authority implementation lives what this package does -=========================== ===================== ============================== -``FP64_ORACLE`` ``kernels/`` declare it, call it at the - lowest noise floor -``SHARED_BACKEND`` external library put it in deterministic mode, - read back to verify -``SELF_WRITTEN`` ``kernels/`` wrap it as a - ``ReferenceImplementation`` -=========================== ===================== ============================== - -Split of duties: ``kernels/`` owns "is the arithmetic right"; this package owns -"how do we put it in a deterministic mode, and how do we prove the setting -actually took effect". The latter is specific to this diagnostic framework and -has nothing to do with the operator itself. +``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 ( diff --git a/rl_engine/mismatch/reference_adapters/settings.py b/rl_engine/mismatch/reference_adapters/settings.py index 95a03dda..526027c9 100644 --- a/rl_engine/mismatch/reference_adapters/settings.py +++ b/rl_engine/mismatch/reference_adapters/settings.py @@ -3,10 +3,9 @@ """Delivering pinned settings by channel, and reading them back. -Neither TransformerEngine nor FlashInfer is deterministic by default, so the -settings have to be pinned explicitly. But pinning alone is not enough: **a -setting that cannot be read back can only be recorded as ``UNOBSERVABLE``** -- -delivered but unverifiable is the same as not delivered. +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 @@ -30,10 +29,8 @@ def apply_required_settings( ) -> dict[SettingChannel, dict[str, Any]]: """Route each setting to the right place for its channel. - Which channel a setting uses decides when it can take effect, and therefore - its rebind cost -- an env var needs a process restart, a call argument does - not. Flattening them into one dict leaves the framework unable to deliver - them at all. + 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 @@ -85,8 +82,7 @@ def verify_required_settings( ) -> tuple[SwitchStatus, tuple[str, ...]]: """Check the values that came back against what was pinned. - Returns ``UNOBSERVABLE`` when a setting declares no readback path: it was - delivered but cannot be proven, which is not evidence. + A setting with no readback path is ``UNOBSERVABLE``: delivered, unproven. """ unobservable: list[str] = [] diff --git a/rl_engine/mismatch/schema/__init__.py b/rl_engine/mismatch/schema/__init__.py index 1dd4c30d..b33c9777 100644 --- a/rl_engine/mismatch/schema/__init__.py +++ b/rl_engine/mismatch/schema/__init__.py @@ -3,31 +3,15 @@ """Pure data structures: public fields, frozen, no meaningful methods. -All behaviour lives in free functions under ``pipeline/``. This is a deliberate -trade-off, not an accident of using dataclasses: +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()``. -====================== ================== ================== - add a new function add a new type -====================== ================== ================== -data + procedural easy hard -objects + polymorphism hard easy -====================== ================== ================== +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. -This framework grows by **adding functions** (new diagnosis rules, new report -views, new ordering strategies, new evidence checks) while the set of types -stays stable -- so procedural is the correct side. - -Two consequences: - -* Do not hang methods on these structures. ``spec.requires_fixed_order()`` turns - a data structure into a half-object hybrid, which is the worst of both worlds. - Write ``requires_fixed_order(spec)`` instead. -* Chained field access such as ``factor.switch.path`` is **fine** here and does - not violate the Law of Demeter -- Demeter constrains an object's internals, - and plain data structures are supposed to expose their fields. Do not wrap - them in getters for the sake of it. - -Modules are ordered by dependency; each only imports from the ones above it. +Modules are ordered by dependency; each imports only from the ones above it. """ from rl_engine.mismatch.schema.collectives import ( diff --git a/rl_engine/mismatch/schema/collectives.py b/rl_engine/mismatch/schema/collectives.py index 4142c28e..44bbd0fa 100644 --- a/rl_engine/mismatch/schema/collectives.py +++ b/rl_engine/mismatch/schema/collectives.py @@ -4,11 +4,9 @@ """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. -``gemm.forward_reduce``, ``gemm.dgrad_reduce``, ``attn.cp_split_k_merge_order``, -``logp.reduce_topology``, ``moe.tp_forces_sp_reduce`` and -``gemm.rollout_all_reduce_backend`` are six instances of *one* semantic model -- -writing each separately guarantees they drift apart. +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 @@ -26,15 +24,13 @@ class CollectiveOp(str, Enum): ALL_TO_ALL = "all_to_all" BROADCAST = "broadcast" POINT_TO_POINT = "point_to_point" - NONE = "none" # single-device path, recorded explicitly rather than left blank + 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.distributed`` type, and this describes a parallel dimension, not the - process group object. + Not ``ProcessGroupKind``: ``ProcessGroup`` is an existing torch type. """ TENSOR = "tensor" @@ -48,8 +44,8 @@ class ParallelDim(str, Enum): class ReductionOrder(str, Enum): """Accumulation order -- the direct root of mismatch, not the collective.""" - ARRIVAL = "arrival" # first to arrive is first added. Control group only - NCCL_ALGORITHM = "nccl_algorithm" # ring/tree, varies with world size and message size + 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" @@ -68,14 +64,9 @@ class DeterminismLevel(str, Enum): class CollectiveContract: """The full numerical semantics of one collective. - Called a contract rather than a spec: it shares the word with - ``OperatorContract``, and it really is "what both sides must agree on", - not merely a description. - - A self-contradictory combination must be rejected at planning time: claiming - ``determinism == STABLE_ACROSS_TOPOLOGY`` while setting ``reduction_order`` - to ``NCCL_ALGORITHM`` or ``ARRIVAL`` produces numbers that mean nothing. - Static check, no run needed -- see ``reject_contradictory_factors()``. + 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 @@ -91,12 +82,10 @@ class CollectiveContract: @dataclass(frozen=True) class CollectiveRewrite: - """A mathematically identical rewrite of a collective -- equal in algebra, - unequal in floating point. + """A rewrite that is equal in algebra and unequal in floating point. - Megatron replacing ``all_reduce`` with ``reduce_scatter + all_gather`` when - sequence parallelism is on is exactly this rule applied. Whether the two - sides apply it is itself a source of mismatch, so it has to be declarable. + Whether each side applies one is itself a source of mismatch, so it has to be + declarable. """ name: str diff --git a/rl_engine/mismatch/schema/contracts.py b/rl_engine/mismatch/schema/contracts.py index 2f252264..4da47f83 100644 --- a/rl_engine/mismatch/schema/contracts.py +++ b/rl_engine/mismatch/schema/contracts.py @@ -24,12 +24,8 @@ class ComparisonRule(str, Enum): class ComparisonIssueCode(str, Enum): """Stable reason codes for the two sides disagreeing. - **These strings go into the artifact schema and callers branch on them -- - renaming one requires a schema version bump.** - - Deliberately excludes "the declaration contradicts itself": that is rejected - by the planner before anything runs (``reject_contradictory_factors()``) and - has nothing to do with comparing the two sides. + 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" @@ -40,8 +36,7 @@ class ComparisonIssueCode(str, Enum): @dataclass(frozen=True) class ComparisonIssue: - """One record of the two sides disagreeing -- an element of what - ``compare_contracts()`` returns.""" + """One record of the two sides disagreeing.""" code: ComparisonIssueCode rule: ComparisonRule # which tier this field was declared at @@ -54,22 +49,16 @@ class ComparisonIssue: class OperatorContract: """One operator's numerical contract on one side. - Only three things are common to every operator and live here; the rest (RoPE - state, vocab shard map, GQA head mapping, ...) belongs to each plugin and - goes in ``extra``. - - Field path convention: the keys of ``comparison_rules`` are paths counted - from the contract root -- + 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 and compares entry by entry. So - a plugin only has to put fields in the right place inside - ``build_contract()``; it never needs to write a function that "collects the - comparable fields". Keep ``extra`` flat -- nesting dicts makes paths long - and unreadable. + 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 diff --git a/rl_engine/mismatch/schema/factors.py b/rl_engine/mismatch/schema/factors.py index 961d6002..beca3fd3 100644 --- a/rl_engine/mismatch/schema/factors.py +++ b/rl_engine/mismatch/schema/factors.py @@ -23,34 +23,28 @@ class FactorCategory(str, Enum): - """Which family a factor belongs to -- decides where to look when it fires. - - Deliberately not ``Layer``: in this domain a layer is a model layer. - """ + """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 choice, fusion, inner precision + 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. + """Evidence every factor must have before a verdict is allowed. - General items only. Operator-specific evidence does not belong in this enum: - putting it here would turn "add an operator" into "change the framework". - Plugins declare their own as plain strings (see the constants below). + 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" # all variants share one set of weights + MODEL_STATE_FINGERPRINT = "model_state_fingerprint" LIBRARY_VERSIONS = "library_versions" -# Operator-specific evidence: plugins define constants in their own module and -# the framework treats them as plain strings. Adding an operator adds a new set -# of constants; the framework does not change. COLLECTIVE_CONTRACT = "collective_contract" BATCH_PLACEMENT = "batch_placement" MODEL_SHAPE = "model_shape" @@ -62,21 +56,18 @@ class Evidence(str, Enum): class ReferenceAuthority(str, Enum): """Where a reference implementation comes from, most authoritative first. - Not ``ReferenceTier``: "tier" says there are levels without saying what - orders them. What orders these is authority. - - This is a decision order, not a description: look for a SHARED_BACKEND - first, and only write SELF_WRITTEN when the first two cannot cover it. + 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, mathematically exact, lowest noise floor only + 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: - """The one definition of a switch. Allowed values and parser declared once.""" + """A switch's one definition: allowed values and parser declared together.""" path: str # "gemm.forward_reduce" rebind_cost: RebindCost @@ -91,10 +82,7 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class Prerequisites: - """What this factor needs in order to run. A whitelist, not a blacklist. - - Capability probing itself is delegated: this only declares what is needed. - """ + """What a factor needs in order to run. A whitelist, not a blacklist.""" required_ops: tuple[str, ...] = () min_gpu_count: int = 1 @@ -105,31 +93,15 @@ class Prerequisites: @dataclass(frozen=True) class ReferenceImplementation: - """What replaces the native implementation, and which execution paths it covers. - - **``covers_paths`` defines the shape of the self-check gate**, the field most - easily overlooked: the gate is not "both sides use the same implementation", - it is "**every path this reference covers must agree bitwise on the same - sequence**". - - Most factors cover two paths (training full-prefill plus rollout - full-prefill) and the gate holds across the two sides. But an attention - FlashInfer reference also covers ROLLOUT_DECODE -- so the gate holds *inside - the rollout side*: with full-prefill and decode both switched to FlashInfer - and ``num_splits=1`` pinned, ``(out, lse)`` must agree on the same sequence. - A disagreement is the reference's fault, unrelated to the training side. - **No decode stub is needed on the training side.** - - Covering several paths also gives the attribution split directly: - - dlogp(training, real rollout) - = dlogp(TRAINING_FULL_PREFILL, ROLLOUT_FULL_PREFILL) cross-framework - + dlogp(ROLLOUT_FULL_PREFILL, ROLLOUT_DECODE) rollout-side path - - The two terms are fixed in completely different ways: the first by swapping - the kernel backend, the second by changing decode's softmax block size to - match prefill. Without the middle path, attribution stops at the combined - "cross-framework plus cross-path" answer. + """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 @@ -139,7 +111,7 @@ class ReferenceImplementation: covers_paths: tuple[ExecutionPath, ...] fp64_oracle: str | None = None required_settings: tuple[RequiredSetting, ...] = () - pinned_libraries: tuple[LibraryPin, ...] = () # **required**, see LibraryPin + pinned_libraries: tuple[LibraryPin, ...] = () @dataclass(frozen=True) @@ -148,9 +120,9 @@ class MismatchFactor: Not ``DivergenceFactor``: in an RL context, divergence means KL divergence. - A factor with ``reference is None`` is a parameter sweep; with a reference it - is an implementation swap. There is no separate ``kind`` field -- derivable - state is state that can disagree with itself. + ``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 @@ -168,28 +140,20 @@ class MismatchFactor: def declared_collectives(factor: MismatchFactor) -> tuple[CollectiveContract, ...]: - """Collectives a factor's reference implementation pins, if any. - - Used by the planner's static check; returns empty when the factor does not - touch collective communication. - """ + """Collectives a factor's reference pins, for the planner's static check.""" reference = factor.reference if reference is None: return () - pinned = tuple( + return tuple( setting.value for setting in reference.required_settings if isinstance(setting.value, CollectiveContract) ) - return pinned def requires_fixed_order(contract: CollectiveContract) -> bool: - """Whether this contract claims its result is independent of topology. - - Claiming it means it must survive the assertion in ``pipeline/planner.py``. - """ + """Whether this contract claims its result is independent of topology.""" return contract.determinism is DeterminismLevel.STABLE_ACROSS_TOPOLOGY diff --git a/rl_engine/mismatch/schema/fingerprints.py b/rl_engine/mismatch/schema/fingerprints.py index 33bbca1b..442de7e1 100644 --- a/rl_engine/mismatch/schema/fingerprints.py +++ b/rl_engine/mismatch/schema/fingerprints.py @@ -20,11 +20,9 @@ @dataclass(frozen=True) class ReuseKey: - """Whether an already-built runtime can be reused. Four parts matching the - four ``RebindCost`` levels, coarse to fine. + """Whether an already-built runtime can be reused. - Compared coarse to fine: a different ``process`` means restarting the whole - process; equal process but different engine means only rebuilding the engine. + Four parts matching the four ``RebindCost`` levels, compared coarse to fine. """ process: str # env vars, determinism switches, compile-time flags @@ -35,8 +33,7 @@ class ReuseKey: @dataclass(frozen=True) class EnvironmentFingerprint: - """The execution environment. When this layer changes, every numerical - conclusion has to be re-run.""" + """The execution environment. Change this layer and every number is stale.""" python_version: str torch_version: str @@ -53,15 +50,11 @@ class ExecutionFingerprint: """One execution's full identity. Any part differing makes two runs incomparable. - Three rules: - - 1. **What goes in is the value read back, not the value requested.** - Requesting ``num_splits=1`` and the backend actually using 1 are two - different facts. - 2. **Library versions belong here, not in a footnote.** - 3. **Thresholds go in too.** Change a threshold and every historical - pass/fail must go stale -- which is why thresholds are code constants: a - configurable value cannot be pinned into an identity. + 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 @@ -75,15 +68,10 @@ class ExecutionFingerprint: @dataclass(frozen=True) class VariantRecord: - """One variant's immutable on-disk record. Append-only, never edited. - - Difference from ``VariantResult``: a Result is what was just computed in - memory, a Record is archived -- it additionally carries the execution - identity and a content hash, the latter being its integrity seal. Only a - record whose hash verifies may be reused on resume. + """One variant's archived record: a ``VariantResult`` plus its identity. - Not ``VariantArtifact`` -- in an ML context, "artifact" usually means model - weights or checkpoints (MLflow artifacts); this stores an execution record. + ``content_hash`` is its integrity seal. Only a record whose hash verifies may + be reused on resume. """ variant_name: str @@ -93,11 +81,7 @@ class VariantRecord: def canonical_fingerprint(payload: Any) -> str: - """Stable hash of a JSON-serialisable payload. - - Key order is normalised so the same content always hashes the same, whatever - order the dict was built in. - """ + """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() diff --git a/rl_engine/mismatch/schema/metrics.py b/rl_engine/mismatch/schema/metrics.py index 4113246f..e74606a2 100644 --- a/rl_engine/mismatch/schema/metrics.py +++ b/rl_engine/mismatch/schema/metrics.py @@ -3,23 +3,14 @@ """Metrics and per-variant results. -``dlogp`` is an intermediate quantity. GRPO's gradient enters the objective -through the importance sampling ratio:: - - rho_{i,t}(theta) = pi_theta(y | x) / pi_old(y | x) = exp(dlogp_{i,t}) - -The objective clips rho, typically at eps = 0.2, i.e. rho in [0.8, 1.2]. Back in -dlogp terms, anything past ``ln(1 + eps) ~= 0.182`` **gets clipped**. - -That is the actual mechanism by which mismatch breaks training: it raises the -clip fraction, and clipped tokens have their gradient signal cut -- **RL discards -samples it should have learned from**, and not random ones, the ones with the -largest mismatch. - -Against the threshold table: dense sits at ``dlogp_mean ~= 0.002-0.008`` and -large MoE at ``0.01-0.03`` -- all far below 0.182. **So the mean never triggers -clipping, and looking only at the mean necessarily concludes "everything is -fine".** The danger is entirely in the tail. +``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 @@ -42,8 +33,7 @@ @dataclass(frozen=True) class WorstToken: - """The largest-deviation token -- the entry point for attribution, since it - often points straight at a layer or an expert.""" + """The largest-deviation token, which often points straight at a layer.""" position: int token_id: int @@ -77,13 +67,10 @@ class RejectedCandidate: @dataclass(frozen=True) class ImplementationResolution: - """Which candidates were tried and why each was rejected -- what you look at - when the status is ``FELL_BACK``. + """Which candidates were tried and why each was rejected. - A single ``fallback_reason: str`` is not enough: what you actually want is - "was the first choice rejected for a missing library, missing devices, or a - version mismatch", plus "which candidate did it land on". Without this trace, - a silent fallback leaves nothing to go on. + 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 @@ -95,14 +82,9 @@ class ImplementationResolution: class LogprobShard: """The slice of logprobs one rank holds under TP/CP. - With TP/CP on, logprobs are not computed by any single rank -- each holds one - slice. **Missing any slice makes the combined result wrong in a way that does - not show**: drop one vocab shard and the LSE denominator loses a chunk, so - logp is systematically too high. Slices must be collected per rank and - checked against ``world_size`` before merging. - - Not ``RankObservation`` -- in RL, an observation is what the agent sees of - the environment. + 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 @@ -149,11 +131,7 @@ def is_silent_failure(metrics: MismatchMetrics) -> bool: def missing_evidence(collected: frozenset[str], required: tuple[str, ...]) -> frozenset[str]: - """Which required evidence is absent. - - A free function rather than a method on a bundle type: making a data - structure reach into a factor's fields would be feature envy. - """ + """Which required evidence is absent.""" return frozenset(required) - collected diff --git a/rl_engine/mismatch/schema/pitfalls.py b/rl_engine/mismatch/schema/pitfalls.py index c6549769..7341f099 100644 --- a/rl_engine/mismatch/schema/pitfalls.py +++ b/rl_engine/mismatch/schema/pitfalls.py @@ -3,9 +3,8 @@ """Known pitfalls, encoded as data. -Every factor in the docs carries a "pitfall" note. Prose pitfalls get read once -and never again -- they have to be data, so the framework can block them before -a run instead of relying on somebody remembering afterwards. +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 @@ -17,12 +16,7 @@ class FailureMode(str, Enum): - """How a pitfall fails -- decides which tool the framework needs against it. - - Borrowed from reliability engineering's failure mode, which is more accurate - than a bare "kind": the useful part is *how it fools you*, not which bucket - it sits in. - """ + """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 @@ -36,19 +30,15 @@ class FailureMode(str, Enum): class KnownPitfall: """A known pitfall together with the assertion that blocks it. - ``symptom`` and ``actual_cause`` are separate on purpose: a pitfall is a - pitfall precisely because its appearance points at the wrong cause. - - Originally split into ``Pitfall`` and ``Precheck``, but the two referenced - each other's ids -- a two-way reference means they were one thing all along: - a pitfall, and the assertion that stops 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 # what it looks like - actual_cause: str # what it actually is - guard: str # the assertion that blocks it, one line + symptom: str + actual_cause: str + guard: str guard_runs_at: NoiseFloor # lowest floor that can run it -- cheap checks first diff --git a/rl_engine/mismatch/schema/rollout_context.py b/rl_engine/mismatch/schema/rollout_context.py index 2ff29a08..c8368a8a 100644 --- a/rl_engine/mismatch/schema/rollout_context.py +++ b/rl_engine/mismatch/schema/rollout_context.py @@ -1,28 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Online RL context. - -The factor model targets **offline pairwise comparison**; real training is an -**online RL loop**. Offline these are placeholders and a factor author never -touches them. They matter only when this framework is wired into a real training -loop for online auditing -- and there, missing them causes two kinds of problem: - -**One, results stop reproducing.** You get ``dlogp = 0.005`` today and ``0.008`` -tomorrow from the same config, and eventually find the batch order changed and -one sample landed in a different microbatch -- **the group it all-reduces with -changed, so the accumulation order changed**. That is what ``BatchPlacement`` -records: not "environment" but part of the **identity**. Changing where a sample -lands is changing the computation. ``DynamicSamplingDecision`` is the same story: -GRPO drops all-correct or all-wrong groups, and which ones were dropped changes -the whole batch composition. - -**Two, metrics get misread.** GRPO normalises advantage within a group: -``A_i = (r_i - mean(r)) / (std(r) + eps)``. So **one rollout's logp drifting -affects all K rollouts in that group through the advantage** -- the error is not -independent per rollout, it amplifies within the group. Averaging over tokens -erases this completely: ``dlogp_mean`` looks fine while a few groups are skewed -wholesale. Hence metrics must be aggregatable by ``RolloutGroup``. +"""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 @@ -34,13 +17,10 @@ class RolloutGroup: """One GRPO group: K rollouts from the same prompt. - Two modes, depending on who is running: - - * **offline ablation** (the main use): identity is built from a fixed token - fixture and this is a placeholder (1 prompt x 1 rollout); group semantics - are unused; - * **online, wired into real training**: a real group comes from the training - loop, and only then does aggregating by group mean anything. + 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 @@ -52,8 +32,9 @@ class RolloutGroup: class BatchPlacement: """Where one sample landed in this training step. - Not ``BatchLayout`` -- in tensor land, layout means memory ordering (NCHW and - friends); this is about placement. + 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 @@ -64,12 +45,7 @@ class BatchPlacement: @dataclass(frozen=True) class DynamicSamplingDecision: - """Record of dropping a group (all-correct or all-wrong groups carry no - learning signal). - - Dropping changes the batch composition and therefore the reduction grouping - -- unrecorded, it cannot be reproduced. - """ + """Record of dropping a group, which changes the batch composition.""" kept: bool reason: str | None = None @@ -77,10 +53,9 @@ class DynamicSamplingDecision: @dataclass(frozen=True) class ComparisonIdentity: - """This comparison's input identity. If these differ, no numerical - comparison means anything. + """This comparison's input identity. - Not ``ScoringIdentity`` -- in RL, "score" means reward scoring. + If these differ, no numerical comparison means anything. """ prompt_token_ids: tuple[int, ...] @@ -89,7 +64,7 @@ class ComparisonIdentity: position_ids: tuple[int, ...] checkpoint_id: str checkpoint_revision: str - model_shape: str # "L=1,H=896,Hq=14,Hkv=2,D=64" -- a trimmed model is a different model + model_shape: str # a trimmed model is a different model group: RolloutGroup batch_placement: BatchPlacement sampling_decision: DynamicSamplingDecision diff --git a/rl_engine/mismatch/schema/thresholds.py b/rl_engine/mismatch/schema/thresholds.py index ca235051..a8f9a8fa 100644 --- a/rl_engine/mismatch/schema/thresholds.py +++ b/rl_engine/mismatch/schema/thresholds.py @@ -68,10 +68,9 @@ def expected_range( ) -> ExpectedRange: """Look up the normal band for one combination. - Reading a low floor with the production table hides real bugs: at - SINGLE_LAYER_ANCHOR the expectation is bitwise, and judging it against - 0.002-0.008 would call a definite operator error "normal". So the lookup is - always keyed by noise floor, and an exact model family beats the wildcard. + 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 = [ @@ -99,10 +98,7 @@ def expected_range( def tolerance_floor(model_family: str, noise_floor: NoiseFloor) -> float: - """The absolute floor below which a difference is not treated as a signal. - - Used by the diagnosis matrix as ``tol_floor`` when deciding convergence. - """ + """The floor below which a difference is not treated as a signal.""" return expected_range(model_family, noise_floor).suspect_above diff --git a/rl_engine/mismatch/schema/tracing.py b/rl_engine/mismatch/schema/tracing.py index 641a0ee6..e9eccb05 100644 --- a/rl_engine/mismatch/schema/tracing.py +++ b/rl_engine/mismatch/schema/tracing.py @@ -16,21 +16,11 @@ @dataclass(frozen=True) class ModuleCorrespondence: - """A training-side module paired with its rollout-side counterpart, plus any - known structural difference between them. + """A training-side module paired with its rollout-side counterpart. - Without this table the two sides' tensors cannot even be matched up, let - alone compared entry by entry. - - A non-empty ``equivalence`` means "different in form, equal in arithmetic" -- - the basis for filtering false positives. Without it, a difference like fused - QKV turns every weight comparison red and buries the real problem. - - Filled once per model by whoever brings the model up, 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. Swap in GLM5 or - DSv3.2 and it has to be redone -- so it lives in ``model_meta/``, not in any - ``operator_checks/``. + 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 @@ -42,20 +32,13 @@ class ModuleCorrespondence: @dataclass(frozen=True) class PropagationEdge: - """One directed edge of the call chain: mismatch at ``upstream`` propagates - to ``downstream``. - - Tracing walks these backwards: start from the last still-aligned module and - look downstream for the first mismatched one. - """ + """One directed edge of the call chain, followed backwards when tracing.""" upstream: str # semantic_name downstream: str class RootCauseCategory(str, Enum): - """How a root cause is classified.""" - MISSING_OPERATOR = "missing_operator" # one side does not have it at all DIFFERENT_IMPLEMENTATION = "different_implementation" DIFFERENT_PARAMETER = "different_parameter" @@ -64,11 +47,10 @@ class RootCauseCategory(str, Enum): @dataclass(frozen=True) class RootCauseHypothesis: - """One root-cause hypothesis, produced by walking the call chain after a - ``NOT_THIS_FACTOR``. + """One hypothesis from walking the call chain after a ``NOT_THIS_FACTOR``. - ``anchor_module`` is the last position where the two sides still agree -- - the root cause must be downstream of it. + The root cause must be downstream of ``anchor_module``, the last position + where the two sides still agree. """ suspected_module: str @@ -81,34 +63,19 @@ class RootCauseHypothesis: @dataclass(frozen=True) class MismatchReport: - """The final report for one run. **Three levels away from a Diagnosis**:: - - one factor's variants -> VariantResult x N - | diagnose() - FactorReport (holds one Diagnosis enum value) - | x 30 factors - trace_root_causes() - v - MismatchReport (ranked hypotheses) - - * ``Diagnosis`` is **one factor's** conclusion, one of eight values; - * ``FactorReport`` is that factor's full record: every variant plus that one - diagnosis plus the reason; - * ``MismatchReport`` is the **whole run**, answering what a Diagnosis cannot. - - Thirty factors give thirty diagnoses -- say three ``CAUSED_BY_TRAINING_SIDE``, - twenty-five ``NOT_THIS_FACTOR``, two ``INSUFFICIENT_EVIDENCE``. That pile is - not the answer. **The report's job is to combine them with the module - correspondence table and the call chain into "the few most suspicious - modules, ranked"** -- the ``hypotheses`` field. The rest is the evidence - supporting it. + """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, ...] # those with equivalence set + filtered_false_positives: tuple[ModuleCorrespondence, ...] failed_guards: tuple[KnownPitfall, ...] diff --git a/rl_engine/mismatch/schema/values.py b/rl_engine/mismatch/schema/values.py index 5382f15b..9fc9db43 100644 --- a/rl_engine/mismatch/schema/values.py +++ b/rl_engine/mismatch/schema/values.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Base value types. This module imports nothing from the project. +"""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/``. @@ -17,22 +17,18 @@ class PolicyRole(str, Enum): """Which of the two compared policies this side plays. - Not ``Side``: that only says "one of the sides" without saying of what. - Not ``Policy`` either -- in RL that means the policy itself (network and - distribution), not the part it plays in this comparison. + Not ``Policy``: in RL that means the policy itself, not the part it plays. """ - ROLLOUT = "rollout" # pi_old -- the policy that produced the trajectory - TRAINING = "training" # pi_theta -- the policy being updated + 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. - The same mathematical quantity (a token's conditional probability) can come - out of several different physical paths: feed the whole sequence at once, - feed it in chunks, or generate token by token. Mathematically equivalent, - not equal in floating point. + Whole sequence at once, in chunks, or token by token: mathematically + equivalent, not equal in floating point. """ TRAINING_FULL_PREFILL = "training_full_prefill" @@ -53,23 +49,17 @@ class Precision(str, Enum): class DowncastPoint(str, Enum): """When a high-precision accumulator is written back at lower precision. - ``downcast`` here means numerical precision reduction (fp32 -> bf16), not - the C++ sense of casting down a class hierarchy. + Numerical precision reduction, not the C++ sense of casting down a hierarchy. """ NEVER = "never" - PER_BLOCK = "per_block" # right after each block (largest error) - PER_PARTIAL = "per_partial" # each shard's partial sum - FINAL_WRITE = "final_write" # only on the final write (smallest error) + 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 this switch. - - Named after the cost rather than an isolation level because its only job is - to answer "is this one expensive?", which drives case ordering and therefore - the wall-clock of the whole run. - """ + """What it costs to change a switch. Drives case ordering.""" PER_REQUEST = "per_request" ENGINE_REBUILD = "engine_rebuild" @@ -78,12 +68,7 @@ class RebindCost(str, Enum): class SettingChannel(str, Enum): - """How a setting reaches the engine, which decides when it takes effect. - - These three are delivered in completely different ways. Flattening them into - one dict leaves the framework unable to deliver them, and unable to derive - the rebind cost that case ordering depends on. - """ + """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 @@ -95,10 +80,9 @@ class SettingChannel(str, Enum): class LibraryPin: """A pinned library version. - TransformerEngine and FlashInfer change kernel selection across versions -- - the same factor can reach the *opposite* conclusion on a different version. - So the version is part of the experiment's identity, not a footnote: it goes - into the fingerprint, and changing it invalidates historical results. + 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 @@ -109,12 +93,10 @@ class LibraryPin: @dataclass(frozen=True) class RequiredSetting: - """A setting that must be pinned: what to pin, how to deliver it, how to - read it back, and which pitfall it guards against. + """A setting that must be pinned, and how to prove it was. - ``readback`` is the main reason this type exists: **a requested value is not - evidence**. A setting that cannot be read back can only be recorded as - ``UNOBSERVABLE`` -- delivered but unverifiable is the same as not delivered. + A setting that cannot be read back can only be recorded ``UNOBSERVABLE``: + delivered but unverifiable is the same as not delivered. """ key: str @@ -128,13 +110,7 @@ class RequiredSetting: class PrecisionProfile: """Which precision this side uses at each point in the computation. - Not ``PrecisionPolicy`` -- in RL, ``Policy`` is pi, the one word this domain - must not borrow. - - Reserved for the precision factors that get wired in later. Today - ``rollout.dtype`` and ``training.compute_dtype`` are two independent - switches with inconsistent names and no way to express "both sides should - match"; this type converges them. + Not ``PrecisionPolicy``: in RL, ``Policy`` is pi. """ compute: Precision @@ -147,10 +123,7 @@ class PrecisionProfile: def choice_parser(*allowed: Any) -> Callable[[Any], Any]: - """Build a parser that accepts only the given values. - - Used by ``Switch`` so the allowed set and the parser are declared once. - """ + """Build a parser accepting only the given values.""" permitted = tuple(allowed) diff --git a/rl_engine/mismatch/schema/variants.py b/rl_engine/mismatch/schema/variants.py index 54216eb4..2588be6b 100644 --- a/rl_engine/mismatch/schema/variants.py +++ b/rl_engine/mismatch/schema/variants.py @@ -15,7 +15,7 @@ class VariantExpansion(str, Enum): """Which variants a factor expands into.""" - STANDARD_FOUR = "standard_four" # swap factors: both_native / both_reference / two one-sided + 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 @@ -26,12 +26,9 @@ class ExpectedOutcome(str, Enum): class SwitchStatus(str, Enum): - """Whether the switch actually reached the engine. A silent fallback is far - more harmful than an error. + """Whether the switch actually reached the engine. - Not ``MaterializationStatus`` -- it describes a **switch**'s state (did it - get delivered), which "materialization" neither reveals nor keeps separate - from attention's execution path. + A silent fallback is far more harmful than an error. """ APPLIED = "applied" @@ -42,13 +39,11 @@ class SwitchStatus(str, Enum): class Diagnosis(str, Enum): - """The conclusion for **one factor** after its variants have run -- an enum - value, not a report. + """One factor's conclusion after its variants have run. - The first three are "cannot judge" and must stay strictly separate from - "judged, nothing here". It answers "is this one factor the culprit"; it - cannot answer "what causes the mismatch" -- that needs cross-factor - synthesis, which is what ``MismatchReport`` is for. + 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" @@ -62,47 +57,41 @@ class Diagnosis(str, Enum): class NoiseFloor(str, Enum): - """The experiment's noise floor: how small a difference this run can resolve. - Orthogonal to factors. A floor that has not passed blocks the next one. + """How small a difference this run can resolve. Orthogonal to factors. - From signal processing: a signal below the noise floor cannot be measured. - Same here -- at the production floor, an e-6 difference drowns in parallel - reduction noise. Deliberately not named after scale: the levels are graded - by noise, not by size. - - The four are arranged so that **each step down introduces exactly one new - noise source**. When a floor starts failing, the suspect set is decided: it - is whatever that floor just added. + 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 fully on / batch=1 / one token. - # Noise sources: none, only the operator itself. - # Called an anchor because it is the reference point for the other three: - # **failing bitwise here is an operator bug, not "training-inference - # mismatch"**, and the other three floors need not run at all. + # 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 / determinism still on. - # New: accumulation over depth. Tests whether error grows linearly or - # exponentially with layer count. + # 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" - # All layers / TP + SP on one node / determinism still on. - # New: reduction-order differences from sharding. - # **This is the first floor with "real" training-inference mismatch.** + # 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 / target batch / 2k-8k, decode path. - # New: everything else (cross-node comms, non-deterministic kernels, the - # real generation path). The floor reports are written from, and the only - # one whose numbers may be read against the threshold table. + # 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 -- the machine-readable form of - "how do I configure this ablation".""" + """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] @@ -110,27 +99,15 @@ class FactorVariant: expected: ExpectedOutcome = ExpectedOutcome.MEASURE_ONLY why: str = "" repeat_under: Mapping[str, tuple[Any, ...]] | None = None - # Run this same variant once under each environment and require bitwise - # equality -- **the only exception to "one variant, one execution"**. The - # runner expands the cartesian product automatically. - # - # repeat_under = {"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")} - # -> four runs, asserted bitwise identical. - # - # It never compares across frameworks, so it is very cheap; and what it - # verifies is the **premise of the self-check gate**: both_reference can - # only serve as an anchor if the "fixed-order implementation" really did fix - # the order. Without this, REFERENCE_ITSELF_IS_BROKEN is itself unreliable. @dataclass(frozen=True) class ExpectedRange: """The normal band for a metric under one (model family, noise floor, config). - **Must be a code constant, never a config file** -- a tunable threshold means - somebody can tune it until the test passes. It enters the execution - fingerprint: changing a threshold requires a code change and review, and - invalidates every historical pass/fail. + 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" / "*" diff --git a/tests/mismatch_cpu_backend.py b/tests/mismatch_cpu_backend.py index 98fb7e69..0bef380f 100644 --- a/tests/mismatch_cpu_backend.py +++ b/tests/mismatch_cpu_backend.py @@ -1,18 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""A CPU scoring backend used to exercise the framework without a GPU. +"""A CPU scoring backend for exercising the framework without a GPU. -This validates plumbing only. It does not claim anything about real Megatron or -vLLM numerics -- those live in ``megatron.py`` and ``vllm.py`` and need devices. - -It is deliberately able to *simulate* the failure modes the framework is built to -catch, so the gates and the matrix can be tested for real: - -* a configurable per-side bias, so one-sided attribution can be checked; -* a switch that silently does nothing, so ``FELL_BACK`` is reachable; -* an unstable mode whose output changes with the environment, so the - topology-independence assertion can actually fail. +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 @@ -33,8 +28,7 @@ class CpuScoringBackend: """Deterministic synthetic logprobs, with injectable deviation. - ``bias`` is the deviation this side adds on top of the shared base signal. - Setting it on one side only is how a one-sided root cause is simulated. + Setting ``bias`` on one side only simulates a one-sided root cause. """ role: PolicyRole @@ -59,16 +53,14 @@ def score( ) -> tuple[Sequence[float], Mapping[str, Any]]: """Produce logprobs for the fixed sequence. - The base signal only depends on the identity, so both sides agree unless - a bias is injected -- which is what makes an injected deviation the only - thing the metrics can see. + 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) - # A replacement callable means the reference implementation is in play; - # the reference is defined to have no bias of its own. + # 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 @@ -86,10 +78,7 @@ def score( return logprobs, effective def reuse_key(self, switch_values: Mapping[str, Any]) -> ReuseKey: - """Group switches by how expensive they are to change. - - Mirrors the four ``RebindCost`` levels so case ordering can be tested. - """ + """Group switches by tier, mirroring the four ``RebindCost`` levels.""" def digest(prefix: str) -> str: payload = sorted( From 443c24423ff118fda40d02e4cc08b2ddd9be11b9 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Sun, 9 Aug 2026 12:24:24 +0000 Subject: [PATCH 4/7] docs(mismatch): remove duplication, keep one home per fact The README, both tutorials and the PR description each restated the four arms, the gates and the noise floors. Duplicated prose goes stale in the copies nobody edits, so each fact now lives in exactly one place: README holds the concepts, the tutorials hold the steps, and the comm tutorial covers only what differs from the kernel one. Also dropped the code that the repository already shows. A tutorial that pastes an entire factor declaration is a second copy to keep in sync; pointing at operator_checks/attention/factors/rope_fusion.py is not. 1176 lines to 629. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01QjJLV6aQZTtm4X6pjex8Ri --- rl_engine/mismatch/README.md | 241 ++++------ rl_engine/mismatch/docs/README.md | 31 +- rl_engine/mismatch/docs/add-a-comm-feature.md | 325 ++++---------- .../mismatch/docs/add-a-kernel-factor.md | 425 +++++------------- 4 files changed, 277 insertions(+), 745 deletions(-) diff --git a/rl_engine/mismatch/README.md b/rl_engine/mismatch/README.md index 3edcf698..d2fea2bf 100644 --- a/rl_engine/mismatch/README.md +++ b/rl_engine/mismatch/README.md @@ -5,192 +5,115 @@ 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 a set of named **factors**, -each of which is a switch that can be flipped one at a time, run, and attributed -to a side. - -The output is not a number. It is a **report**: which factors were measured, -which could not be measured and why, and the few most suspicious modules ranked -by suspicion. - -## Why the mismatch matters - -`dlogp = log π_θ − log π_old` enters the GRPO/PPO objective through the -importance ratio `ρ = exp(dlogp)`, which the objective clips at `1 ± ε`. With the -usual `ε = 0.2`, any token past `|dlogp| > ln(1.2) ≈ 0.182` has its **gradient -signal discarded** — and the discarded tokens are not random, they are the ones -with the largest mismatch. - -The trap: at the production floor a healthy `dlogp_mean` is `0.002–0.008` (dense) -or `0.01–0.03` (large MoE), all far below the clip edge. **Judging on the mean -alone always concludes "everything is fine."** The danger lives in the tail, which -is why `MismatchMetrics` carries `dlogp_p99` / `dlogp_max` / `clip_fraction` / -`approx_kl` / `worst_token`, and why the diagnosis matrix converges on -`clip_fraction`, not on the mean (`pipeline/diagnosis.py::_converged`). - -## What it does, concretely - -1. **Collects factors** from operator plugins at import time, rejecting duplicate - ids, duplicate switch paths, and one contract field claimed at two different - comparison rules. -2. **Rejects self-contradictory declarations statically** — claiming topology - independence while reducing in NCCL's chosen order produces numbers that mean - nothing, and that is knowable before anything runs. -3. **Expands each factor into four arms**, not on/off: - - | arm | rollout | training | what it buys | - |---|---|---|---| - | `both_native` | native | native | the baseline every other arm is measured against | - | `both_reference` | reference | reference | **self-check gate** — must be bitwise identical, or every conclusion from this factor is void | - | `training_reference_only` | native | reference | deviation gone ⇒ the training side is the source | - | `rollout_reference_only` | reference | native | deviation gone ⇒ the rollout side is the source | - - A factor with no reference implementation is a **parameter sweep** instead: - one arm per allowed value. -4. **Orders cases by rebuild cost** so a run reuses engines instead of restarting - the process ~160 times. -5. **Runs four gates before any verdict** — "not measured" and "measured and - clean" are different, and confusing them is the mistake an attribution - framework is most likely to make. -6. **Diagnoses** each factor, then **traces root causes** across all of them into - a ranked `MismatchReport`. - -## The five types that carry the whole design - -| type | what it is | -|---|---| -| `MismatchFactor` | one suspected cause — a switch, its comparison rules, its prerequisites, its pitfalls | -| `FactorVariant` | one ablation arm of that factor, as pasteable switch values | -| `VariantResult` | what one arm produced: status, metrics, evidence, effective config | -| `Diagnosis` | the verdict for one factor: training side, rollout side, both, not this factor, or *cannot tell* | -| `OperatorChecks` | the plugin protocol — one operator's factors plus how to read that operator back from an engine | +"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. -One factor expands into arms, each arm produces a result, the results become one -diagnosis, and all diagnoses become one report. Everything else is detail. +The output is a report: which factors were measured, which could not be and why, +and the few most suspicious modules ranked. -## The pipeline, in order +## Why the tail, not the mean -``` -1 registry collect plugins and factors pipeline/registry.py -2 planner filter by prerequisites pipeline/planner.py -2.5 planner reject contradictory declarations -3 planner expand factors into variants -4 planner order cases by rebind cost -5 runner run each arm on both sides pipeline/runner.py - → contracts compared field by field pipeline/comparison.py - → dlogp / ρ / clip_fraction -6 diagnosis four gates, then the matrix pipeline/diagnosis.py -7 report filter false positives, rank causes pipeline/report.py -``` +`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 | -The four gates in step 6, in order — none of them may be skipped: +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: what was tried, why rejected) | -| 2 · evidence | `required_evidence` is incomplete | `INSUFFICIENT_EVIDENCE` | -| 3 · shards | fewer logprob shards than `world_size` | `INSUFFICIENT_EVIDENCE` — a missing shard is wrong in a way that does not show | +| 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 the dangerous state: the reference was requested, the +`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`. Gate 1 exists for exactly this. +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 dataclasses and enums, no behaviour -│ ├── values.py PolicyRole, ExecutionPath, Precision, RebindCost, RequiredSetting -│ ├── collectives.py CollectiveContract, ReductionOrder, DeterminismLevel, rewrites -│ ├── contracts.py OperatorContract, ComparisonRule, ComparisonIssue -│ ├── factors.py MismatchFactor, Switch, Prerequisites, ReferenceImplementation -│ ├── variants.py FactorVariant, SwitchStatus, Diagnosis, NoiseFloor -│ ├── thresholds.py EXPECTED_RANGES — thresholds as code, not configuration -│ ├── rollout_context.py ComparisonIdentity, RolloutGroup, BatchPlacement -│ ├── metrics.py MismatchMetrics, VariantResult, FactorReport, LogprobShard -│ ├── fingerprints.py ReuseKey, ExecutionFingerprint, VariantRecord -│ ├── pitfalls.py KnownPitfall, FailureMode -│ └── tracing.py ModuleCorrespondence, PropagationEdge, RootCauseHypothesis -├── pipeline/ the seven steps, free functions only, no state -├── engines/ the two sides under test — megatron.py and vllm.py, nothing else -├── reference_adapters/ delivering pinned settings by channel and reading them back -├── model_meta/ per-model module correspondence and call chain (qwen3.py) -├── operator_checks/ plugins, one directory per operator — **empty by design** -└── __main__.py CLI, and the only module that imports operator plugins +├── 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 ``` -### What belongs in `engines/`, and what does not +`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. -`engines/` holds **exactly two modules: `megatron.py` (training side) and -`vllm.py` (rollout side)** — the two policies as they really run. Each one owns -how its engine is constructed, how a switch is delivered to it, and how the -*effective* value is read back off it. They are shared across operators: all of -attention's factors use the one `vllm.py`, and **adding an operator never adds a -file here**. Both are placeholders today; their docstrings list the settings that -must be pinned and the readback path for each. +Three dependency rules keep the plugin seam open: -Anything that merely satisfies the `ScoringBackend` protocol is **not** an -engine. Test harnesses live under `tests/`: - -| module | what it is | -|---|---| -| `tests/mismatch_cpu_backend.py` | CPU stub for exercising the gates and the matrix with no GPU, with an injectable one-sided bias | - -The line is role, not protocol: an engine is *a side under test*; a harness is -what lets you run the framework when that side is not available. - -Three dependency rules hold the plugin seam open: - -1. `schema/` never imports `pipeline/`; inside `schema/` the imports go one way, - with `values.py` at the top importing nothing from the project. -2. `pipeline/` never imports `operator_checks/` — it only sees what the registry - hands it. Break this and "add an operator" becomes "change the framework". -3. Only `__main__` imports `operator_checks/`, which is what triggers plugin - self-registration. +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 # registered operators and their factors -python -m rl_engine.mismatch plan --gpu-count 2 # expand into cases, cheapest rebuild first -python -m rl_engine.mismatch plan --json # same, machine readable +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 ``` -`list` currently prints "no operator plugins registered", and that is the -intended state: **the framework ships without operators.** Each operator is -claimed and written separately, and adding one changes nothing outside its own -directory plus one line in `__main__._OPERATOR_PACKAGES`. - -The plumbing is testable without a GPU: `tests/mismatch_cpu_backend.py` is a -scoring backend that can simulate a one-sided bias, a switch that silently does -nothing, and an implementation that is unstable across environments — so the -gates and the matrix are exercised for real in `tests/test_mismatch_framework.py`. - -## Noise floors: run them in order - -A factor's result only means something at a floor that can resolve it. The four -floors are arranged so **each step down adds exactly one new noise source**, which -is what makes a failure at one floor point at a known suspect set. - -| floor | configuration | new noise source | -|---|---|---| -| `SINGLE_LAYER_ANCHOR` | 1 layer, single device, determinism on, one token | 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 differences — the first floor with *real* mismatch | -| `PRODUCTION` | target TP/CP/PP, determinism off, decode path | everything else; the only floor whose numbers may be read against `EXPECTED_RANGES` | - -A floor that has not passed blocks the next one. +Every `adapter.py` currently raises `NotImplementedError`, but the declaration +layer works — so a factor can be checked as wired before anything is implemented. ## 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 mismatch factor (new operator, or a factor on an existing one) | [`docs/add-a-kernel-factor.md`](docs/add-a-kernel-factor.md) | -| add a communication feature (a collective, a reduction order, a rewrite) | [`docs/add-a-comm-feature.md`](docs/add-a-comm-feature.md) | - -Both are the same shape of work: **adding an operator is adding a directory, -adding a factor is adding a file**, and no existing file changes. If you find -yourself editing the planner, a global dict, or another operator's file, the -abstraction is missing something — say so and fix the framework rather than -patching in place. +| 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/docs/README.md b/rl_engine/mismatch/docs/README.md index 99ee4c57..bd454436 100644 --- a/rl_engine/mismatch/docs/README.md +++ b/rl_engine/mismatch/docs/README.md @@ -1,32 +1,11 @@ -# `rl_engine/mismatch` tutorials +# Tutorials -Start with the package overview in [`../README.md`](../README.md) — what a factor -is, the four arms, the four gates, and the seven pipeline steps. These tutorials -assume it. +Concepts live in [`../README.md`](../README.md); these are the how-to. -| tutorial | when you need it | +| tutorial | when | |---|---| -| [add-a-kernel-factor.md](add-a-kernel-factor.md) | a kernel (attention, GEMM, RoPE, logprob, SwiGLU, …) is suspected of computing something different on the two sides, and you want it measured and attributed | -| [add-a-comm-feature.md](add-a-comm-feature.md) | the suspect is a **collective**: a reduction order, an `all_reduce` → `reduce_scatter + all_gather` rewrite, a CP/split-K merge, a communication backend | - -Both end with the same two commands, which are how you check your work without a -GPU: - -```bash -python -m rl_engine.mismatch list -python -m rl_engine.mismatch plan --gpu-count 2 -``` - -## The rule both tutorials follow - -**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` -when the operator is new. - -If your change needs an edit inside `pipeline/`, a global dict, or another -operator's directory, stop: the framework is missing an abstraction. Raise it and -fix the framework, rather than patching around it — the seam is the only thing -keeping forty-odd factors from turning into forty-odd special cases. +| [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 index 837ef901..cf4cf08f 100644 --- a/rl_engine/mismatch/docs/add-a-comm-feature.md +++ b/rl_engine/mismatch/docs/add-a-comm-feature.md @@ -1,46 +1,25 @@ -# Tutorial: add a communication feature +# Add a communication feature -Mismatch exists because floating-point addition is not associative, and **the -accumulation order is almost entirely decided by collective communication**. So -collectives are not an implementation detail here — they are a first-class -declared object, and `gemm.forward_reduce`, `gemm.dgrad_reduce`, -`attn.cp_split_k_merge_order`, `logp.reduce_topology`, `moe.tp_forces_sp_reduce` -and `gemm.rollout_all_reduce_backend` are six instances of *one* semantic model. -Written separately they drift apart; written against `CollectiveContract` they -cannot. +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 tutorial adds one: `gemm.forward_reduce` — RowParallel forward reduction, -where the training side does `all_reduce` without sequence parallelism and -`reduce_scatter + all_gather` with it, and the two paths accumulate in different -orders. +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`. -Read [add-a-kernel-factor.md](add-a-kernel-factor.md) first. The file layout, the -four arms, and the gates are identical; this tutorial only covers what is -different when the suspect is a collective. - -## Step 1 · Describe the collective, do not just name it - -`CollectiveContract` is the full numerical semantics of one collective. Every -field is load-bearing: +## Describe the collective, do not just name it ```python -from rl_engine.mismatch.schema import ( - CollectiveContract, - CollectiveOp, - DeterminismLevel, - DowncastPoint, - ParallelDim, - Precision, - ReductionOrder, -) - ORDERED_REDUCE_SCATTER = CollectiveContract( op=CollectiveOp.REDUCE_SCATTER, group=ParallelDim.TENSOR, - group_size=2, + group_size=TP_SIZE, reduction_order=ReductionOrder.GLOBAL_RANK_INDEX, accumulate_precision=Precision.FP32, downcast_at=DowncastPoint.FINAL_WRITE, @@ -51,250 +30,122 @@ ORDERED_REDUCE_SCATTER = CollectiveContract( | field | what it decides | |---|---| -| `op` | which collective. `NONE` is a real value — record the single-device path explicitly rather than leaving it blank | -| `group` | which parallel dimension: `TENSOR` / `SEQUENCE` / `CONTEXT` / `EXPERT` / `PIPELINE` / `DATA` | -| `reduction_order` | **the direct root of mismatch.** `ARRIVAL` and `NCCL_ALGORITHM` are order-unstable; `GLOBAL_RANK_INDEX`, `GLOBAL_BLOCK_INDEX` and `GLOBAL_VOCAB_SHARD_INDEX` are the fixed orders | -| `accumulate_precision` + `downcast_at` | how much error the accumulation keeps. `PER_BLOCK` is the largest, `FINAL_WRITE` the smallest | -| `determinism` | how strong a reproducibility guarantee this implementation offers | -| `backend` | `"nccl"` / `"vllm_custom_ipc"` / `"mnnvl"` / `"transformer_engine"` / `"rl_kernel"` | +| `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 -`determinism=STABLE_ACROSS_TOPOLOGY` while reducing with `NCCL_ALGORITHM` or -`ARRIVAL` produces numbers that mean nothing. `reject_contradictory_factors()` -raises `ContradictoryFactor` at planning time — a static check, no GPU, no wasted -run. Do not weaken the claim to get past it; fix whichever half is wrong. +`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. -## Step 2 · Pin the contract so the planner can see it +## Pin the contract so the planner can see it -The planner finds your contract through `declared_collectives(factor)`, which -collects `CollectiveContract` values out of the reference's `required_settings`. -So the contract is pinned like any other setting — with the channel that actually -delivers 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 -# operator_checks/gemm/_common.py -from rl_engine.mismatch.schema import ( - ExecutionPath, - LibraryPin, - ReferenceAuthority, - ReferenceImplementation, - RequiredSetting, - SettingChannel, -) - -DETERMINISTIC_REDUCE_REFERENCE = ReferenceImplementation( - name="rl_kernel", - tier=ReferenceAuthority.SELF_WRITTEN, # justify this in the PR body - 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, # the contract itself, pinned as a value - 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", container_digest="sha256:..."),), +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, so a deterministic `all_reduce` / `reduce_scatter + all_gather` has to -be written. Say that in the PR rather than leaving the tier unexplained. +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. -The channel is not cosmetic — it decides when a setting can take effect and -therefore the rebind cost: `ENV_VAR` and `TORCH_GLOBAL` need a process restart, -`ENGINE_ARG` an engine rebuild, `CALL_ARG` nothing. +## What differs in the factor declaration -## Step 3 · The factor, with the collective in the comparison rules +See `operator_checks/gemm/factors/forward_reduce.py`. Four things are specific to +a comm factor: -```python -# operator_checks/gemm/factors/forward_reduce.py -from rl_engine.mismatch.operator_checks.gemm._common import DETERMINISTIC_REDUCE_REFERENCE -from rl_engine.mismatch.schema import ( - COLLECTIVE_CONTRACT, - ComparisonRule, - Evidence, - FactorCategory, - FailureMode, - KnownPitfall, - MismatchFactor, - NoiseFloor, - PolicyRole, - Prerequisites, - RebindCost, - Switch, -) - -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, - "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=DETERMINISTIC_REDUCE_REFERENCE, - call_sites=("attention.o_linear", "mlp.down_linear", "moe.output"), - 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, - ), - ), -) -``` - -Four things are specific to a comm factor: - -- **Indexed paths reach into the collective**: `collectives[0].reduction_order` +- **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*, not the library. Declaring `backend` as `MUST_MATCH_*` - buries the real finding under a difference you already knew about. -- **`min_gpu_count=2` plus `PROCESS_GROUP_REBUILD`.** The factor is identically - zero on one device, and `suggested_floor_is_lowest()` will tell you it belongs - at `SHARDED_SINGLE_NODE` or above. Running it at the anchor floor wastes time. -- **`call_sites`** records that one factor acts in several physical places — + 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 eating the same accumulation-order problem. One factor, three - sites, not three factors. + parallel linears with the same accumulation-order problem. One factor, three + sites. -`compare_contracts()` adds one check you do not declare: if the two sides' -collectives promise different `DeterminismLevel`s, it emits -`DETERMINISM_INCOMPATIBLE`. Comparing a topology-independent implementation -against one that is not even reproducible across runs measures the weaker side's -noise, not the gap between the sides. +`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. -## Step 4 · The cheapest strong check: rerun under different NCCL settings +## The cheapest strong check -An implementation claiming `STABLE_ACROSS_TOPOLOGY` must produce **bitwise -identical** results when NCCL is told to use a different algorithm. This needs no -cross-framework comparison at all — one side, rerun — which makes it the highest -value-per-minute check in the whole framework. - -Declare it on the arm with `repeat_under`, the single exception to "one variant, -one execution": +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 -from rl_engine.mismatch.schema import ExpectedOutcome, FactorVariant, PolicyRole - -TOPOLOGY_INVARIANCE = FactorVariant( - name="both_reference", - switch_values={"gemm.forward_reduce": "rl_kernel"}, - replace_on={ - PolicyRole.ROLLOUT: "rl_engine.kernels.collectives.ordered_reduce_scatter", - PolicyRole.TRAINING: "rl_engine.kernels.collectives.ordered_reduce_scatter", - }, - expected=ExpectedOutcome.BITWISE_IDENTICAL, - repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")}, - why="a fixed order must survive NCCL choosing a different algorithm", -) +repeat_under={"NCCL_ALGO": ("Ring", "Tree"), "NCCL_PROTO": ("Simple", "LL")} ``` -The runner expands the cartesian product (four runs here) and asserts they agree -bitwise; a disagreement sets `SwitchStatus.ERROR`, which gate 1 turns into -`VARIANT_DID_NOT_APPLY` rather than a numeric verdict. +The runner expands the cartesian product and asserts bitwise agreement; a +disagreement sets `ERROR`, which gate 1 turns into `VARIANT_DID_NOT_APPLY`. -What this protects is **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. +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 arm explicitly through `MismatchFactor.variants` -(a non-empty `variants` tuple is returned as-is by `build_variants`), so declare -all four arms there when you need it on one of them. +To use `repeat_under` you pass the arms explicitly through +`MismatchFactor.variants` — a non-empty tuple is returned as-is by +`build_variants`. -## Step 5 · Observe what actually ran +## Observe what actually ran -```python -# operator_checks/gemm/adapter.py -def observe_collectives(role: PolicyRole, adapter: Any) -> tuple[CollectiveContract, ...]: - """The collectives this operator really performed, this run.""" - return tuple(adapter.collective_trace()) # not what the config asked for -``` +`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. -This feeds the `COLLECTIVE_CONTRACT` evidence item that the factor declares as -required, and gate 2 refuses a verdict without it. The distinction that matters: -`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 the case where the two disagree, and only the second one is -evidence. +## Rewrites -## Declaring a rewrite - -"Mathematically identical, unequal in floating point" is common enough to be its -own type. Two are already declared in `schema/collectives.py`: +"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 rewrite when sequence parallelism is on, -while the rollout side does not, *is* `gemm.forward_reduce`. If your factor is -"one side rewrites this collective and the other does not", add a -`CollectiveRewrite` constant next to those two and reference it from the factor's -`question`, rather than describing the rewrite in prose. +`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. -## When you genuinely need a new enum value +## Adding an enum value -Adding a `CollectiveOp`, `ParallelDim` or `ReductionOrder` value **is** a -framework change, and the plugin seam exists so that this is rare. It is -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. It is *not* justified for a new backend (that is the +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: - -1. Put it in `schema/collectives.py` with a comment saying what orders it. -2. Extend `_NON_DETERMINISTIC_ORDERS` in `pipeline/planner.py` if the new order is - not stable, so the contradiction check keeps working. -3. Add a case to `tests/test_mismatch_framework.py`. -4. Say in the PR why an existing value could not express it. +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 before the PR +## Checklist -- [ ] Every field of `CollectiveContract` is filled from what the code does, not from what the config requests. -- [ ] `determinism` and `reduction_order` do not contradict each other (the planner will tell you, but know why). +- [ ] 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 and the suggested floor is `SHARDED_SINGLE_NODE` or above. -- [ ] `NCCL_ALGO` / `NCCL_PROTO` are pinned as `RequiredSetting`s **with readback**. +- [ ] `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, and `COLLECTIVE_CONTRACT` is in `required_evidence`. -- [ ] `call_sites` lists every physical place this one factor acts. -- [ ] A `SELF_WRITTEN` reference is justified against `SHARED_BACKEND` in the PR body. +- [ ] `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 index a35ecd27..d706e7ea 100644 --- a/rl_engine/mismatch/docs/add-a-kernel-factor.md +++ b/rl_engine/mismatch/docs/add-a-kernel-factor.md @@ -1,98 +1,61 @@ -# Tutorial: add a kernel's mismatch factor +# Add a kernel's mismatch factor -You suspect a kernel — attention, GEMM, RoPE, the logprob path, SwiGLU — of -computing something different on the training side than on the rollout side, and -you want that suspicion **measured and attributed to a side** instead of argued -about. +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`. -This tutorial walks the whole path with one worked example: `attn.rope_fusion`, -"is the difference caused by fused RoPE versus small-operator RoPE, or by -`position_ids` / `theta` / the cast boundary?" +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. -By the end you will have written **declarations only** — no execution logic. The -framework expands the arms, runs them, gates them, and draws the conclusion. +## Three decisions before writing anything -## Before you write anything: three decisions +**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. -**1 · Is this a parameter sweep or an implementation swap?** - -| | you declare | the framework expands into | -|---|---|---| -| **sweep** — nothing is replaced, you scan a setting | `Switch.allowed_values`, `reference=None` | one arm per allowed value | -| **swap** — a reference implementation replaces the native one | `Switch` + `ReferenceImplementation` | the four arms (plus an `fp64_oracle` arm if declared) | - -There is deliberately no `kind` field: `reference is None` *is* the distinction, -because derivable state is state that can disagree with itself. RoPE is a swap; -`logp.precision_downcast` (in `operator_checks/logprob/`) is the worked sweep. - -**Know what a sweep gives up.** `diagnose()` runs the four-arm matrix, which -needs a `both_native` baseline and the two one-sided arms. A sweep has neither, -so it returns `INSUFFICIENT_EVIDENCE` — *by design*: with nothing swapped there is -no side to attribute to. A sweep measures and records; it does not conclude. Read -its numbers against `EXPECTED_RANGES` yourself, or declare explicit `variants` -once a reference implementation exists. - -**2 · Which reference implementation, and are you allowed to write one?** - -`ReferenceAuthority` is a decision order, not a description: +**Which reference?** `ReferenceAuthority` is a decision order: ``` -FP64_ORACLE slow, mathematically exact — gold standard at the lowest floor +FP64_ORACLE slow, exact — gold standard at the lowest floor ↑ validates -SHARED_BACKEND TransformerEngine (training) / FlashInfer (rollout) — look here FIRST +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 (`pytorch/attention/rope.py`) and by -`flashinfer.rope`, so it is `SHARED_BACKEND`. +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`. -**3 · Which noise floor can even show it?** +**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. -Pick the *lowest* floor at which the factor is not identically zero. RoPE shows up -with one layer on one device, so `SINGLE_LAYER_ANCHOR` — cheap, and a bitwise -failure there is an operator bug rather than mismatch. A factor whose switch is -`PROCESS_GROUP_REBUILD` is identical on a single device; running it at the anchor -floor only burns machine time (`planner.suggested_floor_is_lowest`). +## Step 1 — the directory -## Step 1 · Create the operator directory - -Skip to step 3 if the operator already exists — then you are only adding one file. +Skip to step 3 if the operator exists; then you are only adding one file. ``` -rl_engine/mismatch/operator_checks/attention/ -├── __init__.py ~15 lines: operator name + discover_factors, everything else delegated -├── adapter.py the four operator-level methods -├── _common.py shared across factors: reference implementations, contract helpers +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 empty - └── rope_fusion.py one FACTOR constant, 30–50 lines + ├── __init__.py + └── rope_fusion.py ``` -Why one file per factor: an operator has a dozen-plus factors of eight-odd fields -each. In one file that is 800 lines nobody wants to edit. - -**The file name is the factor id with the operator prefix stripped.** +**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 instead of silently dropping the factor. +fails at import rather than silently dropping the factor. -## Step 2 · `_common.py` — what the whole operator shares +## Step 2 — `_common.py` ```python -# operator_checks/attention/_common.py -from rl_engine.mismatch.schema import ( - ExecutionPath, - LibraryPin, - ReferenceAuthority, - ReferenceImplementation, - RequiredSetting, - SettingChannel, -) - TE_ROPE_REFERENCE = ReferenceImplementation( name="transformer_engine", tier=ReferenceAuthority.SHARED_BACKEND, @@ -104,251 +67,91 @@ TE_ROPE_REFERENCE = ReferenceImplementation( ), required_settings=( RequiredSetting( - "NVTE_ALLOW_NONDETERMINISTIC_ALGO", - "0", - SettingChannel.ENV_VAR, - readback="os.environ", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0", + SettingChannel.ENV_VAR, readback="os.environ", ), ), - pinned_libraries=( - LibraryPin("transformer_engine", "2.3.0", container_digest="sha256:..."), - ), + 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.** The gate is not - "both sides use the same implementation", it is "*every path this reference - covers must agree bitwise on the same sequence*". Two paths ⇒ the gate holds - across the two sides. A reference that also covers `ROLLOUT_DECODE` (FlashInfer - attention does) makes the gate hold *inside the rollout side*, and then you do - not need a decode stub on the training side at all. +- **`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 - (`reference_adapters/settings.py::apply_required_settings`) and verified by - `verify_required_settings`. A setting with `readback=None` can only ever be - recorded `UNOBSERVABLE` — delivered but unprovable is the same as not - delivered. + and verified by readback. `readback=None` can only ever be `UNOBSERVABLE`. - **`pinned_libraries` is required.** TE and FlashInfer change kernel selection - across versions; the same factor can reach the *opposite* conclusion on a - different version. The pin goes into the execution fingerprint, so bumping it - invalidates historical results instead of quietly making them incomparable. + across versions, so the same factor can reach the opposite conclusion on a + different one. -## Step 3 · The factor file — declarations, no behaviour - -```python -# operator_checks/attention/factors/rope_fusion.py -from rl_engine.mismatch.operator_checks.attention._common import 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, - "precision.downcast_at": ComparisonRule.MUST_MATCH_SEMANTICALLY, - "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), - 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, - ), - ), -) -``` +## Step 3 — the factor file -### Filling the fields well +One `FACTOR` constant, no behaviour. See +`operator_checks/attention/factors/rope_fusion.py` for the full declaration; the +fields worth thinking about: -**`comparison_rules` — the tier matters more than the field list.** -Keys are dotted paths from the contract root: `precision.accumulate`, -`collectives[0].reduction_order`, `extra.rope_theta`. The framework indexes both -contracts by path and compares entry by entry, so you never write a "collect the -comparable fields" function — you only put fields in the right place in -`build_contract()`. +**`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 | use it when | consequence | +| tier | for | consequence | |---|---|---| -| `MUST_MATCH_BITWISE` | identity: shapes, dtypes, TP size, vocab size | differing ⇒ the case is void, not a finding | -| `MUST_MATCH_SEMANTICALLY` | the two sides may implement it differently but must mean the same thing | differing ⇒ a `SEMANTIC_MISMATCH` issue | -| `RECORD_ONLY` | representation that differs by construction — packed QKV layout, page tables, backend names | recorded, **never compared** | - -`RECORD_ONLY` exists precisely so structural differences do not drown the real -problems. Declare a packed-QKV-style field `MUST_MATCH_*` and your factor -disappears under false positives. - -One more constraint the registry enforces: **two factors may not declare the same -contract field at two different rules.** If that fires, one of the two -declarations is wrong — do not "fix" it by renaming the path. - -**`required_evidence` — what must exist before a verdict is allowed.** -Three items apply to every factor (`Evidence.EFFECTIVE_CONFIG_READBACK`, -`MODEL_STATE_FINGERPRINT`, `LIBRARY_VERSIONS`). Operator-specific evidence is a -plain string constant (`POSITION_CACHE`, `LSE_EXPORT`, `VOCAB_SHARD_MAP`, -`COLLECTIVE_CONTRACT`, …) so that adding an operator never means editing an enum -in the framework. Missing evidence is gate 2: `INSUFFICIENT_EVIDENCE`, which is -**not** the same verdict as "measured and clean". - -**`prerequisites` — a whitelist, declared, not probed by hand.** -`required_ops`, `min_gpu_count`, `required_packages`, `required_model_traits` -(`"moe"`, `"linear_attention"`), `blocked_by` (issues this waits on). The planner -turns each unmet item into a reason string, so `plan` prints *why* a factor was -skipped instead of silently omitting it. - -**`pitfalls` — prose gets read once; data gets enforced.** -`symptom` and `actual_cause` are separate fields on purpose: a pitfall is a -pitfall because its appearance points at the wrong cause. `guard_runs_at` should -be the lowest floor that can run the check — cheap guards first. - -## Step 4 · `adapter.py` — the four operator-level methods - -These are operator-level, not factor-level: how to read config back from an -engine is one piece of logic shared by all of the operator's factors. - -```python -# operator_checks/attention/adapter.py -from typing import Any, Callable, Mapping - -from rl_engine.mismatch.schema import ( - DowncastPoint, - ImplementationResolution, - OperatorContract, - PolicyRole, - Precision, - PrecisionProfile, - RejectedCandidate, -) +| `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. -def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> OperatorContract: - """This side's switch values -> this side's numerical contract.""" - return OperatorContract( - operator="attention", - role=role, - precision=PrecisionProfile( - compute=Precision.BF16, - accumulate=Precision.FP32, - downcast_at=DowncastPoint.FINAL_WRITE, - softmax_accumulate=Precision.FP32, - ), - collectives=(), # see the comm tutorial - extra={ # keep flat: paths stay short and readable - "rope_theta": 1_000_000.0, - "position_ids_digest": "...", - "post_rope_qk_digest": "...", - "fusion_boundary": switch_values.get("attn.rope_fusion", "native"), - }, - ) - - -def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: - """Read switches back **from the engine**. A requested value is not evidence.""" - ... +**`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. -def observe_collectives(role: PolicyRole, adapter: Any): - """Which collectives actually ran. RoPE has none.""" - return () +**`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` -def resolve_implementation( - factor_id: str, role: PolicyRole, impl_name: str -) -> tuple[Callable[..., Any] | None, ImplementationResolution]: - """Resolve the name an arm asks for into a callable. - - Return the trace **even when resolution fails**: which candidates were tried - and why each was rejected. - """ - rejected: list[RejectedCandidate] = [] - for candidate in _candidates_for(impl_name): - try: - return _import(candidate), ImplementationResolution(impl_name, candidate) - except ImportError as exc: - rejected.append(RejectedCandidate(candidate, str(exc))) - return None, ImplementationResolution(impl_name, None, tuple(rejected)) -``` +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 `SwitchStatus.FELL_BACK` with nothing to investigate, and a -fallen-back arm whose deviation "did not change" reads exactly like a clean -`NOT_THIS_FACTOR`. Gate 1 catches the status; only your trace explains it. +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 · `__init__.py` — register, and nothing else +## Step 5 — register ```python -# operator_checks/attention/__init__.py -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__) # scans factors/*.py + 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) -``` - -Adding the next factor drops a file into `factors/`; this file does not change. - -## Step 6 · Make the operator exist - -The framework does not know which operators exist. `__main__` is the only module -that imports plugins, which is what keeps `pipeline/` from ever reaching into -`operator_checks/`: - -```python -# rl_engine/mismatch/__main__.py -_OPERATOR_PACKAGES: tuple[str, ...] = ( - "rl_engine.mismatch.operator_checks.attention", -) + ... ``` -An operator that is not listed here simply does not exist as far as the framework -is concerned. **This one line is the only edit outside your own directory.** +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 7 · Check it, without a GPU +## Step 6 — check it, without a GPU ```bash python -m rl_engine.mismatch list @@ -356,58 +159,34 @@ python -m rl_engine.mismatch list # attn.rope_fusion kernel_implementation ``` -`plan` on a workstation without TransformerEngine reports the factor as skipped, -naming every unmet prerequisite rather than silently omitting it: +`plan` on a machine without the prerequisites names every unmet one rather than +silently omitting the factor: -```bash -python -m rl_engine.mismatch plan --gpu-count 1 -# noise floor: single_layer_anchor -# runnable factors: 0 cases: 0 -# -# skipped (prerequisites not met): -# attn.rope_fusion: operator 'rope' is not dispatchable -# attn.rope_fusion: package 'transformer_engine>=2.0' is not installed ``` - -That is a statement about this machine, not a defect in the factor. Where the -prerequisites are met, the same declaration expands into the four arms, ordered -cheapest-rebuild-first: - +skipped (prerequisites not met): + attn.rope_fusion: operator 'rope' is not dispatchable + attn.rope_fusion: package 'transformer_engine>=2.0' is not installed ``` -# runnable factors: 1 cases: 4 -# cases in execution order (cheapest rebuild first): -# [engine_rebuild ] attn.rope_fusion :: both_native -# [engine_rebuild ] attn.rope_fusion :: both_reference -# [engine_rebuild ] attn.rope_fusion :: training_reference_only -# [engine_rebuild ] attn.rope_fusion :: rollout_reference_only -``` - -`plan --json` gives the same thing machine readably. -Then add a test beside `tests/test_mismatch_framework.py`. The CPU backend in -`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. +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 the framework does so you do not have to +## What you do not have to write -| you might reach for | it already exists | -|---|---| -| writing the four arms out by hand | `build_variants()` | -| a "compare these fields" helper | `compare_contracts()`, driven by `comparison_rules` | -| deciding whether the numbers converged | `diagnose()` — four gates, then the matrix, judged on `clip_fraction` | -| checking prerequisites and printing why one was skipped | `missing_prerequisites()` | -| ordering runs so engines get reused | `order_cases_by_rebind_cost()` | -| re-running under several `NCCL_ALGO` values and asserting bitwise equality | `FactorVariant.repeat_under` + `assert_order_is_topology_independent()` | +`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 before the PR +## Checklist - [ ] File name equals the factor id minus the operator prefix. - [ ] `question` is one line and says what the factor *answers*. -- [ ] Every representation-only field is `RECORD_ONLY`, not `MUST_MATCH_*`. -- [ ] A `SELF_WRITTEN` reference is justified against `SHARED_BACKEND` in the PR body. -- [ ] Every `RequiredSetting` has a `readback`, or you have accepted `UNOBSERVABLE`. -- [ ] `pinned_libraries` names an exact version, ideally with a container digest. -- [ ] Known pitfalls are encoded as `KnownPitfall`, not written in a comment. -- [ ] `python -m rl_engine.mismatch list` and `plan` both show what you expect. +- [ ] 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. From 0c770b365cd20df40370a5b1f3216283c55c0028 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Mon, 10 Aug 2026 20:18:08 +0800 Subject: [PATCH 5/7] fix lint check Signed-off-by: Zhang Jian --- rl_engine/mismatch/__init__.py | 5 ++ rl_engine/mismatch/model_meta/__init__.py | 5 ++ .../operator_checks/attention/_common.py | 4 +- .../operator_checks/attention/adapter.py | 4 +- rl_engine/mismatch/pipeline/__init__.py | 32 ++++++++ .../mismatch/reference_adapters/__init__.py | 6 ++ rl_engine/mismatch/schema/__init__.py | 81 +++++++++++++++++++ rl_engine/mismatch/schema/metrics.py | 7 +- tests/mismatch_cpu_backend.py | 7 +- tests/test_mismatch_framework.py | 2 +- 10 files changed, 134 insertions(+), 19 deletions(-) diff --git a/rl_engine/mismatch/__init__.py b/rl_engine/mismatch/__init__.py index a8284b3c..b99d452e 100644 --- a/rl_engine/mismatch/__init__.py +++ b/rl_engine/mismatch/__init__.py @@ -20,3 +20,8 @@ """ from rl_engine.mismatch import pipeline, schema + +__all__ = [ + "pipeline", + "schema", +] diff --git a/rl_engine/mismatch/model_meta/__init__.py b/rl_engine/mismatch/model_meta/__init__.py index cb66664b..9d977f5b 100644 --- a/rl_engine/mismatch/model_meta/__init__.py +++ b/rl_engine/mismatch/model_meta/__init__.py @@ -10,3 +10,8 @@ """ from rl_engine.mismatch.model_meta.qwen3 import QWEN3_CORRESPONDENCES, QWEN3_EDGES + +__all__ = [ + "QWEN3_CORRESPONDENCES", + "QWEN3_EDGES", +] diff --git a/rl_engine/mismatch/operator_checks/attention/_common.py b/rl_engine/mismatch/operator_checks/attention/_common.py index 8754b8cb..d89dcfdf 100644 --- a/rl_engine/mismatch/operator_checks/attention/_common.py +++ b/rl_engine/mismatch/operator_checks/attention/_common.py @@ -31,7 +31,5 @@ readback="os.environ", ), ), - pinned_libraries=( - LibraryPin("transformer_engine", "2.9.0.dev0", commit="8260f49"), - ), + pinned_libraries=(LibraryPin("transformer_engine", "2.9.0.dev0", commit="8260f49"),), ) diff --git a/rl_engine/mismatch/operator_checks/attention/adapter.py b/rl_engine/mismatch/operator_checks/attention/adapter.py index 9caf36f9..ab6c6875 100644 --- a/rl_engine/mismatch/operator_checks/attention/adapter.py +++ b/rl_engine/mismatch/operator_checks/attention/adapter.py @@ -55,6 +55,4 @@ def resolve_implementation( FlashInfer when it is available, so which one answered is itself evidence. """ - raise NotImplementedError( - "Import the candidates in order and return (callable, resolution)." - ) + raise NotImplementedError("Import the candidates in order and return (callable, resolution).") diff --git a/rl_engine/mismatch/pipeline/__init__.py b/rl_engine/mismatch/pipeline/__init__.py index 9881660f..128cd99a 100644 --- a/rl_engine/mismatch/pipeline/__init__.py +++ b/rl_engine/mismatch/pipeline/__init__.py @@ -49,3 +49,35 @@ 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/reference_adapters/__init__.py b/rl_engine/mismatch/reference_adapters/__init__.py index f03fce17..8c460842 100644 --- a/rl_engine/mismatch/reference_adapters/__init__.py +++ b/rl_engine/mismatch/reference_adapters/__init__.py @@ -13,3 +13,9 @@ apply_required_settings, verify_required_settings, ) + +__all__ = [ + "SettingDeliveryError", + "apply_required_settings", + "verify_required_settings", +] diff --git a/rl_engine/mismatch/schema/__init__.py b/rl_engine/mismatch/schema/__init__.py index b33c9777..b093accf 100644 --- a/rl_engine/mismatch/schema/__init__.py +++ b/rl_engine/mismatch/schema/__init__.py @@ -111,3 +111,84 @@ 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/metrics.py b/rl_engine/mismatch/schema/metrics.py index e74606a2..1267ee89 100644 --- a/rl_engine/mismatch/schema/metrics.py +++ b/rl_engine/mismatch/schema/metrics.py @@ -21,12 +21,7 @@ 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, -) +from rl_engine.mismatch.schema.variants import Diagnosis, FactorVariant, NoiseFloor, SwitchStatus DEFAULT_CLIP_EPS = 0.2 diff --git a/tests/mismatch_cpu_backend.py b/tests/mismatch_cpu_backend.py index 0bef380f..51b8c46b 100644 --- a/tests/mismatch_cpu_backend.py +++ b/tests/mismatch_cpu_backend.py @@ -16,12 +16,7 @@ from dataclasses import dataclass, field from typing import Any, Callable, Mapping, Sequence -from rl_engine.mismatch.schema import ( - ComparisonIdentity, - Evidence, - PolicyRole, - ReuseKey, -) +from rl_engine.mismatch.schema import ComparisonIdentity, Evidence, PolicyRole, ReuseKey @dataclass diff --git a/tests/test_mismatch_framework.py b/tests/test_mismatch_framework.py index c9d2cb09..e8b58125 100644 --- a/tests/test_mismatch_framework.py +++ b/tests/test_mismatch_framework.py @@ -17,7 +17,6 @@ import pytest -from tests.mismatch_cpu_backend import CpuScoringBackend from rl_engine.mismatch.pipeline import ( ContradictoryFactor, PluginRegistry, @@ -79,6 +78,7 @@ reuse_level, tolerance_floor, ) +from tests.mismatch_cpu_backend import CpuScoringBackend # ---------------------------------------------------------------- fixtures -- From cb24c6ccb2fbfbb3faedb69fae09d63a28ce94f4 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 12 Aug 2026 19:36:00 +0800 Subject: [PATCH 6/7] feat(mismatch): implement the logprob adapter and its LSE-merge swap factor Implement logprob's four operator-level methods against the WS2 TP-aware contract semantics (issue #241): build_contract maps each side's partial-LSE merge onto the collective schema (all_reduce in NCCL order on training, full-logits gather on rollout, fixed vocab-shard-order merge for the rl_kernel reference), read_effective_config/observe_collectives report engine state rather than requests, and resolve_implementation returns the rejection trace instead of a bare None. Add logp.lse_merge_order, the swap factor backed by the WS2 deterministic vocab-parallel logprob reference. SELF_WRITTEN because neither TE nor FlashInfer offers a vocab-parallel selected-logprob with a topology-fixed merge order. Its prerequisite gates on the vocab_parallel_logp op, which ships in issue #241 PR3, so plan reports it as skipped until that lands. Run the mismatch test suite in CI; it was not wired into any job. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_012PyjQEqDJwy9Cos4Sb9QBK --- .github/workflows/ci.yml | 4 + rl_engine/mismatch/README.md | 7 +- .../operator_checks/logprob/_common.py | 109 +++++++- .../operator_checks/logprob/adapter.py | 198 +++++++++++++- .../logprob/factors/lse_merge_order.py | 116 +++++++++ tests/test_mismatch_logprob_adapter.py | 242 ++++++++++++++++++ 6 files changed, 661 insertions(+), 15 deletions(-) create mode 100644 rl_engine/mismatch/operator_checks/logprob/factors/lse_merge_order.py create mode 100644 tests/test_mismatch_logprob_adapter.py 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 index d2fea2bf..770f277f 100644 --- a/rl_engine/mismatch/README.md +++ b/rl_engine/mismatch/README.md @@ -102,8 +102,11 @@ python -m rl_engine.mismatch plan --gpu-count 2 # expand into cases, cheape python -m rl_engine.mismatch plan --json ``` -Every `adapter.py` currently raises `NotImplementedError`, but the declaration -layer works — so a factor can be checked as wired before anything is implemented. +The gemm and attention `adapter.py` files still raise `NotImplementedError`, but +the declaration layer works — so a factor can be checked as wired before +anything is implemented. Logprob's four methods are implemented against the WS2 +TP-aware contract semantics (issue #241) and are the worked example of the +adapter layer. ## Adding to this package diff --git a/rl_engine/mismatch/operator_checks/logprob/_common.py b/rl_engine/mismatch/operator_checks/logprob/_common.py index f2475bbc..c490c784 100644 --- a/rl_engine/mismatch/operator_checks/logprob/_common.py +++ b/rl_engine/mismatch/operator_checks/logprob/_common.py @@ -1,15 +1,31 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Values shared by logprob's factors. +"""Values shared by logprob's factors: contracts, shard maps, the reference. -No reference implementation yet: the rollout side has no deterministic -vocab-parallel reduction to swap in, so these factors are parameter sweeps. +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 DowncastPoint, Precision +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] = { @@ -21,3 +37,88 @@ "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 index 142d6de8..60cf9544 100644 --- a/rl_engine/mismatch/operator_checks/logprob/adapter.py +++ b/rl_engine/mismatch/operator_checks/logprob/adapter.py @@ -1,19 +1,66 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Logprob's four operator-level methods. Not implemented.""" +"""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``. @@ -24,25 +71,158 @@ def build_contract(role: PolicyRole, switch_values: Mapping[str, Any]) -> Operat without recording both maps is meaningless. """ - raise NotImplementedError + 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]: - raise NotImplementedError( - "Read the head dtype, the transform chain and the vocab shard " - "boundaries off the live engine." + """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 reduction otherwise.""" + """Return the collective trace: empty at TP=1, the partial-LSE merge otherwise. - raise NotImplementedError + 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]: - """Unused until a logprob factor gains a reference implementation.""" + """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``. + """ - raise NotImplementedError + 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/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/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"] From 3b3919334272e152d35465c9d8fb97ace88c59d8 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 23:35:50 +0800 Subject: [PATCH 7/7] feat(mismatch): validate attention runtime contracts --- rl_engine/mismatch/README.md | 9 +- .../operator_checks/attention/_common.py | 605 +++++++++++++++++- .../operator_checks/attention/adapter.py | 342 +++++++++- .../attention/factors/cp_merge.py | 85 +++ .../attention/factors/precision_downcast.py | 54 ++ .../attention/factors/rope_fusion.py | 14 +- .../attention/factors/split_kv.py | 85 +++ rl_engine/mismatch/pipeline/registry.py | 9 + rl_engine/mismatch/pipeline/runner.py | 17 +- tests/test_mismatch_attention_adapter.py | 314 +++++++++ tests/test_mismatch_attention_factors.py | 153 +++++ tests/test_mismatch_framework.py | 79 +++ 12 files changed, 1727 insertions(+), 39 deletions(-) create mode 100644 rl_engine/mismatch/operator_checks/attention/factors/cp_merge.py create mode 100644 rl_engine/mismatch/operator_checks/attention/factors/precision_downcast.py create mode 100644 rl_engine/mismatch/operator_checks/attention/factors/split_kv.py create mode 100644 tests/test_mismatch_attention_adapter.py create mode 100644 tests/test_mismatch_attention_factors.py diff --git a/rl_engine/mismatch/README.md b/rl_engine/mismatch/README.md index 770f277f..8446577c 100644 --- a/rl_engine/mismatch/README.md +++ b/rl_engine/mismatch/README.md @@ -102,11 +102,10 @@ python -m rl_engine.mismatch plan --gpu-count 2 # expand into cases, cheape python -m rl_engine.mismatch plan --json ``` -The gemm and attention `adapter.py` files still raise `NotImplementedError`, but -the declaration layer works — so a factor can be checked as wired before -anything is implemented. Logprob's four methods are implemented against the WS2 -TP-aware contract semantics (issue #241) and are the worked example of the -adapter layer. +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 diff --git a/rl_engine/mismatch/operator_checks/attention/_common.py b/rl_engine/mismatch/operator_checks/attention/_common.py index d89dcfdf..b60e990f 100644 --- a/rl_engine/mismatch/operator_checks/attention/_common.py +++ b/rl_engine/mismatch/operator_checks/attention/_common.py @@ -1,19 +1,75 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Reference implementations shared by attention's factors.""" +"""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, @@ -33,3 +89,550 @@ ), 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 index ab6c6875..6f104314 100644 --- a/rl_engine/mismatch/operator_checks/attention/adapter.py +++ b/rl_engine/mismatch/operator_checks/attention/adapter.py @@ -1,58 +1,348 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Attention's four operator-level methods. Not implemented. - -They are operator-level rather than factor-level because reading configuration -back from an engine is the same logic for all of attention's factors. -""" +"""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: - raise NotImplementedError( - "Return this side's contract with each field at the path its factor " - "declared in comparison_rules: precision.*, collectives[i].*, and " - "attention's own state under a flat extra.*" + """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" + ) -def read_effective_config(role: PolicyRole, adapter: Any) -> Mapping[str, Any]: - raise NotImplementedError( - "Read actual values off the live engine. Megatron's AttnBackend=auto and " - "vLLM's backend selection both pick per shape, so a requested value is " - "not evidence." + 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, ...]: - raise NotImplementedError("Return the collectives that actually ran.") + """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 an arm's implementation name, returning the trace even on failure. + """Resolve every candidate in order and preserve every rejection reason.""" - 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``. + 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 - Candidates for the RoPE reference, in order: - ``transformer_engine.pytorch.attention.rope:apply_rotary_pos_emb`` on the - training side; ``flashinfer.rope:apply_rope`` then vLLM's - ``rotary_embedding:get_rope`` on the rollout side. vLLM forwards to - FlashInfer when it is available, so which one answered is itself evidence. - """ - raise NotImplementedError("Import the candidates in order and return (callable, resolution).") +__all__ = [ + "AttentionAdapterError", + "build_contract", + "observe_collectives", + "read_effective_config", + "resolve_implementation", +] 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 index 5050b9b5..a48f5de9 100644 --- a/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py +++ b/rl_engine/mismatch/operator_checks/attention/factors/rope_fusion.py @@ -5,7 +5,10 @@ from __future__ import annotations -from rl_engine.mismatch.operator_checks.attention._common import TE_ROPE_REFERENCE +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, @@ -39,6 +42,9 @@ "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. @@ -48,7 +54,11 @@ required_ops=("rope",), required_packages=("transformer_engine>=2.0",), ), - required_evidence=(Evidence.EFFECTIVE_CONFIG_READBACK.value, POSITION_CACHE), + required_evidence=( + Evidence.EFFECTIVE_CONFIG_READBACK.value, + POSITION_CACHE, + POST_ROPE_QK_EVIDENCE, + ), reference=TE_ROPE_REFERENCE, pitfalls=( KnownPitfall( 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/pipeline/registry.py b/rl_engine/mismatch/pipeline/registry.py index e9a39cbe..69faaeae 100644 --- a/rl_engine/mismatch/pipeline/registry.py +++ b/rl_engine/mismatch/pipeline/registry.py @@ -101,11 +101,17 @@ def _check_factor_conflicts(self, plugin: OperatorChecks) -> None: 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: @@ -113,6 +119,9 @@ def _check_factor_conflicts(self, plugin: OperatorChecks) -> None: 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)) diff --git a/rl_engine/mismatch/pipeline/runner.py b/rl_engine/mismatch/pipeline/runner.py index f66c1481..93250f09 100644 --- a/rl_engine/mismatch/pipeline/runner.py +++ b/rl_engine/mismatch/pipeline/runner.py @@ -7,7 +7,7 @@ import itertools import math -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Callable, Mapping, Protocol, Sequence from rl_engine.mismatch.pipeline.comparison import compare_contracts @@ -220,10 +220,17 @@ def run_variant( if not assert_order_is_topology_independent(repeats[role]): status = SwitchStatus.ERROR - contracts = { - role: checks.build_contract(role, variant.switch_values) + 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,) ) @@ -248,8 +255,8 @@ def run_variant( evidence=evidence, effective_config={ f"{role.value}.{key}": value - for role, readback in readbacks.items() - for key, value in readback.items() + for role, effective in effective_configs.items() + for key, value in effective.items() if key != "evidence" }, collectives_observed=tuple( 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 index e8b58125..58f11936 100644 --- a/tests/test_mismatch_framework.py +++ b/tests/test_mismatch_framework.py @@ -629,6 +629,45 @@ def declare_factors(self): 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() == () @@ -695,6 +734,46 @@ def test_runner_detects_a_one_sided_injected_deviation(): 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."""