diff --git a/docs/en/advanced/rl-kernel-operator-adapter.md b/docs/en/advanced/rl-kernel-operator-adapter.md index b65d9cc26..14a9a2103 100644 --- a/docs/en/advanced/rl-kernel-operator-adapter.md +++ b/docs/en/advanced/rl-kernel-operator-adapter.md @@ -22,3 +22,22 @@ The current Phase 1 hooks are: - `provenance()`. Future `linear_logp` integration work should extend this adapter package rather than adding direct `rl_engine` imports to Megatron, rollout, or training modules. + +## Alignment Standard Boundary + +The module mismatch matrix is owned by RL-Kernel. vime reads the public +`rl_engine.alignment.cross_config.debug_matrix.module_debug_matrix()` manifest +through `vime.backends.rl_kernel_utils.standard`; it does not carry a local +copy of Attention, FFN, or logprob mismatch axes. + +`iter_operator_ablation_cases(module)` describes exactly four cases for one +module at a time: `P/P`, `R/R`, `P/R`, and `R/P`. `P` means the production +implementation and `R` means the RL-Kernel implementation, with the training +side written first. The case record includes only the matrix's stable axis IDs; +the detailed probe definitions, comparability gates, and tolerances remain in +RL-Kernel. + +This adapter is intentionally descriptive. It neither changes vime scheduling +nor claims that an operator is installed on a side where vime has no runtime +hook. The runner must record actual train/rollout provenance before treating a +case as a completed measurement. diff --git a/tests/test_rl_kernel_standard_adapter.py b/tests/test_rl_kernel_standard_adapter.py new file mode 100644 index 000000000..bb8c31b02 --- /dev/null +++ b/tests/test_rl_kernel_standard_adapter.py @@ -0,0 +1,210 @@ +import sys +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from enum import Enum +from types import ModuleType +from typing import Any + +import pytest +import torch + +from vime.backends.rl_kernel_utils import ( + RlkAlignmentStandard, + RlkAlignmentStandardUnavailable, + alignment_standard_metadata, + build_standard_score_artifact, + compare_standard_score_records, + iter_operator_ablation_cases, + load_alignment_standard, + load_operator_ablation_matrix, + select_operator_ablation_case, +) + + +def _matrix() -> dict[str, object]: + return { + "schema_version": "rlkernel.debug_matrix.v1", + "modules": { + "attention": {"axes": [{"id": "position_rope"}]}, + "ffn": {"axes": [{"id": "gemm_reduction"}]}, + "logp": {"axes": [{"id": "vocab_lse_reduction"}]}, + }, + } + + +@dataclass(frozen=True) +class _FakeStandard: + module_debug_matrix: Mapping[str, Any] = field(default_factory=_matrix) + compare_score_artifacts: Callable[..., Any] | None = None + schema_types: Mapping[str, Any] = field(default_factory=dict) + source: str = "rl_kernel.fake" + standard_id: str = "rl_kernel.cross_config.module_debug_matrix" + tolerance_fingerprint: str = "tolerance-sha" + + +@dataclass(frozen=True) +class _FakeIdentity: + checkpoint_id: str + model_version: str + + +class _FakeScoreSide(str, Enum): + ROLLOUT = "rollout" + TRAINING = "training" + + +@dataclass(frozen=True) +class _FakeScorer: + side: _FakeScoreSide + backend_id: str + dtype: str + + +@dataclass(frozen=True) +class _FakeProvenance: + requested: dict[str, object] + actual: dict[str, object] + + +@dataclass(frozen=True) +class _FakeScoreArtifact: + case_id: str + attempt_id: str + side: _FakeScoreSide + identity: _FakeIdentity + scorer: _FakeScorer + selected_logprobs: torch.Tensor + active_mask: torch.Tensor + provenance: _FakeProvenance + + +@pytest.mark.unit +def test_module_matrix_is_consumed_without_a_vime_copy(): + standard = _FakeStandard() + + assert load_operator_ablation_matrix(standard) is standard.module_debug_matrix + assert alignment_standard_metadata(standard)["alignment_matrix_schema_version"] == "rlkernel.debug_matrix.v1" + + cases = iter_operator_ablation_cases("attention", provider=standard) + + assert [(case.case_id, case.training_implementation, case.rollout_implementation) for case in cases] == [ + ("P/P", "production", "production"), + ("R/R", "rl_kernel", "rl_kernel"), + ("P/R", "production", "rl_kernel"), + ("R/P", "rl_kernel", "production"), + ] + assert all(case.diagnostic_axes == ("position_rope",) for case in cases) + assert select_operator_ablation_case("attention", "p/r", provider=standard).purpose == "rollout-only mismatch" + + +@pytest.mark.unit +def test_load_alignment_standard_reads_rl_kernel_public_contracts(monkeypatch): + compared = {} + + def compare(rollout, training): + compared["rollout"] = rollout + compared["training"] = training + return "compared-through-rl-kernel" + + rl_engine = ModuleType("rl_engine") + alignment = ModuleType("rl_engine.alignment") + kernels = ModuleType("rl_engine.kernels") + gtest = ModuleType("rl_engine.kernels.gtest") + cross_config = ModuleType("rl_engine.alignment.cross_config") + debug_matrix = ModuleType("rl_engine.alignment.cross_config.debug_matrix") + schema = ModuleType("rl_engine.alignment.cross_config.schema") + tolerance = ModuleType("rl_engine.kernels.gtest.tolerance") + cross_config.compare_score_artifacts = compare + debug_matrix.module_debug_matrix = _matrix + schema.RuntimeProvenance = _FakeProvenance + schema.ScoreArtifact = _FakeScoreArtifact + schema.ScorerSpec = _FakeScorer + schema.ScoreSide = _FakeScoreSide + schema.SemanticIdentitySpec = _FakeIdentity + tolerance.tolerance_contract_fingerprint = lambda: "tolerance-fingerprint" + + monkeypatch.setitem(sys.modules, "rl_engine", rl_engine) + monkeypatch.setitem(sys.modules, "rl_engine.alignment", alignment) + monkeypatch.setitem(sys.modules, "rl_engine.kernels", kernels) + monkeypatch.setitem(sys.modules, "rl_engine.kernels.gtest", gtest) + monkeypatch.setitem(sys.modules, "rl_engine.alignment.cross_config", cross_config) + monkeypatch.setitem(sys.modules, "rl_engine.alignment.cross_config.debug_matrix", debug_matrix) + monkeypatch.setitem(sys.modules, "rl_engine.alignment.cross_config.schema", schema) + monkeypatch.setitem(sys.modules, "rl_engine.kernels.gtest.tolerance", tolerance) + + standard = load_alignment_standard() + + assert isinstance(standard, RlkAlignmentStandard) + assert standard.compare_score_artifacts is compare + assert standard.schema_types["ScoreArtifact"] is _FakeScoreArtifact + assert alignment_standard_metadata(standard) == { + "alignment_standard_source": "rl_kernel", + "alignment_standard_id": "rl_kernel.cross_config.module_debug_matrix", + "alignment_matrix_schema_version": "rlkernel.debug_matrix.v1", + "alignment_standard_fingerprint": standard.fingerprint, + "alignment_tolerance_fingerprint": "tolerance-fingerprint", + } + + +@pytest.mark.unit +def test_invalid_or_unknown_matrix_cases_are_rejected(): + with pytest.raises(ValueError, match="unknown RL-Kernel ablation module"): + iter_operator_ablation_cases("unknown", provider=_FakeStandard()) + with pytest.raises(ValueError, match="unknown RL-Kernel ablation case"): + select_operator_ablation_case("logp", "P/R/R", provider=_FakeStandard()) + with pytest.raises(RlkAlignmentStandardUnavailable, match="module debug matrix"): + load_alignment_standard({"module_debug_matrix": {}}) + + +@pytest.mark.unit +def test_framework_score_records_compare_through_standard_comparator(): + compared = {} + + def compare(rollout, training): + compared["rollout"] = rollout + compared["training"] = training + return { + "case_id": rollout.case_id, + "mismatch_count": int( + (training.selected_logprobs[training.active_mask] != rollout.selected_logprobs[rollout.active_mask]) + .sum() + .item() + ), + } + + standard = _FakeStandard( + compare_score_artifacts=compare, + schema_types={ + "RuntimeProvenance": _FakeProvenance, + "ScoreArtifact": _FakeScoreArtifact, + "ScorerSpec": _FakeScorer, + "ScoreSide": _FakeScoreSide, + "SemanticIdentitySpec": _FakeIdentity, + }, + ) + base_record = { + "case_id": "case-1", + "attempt_id": "attempt-1", + "identity": {"checkpoint_id": "ckpt", "model_version": "weights"}, + "scorer": {"backend_id": "vime.native", "dtype": "float32"}, + "selected_logprobs": torch.tensor([0.0, -1.0, -2.0]), + "active_mask": torch.tensor([1, 0, 1], dtype=torch.bool), + "provenance": {"requested": {"backend": "rlk"}, "actual": {"backend": "rlk"}}, + } + + result = compare_standard_score_records( + {**base_record, "side": "rollout"}, + {**base_record, "side": "training", "selected_logprobs": torch.tensor([0.0, -9.0, -2.5])}, + standard=standard, + ) + + assert result == {"case_id": "case-1", "mismatch_count": 1} + assert compared["rollout"].side is _FakeScoreSide.ROLLOUT + assert compared["training"].scorer.side is _FakeScoreSide.TRAINING + assert compared["training"].active_mask.dtype is torch.bool + + +@pytest.mark.unit +def test_standard_score_artifact_export_requires_rl_kernel_schema(): + with pytest.raises(RlkAlignmentStandardUnavailable, match="score-artifact schema"): + build_standard_score_artifact({}, standard=_FakeStandard()) diff --git a/vime/backends/rl_kernel_utils/__init__.py b/vime/backends/rl_kernel_utils/__init__.py index af9ebcc13..3b8857e41 100644 --- a/vime/backends/rl_kernel_utils/__init__.py +++ b/vime/backends/rl_kernel_utils/__init__.py @@ -47,6 +47,20 @@ RlkOperatorComparisonUnavailable, load_operator_comparison_module, ) +from vime.backends.rl_kernel_utils.standard import ( + OperatorAblationCase, + RlkAlignmentStandard, + RlkAlignmentStandardUnavailable, + alignment_standard_metadata, + build_standard_score_artifact, + compare_standard_score_records, + iter_alignment_profiles, + iter_operator_ablation_cases, + load_alignment_standard, + load_operator_ablation_matrix, + select_least_restrictive_passing_profile, + select_operator_ablation_case, +) _OPERATOR_COMPARISON_EXPORTS = frozenset( { @@ -142,6 +156,8 @@ def __getattr__(name: str) -> Any: "runtime_batch_metadata_from_vime_batch", "RlkOperatorComparisonUnavailable", "load_operator_comparison_module", + "OperatorAblationCase", + "RlkAlignmentStandard", "BatchInvarianceCase", "ForwardChainComparisonResult", "ForwardChainStep", @@ -185,4 +201,14 @@ def __getattr__(name: str) -> Any: "run_deterministic_repeatability_check", "run_forward_chain_comparison", "run_reference_operator", + "RlkAlignmentStandardUnavailable", + "alignment_standard_metadata", + "build_standard_score_artifact", + "compare_standard_score_records", + "iter_alignment_profiles", + "iter_operator_ablation_cases", + "load_alignment_standard", + "load_operator_ablation_matrix", + "select_least_restrictive_passing_profile", + "select_operator_ablation_case", ] diff --git a/vime/backends/rl_kernel_utils/standard.py b/vime/backends/rl_kernel_utils/standard.py new file mode 100644 index 000000000..0c7a0ef78 --- /dev/null +++ b/vime/backends/rl_kernel_utils/standard.py @@ -0,0 +1,349 @@ +"""Thin access layer for RL-Kernel-owned alignment contracts. + +vime consumes the public score comparator and module mismatch matrix, but does +not redefine either contract. The matrix remains a fixed-replay diagnostic +manifest; this module only makes its module-level P/R cases available to vime +callers and external runners. +""" + +from __future__ import annotations + +import hashlib +import importlib +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + + +class RlkAlignmentStandardUnavailable(RuntimeError): + """Raised when RL-Kernel's public alignment contracts cannot be used.""" + + +@dataclass(frozen=True) +class RlkAlignmentStandard: + """Normalized view of the public RL-Kernel alignment contracts.""" + + module_debug_matrix: Mapping[str, Any] + compare_score_artifacts: Any + schema_types: Mapping[str, Any] + source: str = "rl_kernel" + standard_id: str = "rl_kernel.cross_config.module_debug_matrix" + tolerance_fingerprint: str = "" + + @property + def profile_version(self) -> str: + return str(self.module_debug_matrix["schema_version"]) + + @property + def fingerprint(self) -> str: + return _stable_fingerprint(self.module_debug_matrix) + + def to_metadata(self) -> dict[str, str]: + return { + "alignment_standard_source": self.source, + "alignment_standard_id": self.standard_id, + "alignment_matrix_schema_version": self.profile_version, + "alignment_standard_fingerprint": self.fingerprint, + "alignment_tolerance_fingerprint": self.tolerance_fingerprint, + } + + +@dataclass(frozen=True) +class OperatorAblationCase: + """One module-local production/RL-Kernel comparison case.""" + + module: str + case_id: str + training_implementation: str + rollout_implementation: str + purpose: str + diagnostic_axes: tuple[str, ...] + matrix_schema_version: str + + def to_dict(self) -> dict[str, Any]: + return { + "module": self.module, + "case_id": self.case_id, + "training_implementation": self.training_implementation, + "rollout_implementation": self.rollout_implementation, + "purpose": self.purpose, + "diagnostic_axes": self.diagnostic_axes, + "matrix_schema_version": self.matrix_schema_version, + } + + +_CASE_DEFINITIONS = ( + ("P/P", "production", "production", "native baseline"), + ("R/R", "rl_kernel", "rl_kernel", "RL-Kernel control"), + ("P/R", "production", "rl_kernel", "rollout-only mismatch"), + ("R/P", "rl_kernel", "production", "training-only mismatch"), +) + + +def load_alignment_standard(provider: Any = None) -> Any: + """Return the normalized RL-Kernel public alignment contracts.""" + + if provider is not None: + get_standard = getattr(provider, "get_alignment_standard", None) + standard = get_standard() if get_standard is not None else provider + if _standard_value(standard, "module_debug_matrix") is not None: + _require_module_debug_matrix(_standard_value(standard, "module_debug_matrix")) + elif not _has_legacy_profiles(standard): + raise RlkAlignmentStandardUnavailable("RL-Kernel did not provide a module debug matrix.") + return standard + + try: + cross_config = importlib.import_module("rl_engine.alignment.cross_config") + debug_matrix = importlib.import_module("rl_engine.alignment.cross_config.debug_matrix") + schema = importlib.import_module("rl_engine.alignment.cross_config.schema") + tolerance = importlib.import_module("rl_engine.kernels.gtest.tolerance") + except Exception as exc: + raise RlkAlignmentStandardUnavailable( + "RL-Kernel alignment contracts are unavailable; install RL-Kernel or pass an explicit test provider." + ) from exc + + module_debug_matrix = getattr(debug_matrix, "module_debug_matrix", None) + if module_debug_matrix is None: + raise RlkAlignmentStandardUnavailable("RL-Kernel does not expose module_debug_matrix().") + compare_score_artifacts = getattr(cross_config, "compare_score_artifacts", None) + if compare_score_artifacts is None: + raise RlkAlignmentStandardUnavailable("RL-Kernel does not expose compare_score_artifacts().") + + matrix = module_debug_matrix() + _require_module_debug_matrix(matrix) + schema_types = { + name: getattr(schema, name, None) + for name in ("RuntimeProvenance", "ScoreArtifact", "ScorerSpec", "ScoreSide", "SemanticIdentitySpec") + } + if not _has_score_schema(schema_types): + raise RlkAlignmentStandardUnavailable("RL-Kernel cross_config score-artifact schema is unavailable.") + tolerance_fingerprint = getattr(tolerance, "tolerance_contract_fingerprint", lambda: "")() + return RlkAlignmentStandard( + module_debug_matrix=matrix, + compare_score_artifacts=compare_score_artifacts, + schema_types=schema_types, + tolerance_fingerprint=str(tolerance_fingerprint), + ) + + +def load_operator_ablation_matrix(provider: Any = None) -> Mapping[str, Any]: + """Return RL-Kernel's module mismatch manifest without redefining it.""" + + standard = load_alignment_standard(provider) + matrix = _standard_value(standard, "module_debug_matrix") + _require_module_debug_matrix(matrix) + return matrix + + +def iter_operator_ablation_cases( + module: str, + *, + provider: Any = None, +) -> tuple[OperatorAblationCase, ...]: + """Build the four one-module P/R cases for a manifest module.""" + + matrix = load_operator_ablation_matrix(provider) + try: + module_manifest = matrix["modules"][module] + except (KeyError, TypeError) as exc: + raise ValueError(f"unknown RL-Kernel ablation module {module!r}") from exc + axes = tuple(str(axis["id"]) for axis in module_manifest.get("axes", ())) + return tuple( + OperatorAblationCase( + module=module, + case_id=case_id, + training_implementation=training_implementation, + rollout_implementation=rollout_implementation, + purpose=purpose, + diagnostic_axes=axes, + matrix_schema_version=str(matrix["schema_version"]), + ) + for case_id, training_implementation, rollout_implementation, purpose in _CASE_DEFINITIONS + ) + + +def select_operator_ablation_case( + module: str, + case_id: str, + *, + provider: Any = None, +) -> OperatorAblationCase: + """Select one fixed P/R case by its stable table label.""" + + normalized_case_id = case_id.strip().upper() + for case in iter_operator_ablation_cases(module, provider=provider): + if case.case_id == normalized_case_id: + return case + raise ValueError(f"unknown RL-Kernel ablation case {case_id!r}") + + +def iter_alignment_profiles(provider: Any = None) -> tuple[Any, ...]: + """Compatibility reader for an explicitly supplied legacy provider. + + The default RL-Kernel path has no profile catalog; callers should use + ``iter_operator_ablation_cases`` for the current module-level contract. + """ + + standard = load_alignment_standard(provider) + iterator = getattr(standard, "iter_profiles", None) + if iterator is not None: + return tuple(iterator()) + profiles = _standard_value(standard, "profiles") + if isinstance(profiles, Mapping): + return tuple(profiles.values()) + raise RlkAlignmentStandardUnavailable("RL-Kernel alignment profiles are not part of the current public contract.") + + +def select_least_restrictive_passing_profile( + profile_names: set[str], + *, + provider: Any = None, +) -> Any | None: + """Compatibility selector for an explicitly supplied legacy provider.""" + + for profile in reversed(iter_alignment_profiles(provider)): + name = profile.get("name") if isinstance(profile, Mapping) else getattr(profile, "name", None) + if str(name) in profile_names: + return profile + return None + + +def alignment_standard_metadata(standard: Any | None = None) -> dict[str, Any]: + """Return stable metadata that identifies the consumed RL-Kernel contract.""" + + standard = load_alignment_standard() if standard is None else standard + to_metadata = getattr(standard, "to_metadata", None) + if to_metadata is not None: + return {key: value for key, value in dict(to_metadata()).items() if value != ""} + + matrix = _standard_value(standard, "module_debug_matrix") + _require_module_debug_matrix(matrix) + return { + "alignment_standard_source": _standard_value(standard, "source", "rl_kernel"), + "alignment_standard_id": _standard_value(standard, "standard_id", "rl_kernel.cross_config.module_debug_matrix"), + "alignment_matrix_schema_version": str(matrix["schema_version"]), + "alignment_standard_fingerprint": _standard_value(standard, "fingerprint", _stable_fingerprint(matrix)), + "alignment_tolerance_fingerprint": _standard_value(standard, "tolerance_fingerprint", ""), + } + + +def build_standard_score_artifact( + record: Mapping[str, Any], + *, + standard: Any | None = None, +) -> Any: + """Convert a vime-owned score record into RL-Kernel's ScoreArtifact.""" + + standard = load_alignment_standard() if standard is None else standard + schema_types = _schema_types(standard) + if not _has_score_schema(schema_types): + raise RlkAlignmentStandardUnavailable("RL-Kernel cross_config score-artifact schema is unavailable.") + + score_side_type = schema_types["ScoreSide"] + side = record["side"] + side_value = side if isinstance(side, score_side_type) else score_side_type(str(side)) + identity = _coerce_schema_value(schema_types["SemanticIdentitySpec"], record["identity"]) + scorer_type = schema_types["ScorerSpec"] + scorer_value = record["scorer"] + if isinstance(scorer_value, scorer_type): + scorer = scorer_value + else: + scorer_record = dict(scorer_value) + scorer_record.setdefault("side", side_value) + scorer = _coerce_schema_value(scorer_type, scorer_record) + provenance = _coerce_schema_value(schema_types["RuntimeProvenance"], record["provenance"]) + return schema_types["ScoreArtifact"]( + case_id=record["case_id"], + attempt_id=record["attempt_id"], + side=side_value, + identity=identity, + scorer=scorer, + selected_logprobs=_tensor(record["selected_logprobs"]), + active_mask=_tensor(record["active_mask"]).to(dtype=torch.bool), + provenance=provenance, + ) + + +def compare_standard_score_records( + rollout: Mapping[str, Any], + training: Mapping[str, Any], + *, + standard: Any | None = None, +) -> Any: + """Compare vime-owned rollout/training score records through RL-Kernel.""" + + standard = load_alignment_standard() if standard is None else standard + compare_score_artifacts = _standard_value(standard, "compare_score_artifacts") + if compare_score_artifacts is None: + raise RlkAlignmentStandardUnavailable("RL-Kernel score comparator is unavailable.") + return compare_score_artifacts( + build_standard_score_artifact(rollout, standard=standard), + build_standard_score_artifact(training, standard=standard), + ) + + +def _require_module_debug_matrix(matrix: Any) -> None: + if not isinstance(matrix, Mapping) or not isinstance(matrix.get("schema_version"), str): + raise RlkAlignmentStandardUnavailable("RL-Kernel did not provide a valid module debug matrix.") + modules = matrix.get("modules") + if not isinstance(modules, Mapping) or not {"attention", "ffn", "logp"}.issubset(modules): + raise RlkAlignmentStandardUnavailable("RL-Kernel module debug matrix is missing attention, ffn, or logp.") + + +def _has_legacy_profiles(standard: Any) -> bool: + iterator = getattr(standard, "iter_profiles", None) + if iterator is not None: + return True + return isinstance(_standard_value(standard, "profiles"), Mapping) + + +def _schema_types(standard: Any) -> Mapping[str, Any]: + schema_types = _standard_value(standard, "schema_types", {}) + return schema_types if isinstance(schema_types, Mapping) else {} + + +def _has_score_schema(schema_types: Mapping[str, Any]) -> bool: + required = {"RuntimeProvenance", "ScoreArtifact", "ScorerSpec", "ScoreSide", "SemanticIdentitySpec"} + return required.issubset({key for key, value in schema_types.items() if value is not None}) + + +def _coerce_schema_value(schema_type: Any, value: Any) -> Any: + if isinstance(value, schema_type): + return value + if not isinstance(value, Mapping): + raise TypeError(f"expected mapping for {schema_type!r}, got {type(value)!r}.") + return schema_type(**dict(value)) + + +def _standard_value(standard: Any, name: str, default: Any = None) -> Any: + if isinstance(standard, Mapping): + return standard.get(name, default) + return getattr(standard, name, default) + + +def _stable_fingerprint(value: Mapping[str, Any]) -> str: + encoded = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _tensor(value: Any) -> torch.Tensor: + return value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + + +__all__ = [ + "OperatorAblationCase", + "RlkAlignmentStandard", + "RlkAlignmentStandardUnavailable", + "alignment_standard_metadata", + "build_standard_score_artifact", + "compare_standard_score_records", + "iter_alignment_profiles", + "iter_operator_ablation_cases", + "load_alignment_standard", + "load_operator_ablation_matrix", + "select_least_restrictive_passing_profile", + "select_operator_ablation_case", +]