From b32e34381e0986d209fc90a2783d3e3830f5c293 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Fri, 24 Jul 2026 20:46:52 +0800 Subject: [PATCH] Add single-card operator comparison harness Signed-off-by: inaniloquentee <3051000145@qq.com> --- tests/test_rl_kernel_operator_comparison.py | 119 ++++++++++++++++++ tests/utils/test_consistency_audit.py | 22 ++-- vime/backends/rl_kernel_utils/__init__.py | 109 ++++++++++++++++ .../rl_kernel_utils/operator_comparison.py | 84 +++++++++++++ vime/utils/consistency_audit.py | 14 +-- 5 files changed, 328 insertions(+), 20 deletions(-) create mode 100644 tests/test_rl_kernel_operator_comparison.py create mode 100644 vime/backends/rl_kernel_utils/operator_comparison.py diff --git a/tests/test_rl_kernel_operator_comparison.py b/tests/test_rl_kernel_operator_comparison.py new file mode 100644 index 00000000..f65880ca --- /dev/null +++ b/tests/test_rl_kernel_operator_comparison.py @@ -0,0 +1,119 @@ +import importlib +import sys +from types import ModuleType + +import pytest + + +def _drop_rl_engine_modules() -> None: + for name in list(sys.modules): + if name == "rl_engine" or name.startswith("rl_engine."): + sys.modules.pop(name, None) + + +def _install_fake_rl_kernel_operator_comparison(monkeypatch) -> ModuleType: + rl_engine = ModuleType("rl_engine") + alignment = ModuleType("rl_engine.alignment") + cross_config = ModuleType("rl_engine.alignment.cross_config") + operator_comparison = ModuleType("rl_engine.alignment.cross_config.operator_comparison") + + class FakeOperatorTolerance: + def __init__(self, *, atol=1e-6, rtol=1e-6): + self.atol = atol + self.rtol = rtol + + operator_comparison.RLK_OP_LOGP = "logp" + operator_comparison.PHASE4_TARGET_OPERATORS = ("logp",) + operator_comparison.OPERATOR_COMPARISON_SPECS = {"logp": "spec"} + operator_comparison.OperatorTolerance = FakeOperatorTolerance + operator_comparison.iter_operator_comparison_specs = lambda: ("spec",) + operator_comparison.compare_operator_outputs = lambda op_name, train, infer, **kwargs: { + "op_name": op_name, + "train": train, + "infer": infer, + "kwargs": kwargs, + } + + monkeypatch.setitem(sys.modules, "rl_engine", rl_engine) + monkeypatch.setitem(sys.modules, "rl_engine.alignment", alignment) + monkeypatch.setitem(sys.modules, "rl_engine.alignment.cross_config", cross_config) + monkeypatch.setitem( + sys.modules, + "rl_engine.alignment.cross_config.operator_comparison", + operator_comparison, + ) + return operator_comparison + + +@pytest.mark.unit +def test_operator_comparison_adapter_import_does_not_import_rl_engine(): + _drop_rl_engine_modules() + sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None) + + importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison") + + assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules) + + +@pytest.mark.unit +def test_operator_comparison_adapter_delegates_to_rl_kernel(monkeypatch): + _drop_rl_engine_modules() + sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None) + fake = _install_fake_rl_kernel_operator_comparison(monkeypatch) + adapter = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison") + + assert adapter.RLK_OP_LOGP == "logp" + assert adapter.PHASE4_TARGET_OPERATORS == ("logp",) + assert adapter.OPERATOR_COMPARISON_SPECS == {"logp": "spec"} + assert adapter.iter_operator_comparison_specs() == ("spec",) + assert adapter.OperatorTolerance(atol=0.5).atol == pytest.approx(0.5) + assert adapter.compare_operator_outputs("logp", "train", "infer") == { + "op_name": "logp", + "train": "train", + "infer": "infer", + "kwargs": {}, + } + assert adapter.load_operator_comparison_module() is fake + + +@pytest.mark.unit +def test_package_reexports_operator_comparison_lazily(monkeypatch): + _drop_rl_engine_modules() + sys.modules.pop("vime.backends.rl_kernel_utils", None) + sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None) + + package = importlib.import_module("vime.backends.rl_kernel_utils") + + assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules) + + _install_fake_rl_kernel_operator_comparison(monkeypatch) + + assert package.RLK_OP_LOGP == "logp" + assert package.iter_operator_comparison_specs() == ("spec",) + + +@pytest.mark.unit +def test_unavailable_operator_comparison_error_is_clear(monkeypatch): + _drop_rl_engine_modules() + sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None) + adapter = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison") + + def fail_import(name, *args, **kwargs): + if name == "rl_engine.alignment.cross_config.operator_comparison": + raise ModuleNotFoundError("No module named 'rl_engine'") + return importlib.import_module(name, *args, **kwargs) + + monkeypatch.setattr(adapter.importlib, "import_module", fail_import) + + with pytest.raises(adapter.RlkOperatorComparisonUnavailable, match="operator comparison standard"): + adapter.load_operator_comparison_module() + + +@pytest.mark.unit +def test_vime_operator_comparison_does_not_carry_reference_implementations(): + module = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison") + source = module.__loader__.get_source(module.__name__) + + assert "rl_engine.alignment.cross_config.operator_comparison" in source + assert "def reference_rmsnorm" not in source + assert "torch.nn.functional" not in source diff --git a/tests/utils/test_consistency_audit.py b/tests/utils/test_consistency_audit.py index 8f3edaba..8b36ff21 100644 --- a/tests/utils/test_consistency_audit.py +++ b/tests/utils/test_consistency_audit.py @@ -77,15 +77,19 @@ def _metadata(sample: Sample | None = None, *, args: Namespace | None = None, ** def _batch(*, metadata=None, layout=None): metadata = [_metadata()] if metadata is None else metadata - layout = [ - { - "fingerprint": "layout-1", - "active_mask_density": 1.0, - "dp_rank": 0, - "microbatch_id": 0, - "microbatch_offset": 0, - } - ] if layout is None else layout + layout = ( + [ + { + "fingerprint": "layout-1", + "active_mask_density": 1.0, + "dp_rank": 0, + "microbatch_id": 0, + "microbatch_offset": 0, + } + ] + if layout is None + else layout + ) return { "tokens": [torch.tensor([101, 201, 202])], "unconcat_tokens": [torch.tensor([101, 201, 202])], diff --git a/vime/backends/rl_kernel_utils/__init__.py b/vime/backends/rl_kernel_utils/__init__.py index fcb790c2..af9ebcc1 100644 --- a/vime/backends/rl_kernel_utils/__init__.py +++ b/vime/backends/rl_kernel_utils/__init__.py @@ -1,5 +1,8 @@ """RL-Kernel adapter and execution utilities owned by vime.""" +import importlib +from typing import Any + from vime.backends.rl_kernel_utils.adapter import ( RLK_ALL_OPERATORS, RLK_OP_LINEAR_LOGP, @@ -40,6 +43,67 @@ select_execution_decision, ) +from vime.backends.rl_kernel_utils.operator_comparison import ( + RlkOperatorComparisonUnavailable, + load_operator_comparison_module, +) + +_OPERATOR_COMPARISON_EXPORTS = frozenset( + { + "OPERATOR_COMPARISON_SPECS", + "PHASE4_TARGET_OPERATORS", + "RLK_OP_ATTENTION", + "RLK_OP_DPO_FRAGMENT", + "RLK_OP_EMBEDDING", + "RLK_OP_GRPO_FRAGMENT", + "RLK_OP_LM_HEAD", + "RLK_OP_LOGP", + "RLK_OP_MATMUL_PROJECTION", + "RLK_OP_PPO_FRAGMENT", + "RLK_OP_RATIO_KL", + "RLK_OP_RMSNORM", + "RLK_OP_ROPE", + "RLK_OP_SWIGLU", + "BatchInvarianceCase", + "ForwardChainComparisonResult", + "ForwardChainStep", + "OperatorComparisonResult", + "OperatorComparisonSpec", + "OperatorPair", + "OperatorTolerance", + "StrictBackendAdmissionReport", + "build_single_card_batch_invariance_cases", + "build_strict_backend_admission_report", + "compare_batch_invariance", + "compare_operator_outputs", + "compare_operator_pair", + "get_operator_comparison_spec", + "iter_operator_comparison_specs", + "reference_attention", + "reference_embedding", + "reference_linear_logp", + "reference_lm_head", + "reference_matmul_projection", + "reference_ppo_fragment", + "reference_ratio_kl", + "reference_rmsnorm", + "reference_rope", + "reference_selected_logprobs", + "reference_swiglu", + "run_deterministic_repeatability_check", + "run_forward_chain_comparison", + "run_reference_operator", + } +) + + +def __getattr__(name: str) -> Any: + if name not in _OPERATOR_COMPARISON_EXPORTS: + raise AttributeError(name) + module = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison") + return getattr(module, name) + + __all__ = [ "BackendCapability", "CapabilityQueryResult", @@ -76,4 +140,49 @@ "linear_logp_inputs_from_vime", "rlk_policy_context_from_args", "runtime_batch_metadata_from_vime_batch", + "RlkOperatorComparisonUnavailable", + "load_operator_comparison_module", + "BatchInvarianceCase", + "ForwardChainComparisonResult", + "ForwardChainStep", + "OPERATOR_COMPARISON_SPECS", + "OperatorComparisonResult", + "OperatorComparisonSpec", + "OperatorPair", + "OperatorTolerance", + "PHASE4_TARGET_OPERATORS", + "RLK_OP_ATTENTION", + "RLK_OP_DPO_FRAGMENT", + "RLK_OP_EMBEDDING", + "RLK_OP_GRPO_FRAGMENT", + "RLK_OP_LM_HEAD", + "RLK_OP_LOGP", + "RLK_OP_MATMUL_PROJECTION", + "RLK_OP_PPO_FRAGMENT", + "RLK_OP_RATIO_KL", + "RLK_OP_RMSNORM", + "RLK_OP_ROPE", + "RLK_OP_SWIGLU", + "StrictBackendAdmissionReport", + "build_single_card_batch_invariance_cases", + "build_strict_backend_admission_report", + "compare_batch_invariance", + "compare_operator_outputs", + "compare_operator_pair", + "get_operator_comparison_spec", + "iter_operator_comparison_specs", + "reference_attention", + "reference_embedding", + "reference_linear_logp", + "reference_lm_head", + "reference_matmul_projection", + "reference_ppo_fragment", + "reference_ratio_kl", + "reference_rmsnorm", + "reference_rope", + "reference_selected_logprobs", + "reference_swiglu", + "run_deterministic_repeatability_check", + "run_forward_chain_comparison", + "run_reference_operator", ] diff --git a/vime/backends/rl_kernel_utils/operator_comparison.py b/vime/backends/rl_kernel_utils/operator_comparison.py new file mode 100644 index 00000000..e8d1606a --- /dev/null +++ b/vime/backends/rl_kernel_utils/operator_comparison.py @@ -0,0 +1,84 @@ +"""Lazy re-export of RL-Kernel-owned operator comparison helpers.""" + +from __future__ import annotations + +import importlib +from typing import Any + +_RLK_OPERATOR_COMPARISON_MODULE = "rl_engine.alignment.cross_config.operator_comparison" + +_OPERATOR_COMPARISON_EXPORTS = frozenset( + { + "OPERATOR_COMPARISON_SPECS", + "PHASE4_TARGET_OPERATORS", + "RLK_OP_ATTENTION", + "RLK_OP_DPO_FRAGMENT", + "RLK_OP_EMBEDDING", + "RLK_OP_GRPO_FRAGMENT", + "RLK_OP_LM_HEAD", + "RLK_OP_LOGP", + "RLK_OP_MATMUL_PROJECTION", + "RLK_OP_PPO_FRAGMENT", + "RLK_OP_RATIO_KL", + "RLK_OP_RMSNORM", + "RLK_OP_ROPE", + "RLK_OP_SWIGLU", + "BatchInvarianceCase", + "ForwardChainComparisonResult", + "ForwardChainStep", + "OperatorComparisonResult", + "OperatorComparisonSpec", + "OperatorPair", + "OperatorTolerance", + "StrictBackendAdmissionReport", + "build_single_card_batch_invariance_cases", + "build_strict_backend_admission_report", + "compare_batch_invariance", + "compare_operator_outputs", + "compare_operator_pair", + "get_operator_comparison_spec", + "iter_operator_comparison_specs", + "reference_attention", + "reference_embedding", + "reference_linear_logp", + "reference_lm_head", + "reference_matmul_projection", + "reference_ppo_fragment", + "reference_ratio_kl", + "reference_rmsnorm", + "reference_rope", + "reference_selected_logprobs", + "reference_swiglu", + "run_deterministic_repeatability_check", + "run_forward_chain_comparison", + "run_reference_operator", + } +) + + +class RlkOperatorComparisonUnavailable(RuntimeError): + """Raised when RL-Kernel's operator comparison standard cannot be imported.""" + + +def load_operator_comparison_module() -> Any: + try: + return importlib.import_module(_RLK_OPERATOR_COMPARISON_MODULE) + except Exception as exc: + raise RlkOperatorComparisonUnavailable(f"RL-Kernel operator comparison standard is unavailable; install an RL-Kernel build that exposes {_RLK_OPERATOR_COMPARISON_MODULE!r}.") from exc + + +def __getattr__(name: str) -> Any: + if name not in _OPERATOR_COMPARISON_EXPORTS: + raise AttributeError(name) + return getattr(load_operator_comparison_module(), name) + + +def __dir__() -> list[str]: + return sorted((*globals(), *_OPERATOR_COMPARISON_EXPORTS)) + + +__all__ = [ + "RlkOperatorComparisonUnavailable", + "load_operator_comparison_module", + *_OPERATOR_COMPARISON_EXPORTS, +] diff --git a/vime/utils/consistency_audit.py b/vime/utils/consistency_audit.py index c7dec738..c34cdc2f 100644 --- a/vime/utils/consistency_audit.py +++ b/vime/utils/consistency_audit.py @@ -22,11 +22,7 @@ stable_fingerprint, validate_samples_consistency_metadata, ) -from vime.utils.dlogp_diagnostics import ( - DlogpAuditReport, - compute_dlogp_diagnostics, - get_rlk_consistency_mode, -) +from vime.utils.dlogp_diagnostics import DlogpAuditReport, compute_dlogp_diagnostics, get_rlk_consistency_mode CONSISTENCY_AUDIT_SCHEMA_VERSION = 1 AUDIT_REQUIRED_METADATA_FIELDS = ( @@ -289,8 +285,7 @@ def build_consistency_replay_manifest( "total_length": _int_or_none(_sequence_value(_batch_get(batch, "total_lengths"), position)), "response_length": response_length, "active_token_count": _active_token_count(loss_mask, response_length), - "has_rollout_log_probs": _batch_sequence_value(batch, "rollout_log_probs", position) - is not None, + "has_rollout_log_probs": _batch_sequence_value(batch, "rollout_log_probs", position) is not None, "consistency_metadata_fingerprint": None if record is None else record.get("fingerprint"), "batch_layout_fingerprint": _first_present( None if layout is None else layout.get("fingerprint"), @@ -545,10 +540,7 @@ def _with_audit_required_metadata_issues( warnings = list(validation.warnings) failures = list(validation.failures) - seen = { - (issue.code, issue.sample_index, issue.rollout_id, issue.field) - for issue in (*warnings, *failures) - } + seen = {(issue.code, issue.sample_index, issue.rollout_id, issue.field) for issue in (*warnings, *failures)} for position, record in enumerate(records): if record is None: