Skip to content

Commit 1649d82

Browse files
inaniloquenteeinaniloquentee
committed
Add single-card operator comparison harness
Signed-off-by: inaniloquentee <3051000145@qq.com>
1 parent dc5837d commit 1649d82

5 files changed

Lines changed: 328 additions & 20 deletions

File tree

‎tests/test_rl_kernel_operator_comparison.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import importlib
2+
import sys
3+
from types import ModuleType
4+
5+
import pytest
6+
7+
8+
def _drop_rl_engine_modules() -> None:
9+
for name in list(sys.modules):
10+
if name == "rl_engine" or name.startswith("rl_engine."):
11+
sys.modules.pop(name, None)
12+
13+
14+
def _install_fake_rl_kernel_operator_comparison(monkeypatch) -> ModuleType:
15+
rl_engine = ModuleType("rl_engine")
16+
alignment = ModuleType("rl_engine.alignment")
17+
cross_config = ModuleType("rl_engine.alignment.cross_config")
18+
operator_comparison = ModuleType("rl_engine.alignment.cross_config.operator_comparison")
19+
20+
class FakeOperatorTolerance:
21+
def __init__(self, *, atol=1e-6, rtol=1e-6):
22+
self.atol = atol
23+
self.rtol = rtol
24+
25+
operator_comparison.RLK_OP_LOGP = "logp"
26+
operator_comparison.PHASE4_TARGET_OPERATORS = ("logp",)
27+
operator_comparison.OPERATOR_COMPARISON_SPECS = {"logp": "spec"}
28+
operator_comparison.OperatorTolerance = FakeOperatorTolerance
29+
operator_comparison.iter_operator_comparison_specs = lambda: ("spec",)
30+
operator_comparison.compare_operator_outputs = lambda op_name, train, infer, **kwargs: {
31+
"op_name": op_name,
32+
"train": train,
33+
"infer": infer,
34+
"kwargs": kwargs,
35+
}
36+
37+
monkeypatch.setitem(sys.modules, "rl_engine", rl_engine)
38+
monkeypatch.setitem(sys.modules, "rl_engine.alignment", alignment)
39+
monkeypatch.setitem(sys.modules, "rl_engine.alignment.cross_config", cross_config)
40+
monkeypatch.setitem(
41+
sys.modules,
42+
"rl_engine.alignment.cross_config.operator_comparison",
43+
operator_comparison,
44+
)
45+
return operator_comparison
46+
47+
48+
@pytest.mark.unit
49+
def test_operator_comparison_adapter_import_does_not_import_rl_engine():
50+
_drop_rl_engine_modules()
51+
sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
52+
53+
importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
54+
55+
assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules)
56+
57+
58+
@pytest.mark.unit
59+
def test_operator_comparison_adapter_delegates_to_rl_kernel(monkeypatch):
60+
_drop_rl_engine_modules()
61+
sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
62+
fake = _install_fake_rl_kernel_operator_comparison(monkeypatch)
63+
adapter = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
64+
65+
assert adapter.RLK_OP_LOGP == "logp"
66+
assert adapter.PHASE4_TARGET_OPERATORS == ("logp",)
67+
assert adapter.OPERATOR_COMPARISON_SPECS == {"logp": "spec"}
68+
assert adapter.iter_operator_comparison_specs() == ("spec",)
69+
assert adapter.OperatorTolerance(atol=0.5).atol == pytest.approx(0.5)
70+
assert adapter.compare_operator_outputs("logp", "train", "infer") == {
71+
"op_name": "logp",
72+
"train": "train",
73+
"infer": "infer",
74+
"kwargs": {},
75+
}
76+
assert adapter.load_operator_comparison_module() is fake
77+
78+
79+
@pytest.mark.unit
80+
def test_package_reexports_operator_comparison_lazily(monkeypatch):
81+
_drop_rl_engine_modules()
82+
sys.modules.pop("vime.backends.rl_kernel_utils", None)
83+
sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
84+
85+
package = importlib.import_module("vime.backends.rl_kernel_utils")
86+
87+
assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules)
88+
89+
_install_fake_rl_kernel_operator_comparison(monkeypatch)
90+
91+
assert package.RLK_OP_LOGP == "logp"
92+
assert package.iter_operator_comparison_specs() == ("spec",)
93+
94+
95+
@pytest.mark.unit
96+
def test_unavailable_operator_comparison_error_is_clear(monkeypatch):
97+
_drop_rl_engine_modules()
98+
sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
99+
adapter = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
100+
101+
def fail_import(name, *args, **kwargs):
102+
if name == "rl_engine.alignment.cross_config.operator_comparison":
103+
raise ModuleNotFoundError("No module named 'rl_engine'")
104+
return importlib.import_module(name, *args, **kwargs)
105+
106+
monkeypatch.setattr(adapter.importlib, "import_module", fail_import)
107+
108+
with pytest.raises(adapter.RlkOperatorComparisonUnavailable, match="operator comparison standard"):
109+
adapter.load_operator_comparison_module()
110+
111+
112+
@pytest.mark.unit
113+
def test_vime_operator_comparison_does_not_carry_reference_implementations():
114+
module = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
115+
source = module.__loader__.get_source(module.__name__)
116+
117+
assert "rl_engine.alignment.cross_config.operator_comparison" in source
118+
assert "def reference_rmsnorm" not in source
119+
assert "torch.nn.functional" not in source

‎tests/utils/test_consistency_audit.py‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -77,15 +77,19 @@ def _metadata(sample: Sample | None = None, *, args: Namespace | None = None, **
7777

7878
def _batch(*, metadata=None, layout=None):
7979
metadata = [_metadata()] if metadata is None else metadata
80-
layout = [
81-
{
82-
"fingerprint": "layout-1",
83-
"active_mask_density": 1.0,
84-
"dp_rank": 0,
85-
"microbatch_id": 0,
86-
"microbatch_offset": 0,
87-
}
88-
] if layout is None else layout
80+
layout = (
81+
[
82+
{
83+
"fingerprint": "layout-1",
84+
"active_mask_density": 1.0,
85+
"dp_rank": 0,
86+
"microbatch_id": 0,
87+
"microbatch_offset": 0,
88+
}
89+
]
90+
if layout is None
91+
else layout
92+
)
8993
return {
9094
"tokens": [torch.tensor([101, 201, 202])],
9195
"unconcat_tokens": [torch.tensor([101, 201, 202])],

‎vime/backends/rl_kernel_utils/__init__.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
"""RL-Kernel adapter and execution utilities owned by vime."""
22

3+
import importlib
4+
from typing import Any
5+
36
from vime.backends.rl_kernel_utils.adapter import (
47
RLK_ALL_OPERATORS,
58
RLK_OP_LINEAR_LOGP,
@@ -40,6 +43,67 @@
4043
select_execution_decision,
4144
)
4245

46+
from vime.backends.rl_kernel_utils.operator_comparison import (
47+
RlkOperatorComparisonUnavailable,
48+
load_operator_comparison_module,
49+
)
50+
51+
_OPERATOR_COMPARISON_EXPORTS = frozenset(
52+
{
53+
"OPERATOR_COMPARISON_SPECS",
54+
"PHASE4_TARGET_OPERATORS",
55+
"RLK_OP_ATTENTION",
56+
"RLK_OP_DPO_FRAGMENT",
57+
"RLK_OP_EMBEDDING",
58+
"RLK_OP_GRPO_FRAGMENT",
59+
"RLK_OP_LM_HEAD",
60+
"RLK_OP_LOGP",
61+
"RLK_OP_MATMUL_PROJECTION",
62+
"RLK_OP_PPO_FRAGMENT",
63+
"RLK_OP_RATIO_KL",
64+
"RLK_OP_RMSNORM",
65+
"RLK_OP_ROPE",
66+
"RLK_OP_SWIGLU",
67+
"BatchInvarianceCase",
68+
"ForwardChainComparisonResult",
69+
"ForwardChainStep",
70+
"OperatorComparisonResult",
71+
"OperatorComparisonSpec",
72+
"OperatorPair",
73+
"OperatorTolerance",
74+
"StrictBackendAdmissionReport",
75+
"build_single_card_batch_invariance_cases",
76+
"build_strict_backend_admission_report",
77+
"compare_batch_invariance",
78+
"compare_operator_outputs",
79+
"compare_operator_pair",
80+
"get_operator_comparison_spec",
81+
"iter_operator_comparison_specs",
82+
"reference_attention",
83+
"reference_embedding",
84+
"reference_linear_logp",
85+
"reference_lm_head",
86+
"reference_matmul_projection",
87+
"reference_ppo_fragment",
88+
"reference_ratio_kl",
89+
"reference_rmsnorm",
90+
"reference_rope",
91+
"reference_selected_logprobs",
92+
"reference_swiglu",
93+
"run_deterministic_repeatability_check",
94+
"run_forward_chain_comparison",
95+
"run_reference_operator",
96+
}
97+
)
98+
99+
100+
def __getattr__(name: str) -> Any:
101+
if name not in _OPERATOR_COMPARISON_EXPORTS:
102+
raise AttributeError(name)
103+
module = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
104+
return getattr(module, name)
105+
106+
43107
__all__ = [
44108
"BackendCapability",
45109
"CapabilityQueryResult",
@@ -76,4 +140,49 @@
76140
"linear_logp_inputs_from_vime",
77141
"rlk_policy_context_from_args",
78142
"runtime_batch_metadata_from_vime_batch",
143+
"RlkOperatorComparisonUnavailable",
144+
"load_operator_comparison_module",
145+
"BatchInvarianceCase",
146+
"ForwardChainComparisonResult",
147+
"ForwardChainStep",
148+
"OPERATOR_COMPARISON_SPECS",
149+
"OperatorComparisonResult",
150+
"OperatorComparisonSpec",
151+
"OperatorPair",
152+
"OperatorTolerance",
153+
"PHASE4_TARGET_OPERATORS",
154+
"RLK_OP_ATTENTION",
155+
"RLK_OP_DPO_FRAGMENT",
156+
"RLK_OP_EMBEDDING",
157+
"RLK_OP_GRPO_FRAGMENT",
158+
"RLK_OP_LM_HEAD",
159+
"RLK_OP_LOGP",
160+
"RLK_OP_MATMUL_PROJECTION",
161+
"RLK_OP_PPO_FRAGMENT",
162+
"RLK_OP_RATIO_KL",
163+
"RLK_OP_RMSNORM",
164+
"RLK_OP_ROPE",
165+
"RLK_OP_SWIGLU",
166+
"StrictBackendAdmissionReport",
167+
"build_single_card_batch_invariance_cases",
168+
"build_strict_backend_admission_report",
169+
"compare_batch_invariance",
170+
"compare_operator_outputs",
171+
"compare_operator_pair",
172+
"get_operator_comparison_spec",
173+
"iter_operator_comparison_specs",
174+
"reference_attention",
175+
"reference_embedding",
176+
"reference_linear_logp",
177+
"reference_lm_head",
178+
"reference_matmul_projection",
179+
"reference_ppo_fragment",
180+
"reference_ratio_kl",
181+
"reference_rmsnorm",
182+
"reference_rope",
183+
"reference_selected_logprobs",
184+
"reference_swiglu",
185+
"run_deterministic_repeatability_check",
186+
"run_forward_chain_comparison",
187+
"run_reference_operator",
79188
]

‎vime/backends/rl_kernel_utils/operator_comparison.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
"""Lazy re-export of RL-Kernel-owned operator comparison helpers."""
2+
3+
from __future__ import annotations
4+
5+
import importlib
6+
from typing import Any
7+
8+
_RLK_OPERATOR_COMPARISON_MODULE = "rl_engine.alignment.cross_config.operator_comparison"
9+
10+
_OPERATOR_COMPARISON_EXPORTS = frozenset(
11+
{
12+
"OPERATOR_COMPARISON_SPECS",
13+
"PHASE4_TARGET_OPERATORS",
14+
"RLK_OP_ATTENTION",
15+
"RLK_OP_DPO_FRAGMENT",
16+
"RLK_OP_EMBEDDING",
17+
"RLK_OP_GRPO_FRAGMENT",
18+
"RLK_OP_LM_HEAD",
19+
"RLK_OP_LOGP",
20+
"RLK_OP_MATMUL_PROJECTION",
21+
"RLK_OP_PPO_FRAGMENT",
22+
"RLK_OP_RATIO_KL",
23+
"RLK_OP_RMSNORM",
24+
"RLK_OP_ROPE",
25+
"RLK_OP_SWIGLU",
26+
"BatchInvarianceCase",
27+
"ForwardChainComparisonResult",
28+
"ForwardChainStep",
29+
"OperatorComparisonResult",
30+
"OperatorComparisonSpec",
31+
"OperatorPair",
32+
"OperatorTolerance",
33+
"StrictBackendAdmissionReport",
34+
"build_single_card_batch_invariance_cases",
35+
"build_strict_backend_admission_report",
36+
"compare_batch_invariance",
37+
"compare_operator_outputs",
38+
"compare_operator_pair",
39+
"get_operator_comparison_spec",
40+
"iter_operator_comparison_specs",
41+
"reference_attention",
42+
"reference_embedding",
43+
"reference_linear_logp",
44+
"reference_lm_head",
45+
"reference_matmul_projection",
46+
"reference_ppo_fragment",
47+
"reference_ratio_kl",
48+
"reference_rmsnorm",
49+
"reference_rope",
50+
"reference_selected_logprobs",
51+
"reference_swiglu",
52+
"run_deterministic_repeatability_check",
53+
"run_forward_chain_comparison",
54+
"run_reference_operator",
55+
}
56+
)
57+
58+
59+
class RlkOperatorComparisonUnavailable(RuntimeError):
60+
"""Raised when RL-Kernel's operator comparison standard cannot be imported."""
61+
62+
63+
def load_operator_comparison_module() -> Any:
64+
try:
65+
return importlib.import_module(_RLK_OPERATOR_COMPARISON_MODULE)
66+
except Exception as exc:
67+
raise RlkOperatorComparisonUnavailable(f"RL-Kernel operator comparison standard is unavailable; install an RL-Kernel build that exposes {_RLK_OPERATOR_COMPARISON_MODULE!r}.") from exc
68+
69+
70+
def __getattr__(name: str) -> Any:
71+
if name not in _OPERATOR_COMPARISON_EXPORTS:
72+
raise AttributeError(name)
73+
return getattr(load_operator_comparison_module(), name)
74+
75+
76+
def __dir__() -> list[str]:
77+
return sorted((*globals(), *_OPERATOR_COMPARISON_EXPORTS))
78+
79+
80+
__all__ = [
81+
"RlkOperatorComparisonUnavailable",
82+
"load_operator_comparison_module",
83+
*_OPERATOR_COMPARISON_EXPORTS,
84+
]

‎vime/utils/consistency_audit.py‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,7 @@
2222
stable_fingerprint,
2323
validate_samples_consistency_metadata,
2424
)
25-
from vime.utils.dlogp_diagnostics import (
26-
DlogpAuditReport,
27-
compute_dlogp_diagnostics,
28-
get_rlk_consistency_mode,
29-
)
25+
from vime.utils.dlogp_diagnostics import DlogpAuditReport, compute_dlogp_diagnostics, get_rlk_consistency_mode
3026

3127
CONSISTENCY_AUDIT_SCHEMA_VERSION = 1
3228
AUDIT_REQUIRED_METADATA_FIELDS = (
@@ -289,8 +285,7 @@ def build_consistency_replay_manifest(
289285
"total_length": _int_or_none(_sequence_value(_batch_get(batch, "total_lengths"), position)),
290286
"response_length": response_length,
291287
"active_token_count": _active_token_count(loss_mask, response_length),
292-
"has_rollout_log_probs": _batch_sequence_value(batch, "rollout_log_probs", position)
293-
is not None,
288+
"has_rollout_log_probs": _batch_sequence_value(batch, "rollout_log_probs", position) is not None,
294289
"consistency_metadata_fingerprint": None if record is None else record.get("fingerprint"),
295290
"batch_layout_fingerprint": _first_present(
296291
None if layout is None else layout.get("fingerprint"),
@@ -545,10 +540,7 @@ def _with_audit_required_metadata_issues(
545540

546541
warnings = list(validation.warnings)
547542
failures = list(validation.failures)
548-
seen = {
549-
(issue.code, issue.sample_index, issue.rollout_id, issue.field)
550-
for issue in (*warnings, *failures)
551-
}
543+
seen = {(issue.code, issue.sample_index, issue.rollout_id, issue.field) for issue in (*warnings, *failures)}
552544

553545
for position, record in enumerate(records):
554546
if record is None:

0 commit comments

Comments
 (0)