From 611d0c8c2739c807aef17210f636552b177bfbcd Mon Sep 17 00:00:00 2001 From: lambdabaa Date: Sun, 26 Apr 2026 14:38:30 -0700 Subject: [PATCH] integration: add TransformersProcessRewardModel with RoPE buffer repair MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Transformers ≥5 uses meta-tensor initialisation during from_pretrained. Non-persistent buffers (inv_freq, cos_cached, sin_cached) on Qwen2RotaryEmbedding are zeroed out and skipped by load_state_dict, causing every forward pass to return a constant ~0.50003 score. _repair_rotary_embeddings() recomputes inv_freq from module.base and module.dim for every RotaryEmbedding submodule after loading, restoring correct per-step scores for Qwen2.5-Math-PRM-7B and compatible models. Also adds: - Classifier-head forward pass: logits[:, extra0_positions, 1] softmax - Trailing separator so N steps produce exactly N score tokens - bfloat16 default dtype (avoids NaN hidden states on MPS with fp16) - use_cache=False to skip DynamicCache.from_legacy_cache in transformers 5.x - Tests covering import guard, rotary repair, score interface, and export Signed-off-by: lambdabaa --- its_hub/integration/__init__.py | 5 + its_hub/integration/transformers_prm.py | 192 ++++++++++++++++ tests/test_transformers_prm.py | 291 ++++++++++++++++++++++++ 3 files changed, 488 insertions(+) create mode 100644 its_hub/integration/transformers_prm.py create mode 100644 tests/test_transformers_prm.py diff --git a/its_hub/integration/__init__.py b/its_hub/integration/__init__.py index e69de29b..2c557377 100644 --- a/its_hub/integration/__init__.py +++ b/its_hub/integration/__init__.py @@ -0,0 +1,5 @@ +from .transformers_prm import TransformersProcessRewardModel + +__all__ = [ + "TransformersProcessRewardModel", +] diff --git a/its_hub/integration/transformers_prm.py b/its_hub/integration/transformers_prm.py new file mode 100644 index 00000000..7a35b398 --- /dev/null +++ b/its_hub/integration/transformers_prm.py @@ -0,0 +1,192 @@ +"""Transformers-based Process Reward Model for classifier-head PRMs. + +Implements the correct scoring algorithm for Qwen2.5-Math-PRM-7B: + 1. Steps are joined with (the model's canonical step separator). + 2. One forward pass through the base transformer yields hidden states. + 3. model.score(hidden_states) applies the 2-class head → [batch, seq, 2]. + 4. Softmax(positive class) at each position = per-step score. + +Unlike MLXProcessRewardModel (which targets Math-Shepherd-style "+/-" token +scoring), this implementation handles classifier-head PRMs that cannot be +loaded by mlx_lm. +""" + +from __future__ import annotations + +import asyncio + +from its_hub.base import AbstractProcessRewardModel +from its_hub.types import ChatMessage, ChatMessages +from its_hub.utils import QWEN_SYSTEM_PROMPT + +_STEP_SEP_TOKEN = "" # Qwen2.5-Math-PRM canonical step separator + + +class TransformersProcessRewardModel(AbstractProcessRewardModel): + """Process Reward Model using transformers + MPS/CUDA for Qwen2.5-Math-PRM-7B. + + Input columns consumed by the companion sdg_hub block: ``problem`` (str) + and ``steps`` (``list[str]``). One forward pass returns one score per step. + + Args: + model_name: HuggingFace model path or local directory. + device: PyTorch device string. ``"mps"`` for Apple Silicon, + ``"cuda"`` for NVIDIA, ``"cpu"`` for CPU-only. + dtype: ``torch.dtype`` for model weights. Defaults to + ``torch.bfloat16``; bfloat16 shares float32's exponent range so + it avoids the NaN hidden states that fp16 produces on MPS for + Qwen2ForProcessRewardModel's deep backbone. + """ + + def __init__( + self, + model_name: str = "Qwen/Qwen2.5-Math-PRM-7B", + device: str = "mps", + dtype=None, + ): + try: + import torch + from transformers import AutoModel, AutoTokenizer + except ImportError as exc: + raise ImportError( + "TransformersProcessRewardModel requires transformers and torch. " + "Install with: pip install transformers torch" + ) from exc + + _dtype = dtype if dtype is not None else torch.bfloat16 + self._device = device + + self._tokenizer = AutoTokenizer.from_pretrained(model_name) + + # Qwen2RMConfig does not always carry pad_token_id from the JSON; + # load the config first and backfill it before constructing the model. + from transformers import AutoConfig + + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + if not hasattr(config, "pad_token_id") or config.pad_token_id is None: + config.pad_token_id = self._tokenizer.eos_token_id + + self._model = AutoModel.from_pretrained( + model_name, + config=config, + trust_remote_code=True, + torch_dtype=_dtype, + ).eval().to(device) + + # Transformers 5.x uses meta-tensor initialisation: non-persistent buffers + # (inv_freq, cos_cached, sin_cached) on Qwen2RotaryEmbedding are materialised + # as zeros rather than being recomputed. Force-recompute them now. + self._repair_rotary_embeddings() + + self._extra0_token_id: int = self._tokenizer.convert_tokens_to_ids(_STEP_SEP_TOKEN) + + # ------------------------------------------------------------------ + # Rotary embedding repair + # ------------------------------------------------------------------ + + def _repair_rotary_embeddings(self) -> None: + """Recompute inv_freq / cos_cached / sin_cached for every rotary layer. + + Transformers ≥5 uses meta-tensor initialisation during from_pretrained. + Non-persistent buffers (inv_freq, cos_cached, sin_cached) are skipped by + load_state_dict and end up as zeros. We recompute them from the module's + stored base / dim / max_seq_len_cached. + """ + import torch + + device = next(self._model.parameters()).device + repaired = 0 + for module in self._model.modules(): + if "RotaryEmbedding" not in type(module).__name__: + continue + if not (hasattr(module, "base") and hasattr(module, "dim") and hasattr(module, "inv_freq")): + continue + inv_freq = 1.0 / ( + module.base + ** (torch.arange(0, module.dim, 2, dtype=torch.int64).float().to(device) / module.dim) + ) + module.register_buffer("inv_freq", inv_freq, persistent=False) + if hasattr(module, "_set_cos_sin_cache"): + module._set_cos_sin_cache( + seq_len=module.max_seq_len_cached, + device=device, + dtype=torch.float32, + ) + repaired += 1 + + # ------------------------------------------------------------------ + # Core scoring + # ------------------------------------------------------------------ + + def _score_steps(self, prompt: str, steps: list[str]) -> list[float]: + """Single forward pass; returns one score per step.""" + import torch + + # Append a trailing separator so each step has exactly one token. + # "".join(steps) + "" produces N separators for N steps. + assistant_content = _STEP_SEP_TOKEN.join(steps) + _STEP_SEP_TOKEN + messages = [ + {"role": "system", "content": QWEN_SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + {"role": "assistant", "content": assistant_content}, + ] + text = self._tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=False + ) + enc = self._tokenizer(text, return_tensors="pt") + input_ids = enc["input_ids"].to(self._device) + attention_mask = enc.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(self._device) + + # Locate positions in the token sequence + extra0_positions = ( + (input_ids[0] == self._extra0_token_id).nonzero(as_tuple=True)[0] + ) + + with torch.no_grad(): + # Qwen2ForProcessRewardModel.forward() returns TokenClassifierOutput + # with logits of shape (1, seq, 2) already computed via the score head. + output = self._model( + input_ids=input_ids, + attention_mask=attention_mask, + use_cache=False, # avoids DynamicCache.from_legacy_cache in transformers 5.x + ) + logits = output.logits # (1, seq, 2) + + if len(extra0_positions) == 0: + # No step boundaries found — use last-token score as fallback + last_logits = logits[0, -1, :].float() + score = float(torch.softmax(last_logits, dim=-1)[1]) + return [score] * len(steps) + + scores: list[float] = [] + for pos in extra0_positions: + step_logits = logits[0, int(pos), :].float() + scores.append(float(torch.softmax(step_logits, dim=-1)[1])) + + # Align to step count in case of truncation or tokenisation quirks + while len(scores) < len(steps): + scores.append(scores[-1]) + return scores[: len(steps)] + + # ------------------------------------------------------------------ + # AbstractProcessRewardModel interface + # ------------------------------------------------------------------ + + def score( + self, + prompt_or_messages: str | list[ChatMessage] | ChatMessages, + steps: list[str], + ) -> list[float]: + chat_messages = ChatMessages.from_prompt_or_messages(prompt_or_messages) + return self._score_steps(chat_messages.to_prompt(), steps) + + async def ascore( + self, + prompt_or_messages: str | list[ChatMessage] | ChatMessages, + steps: list[str], + ) -> list[float]: + chat_messages = ChatMessages.from_prompt_or_messages(prompt_or_messages) + prompt = chat_messages.to_prompt() + return await asyncio.to_thread(self._score_steps, prompt, steps) diff --git a/tests/test_transformers_prm.py b/tests/test_transformers_prm.py new file mode 100644 index 00000000..7df483c9 --- /dev/null +++ b/tests/test_transformers_prm.py @@ -0,0 +1,291 @@ +"""Tests for TransformersProcessRewardModel. + +Covers: + - ImportError guard when torch/transformers are absent + - _repair_rotary_embeddings() recomputes zeroed inv_freq buffers + - score() returns a list of floats in [0, 1] of length == len(steps) + - Trailing separator is appended (N steps → N separators) + - ascore() delegates to score() via asyncio.to_thread +""" + +from __future__ import annotations + +import asyncio +import sys +import types +from unittest.mock import MagicMock, patch + +import pytest + + +# --------------------------------------------------------------------------- +# Helpers to build minimal fake torch / transformers stubs +# --------------------------------------------------------------------------- + + +def _make_torch_stub(): + """Minimal torch stub sufficient to import and instantiate the PRM.""" + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.float32 = "float32" + + # Tensor-like with basic arithmetic for inv_freq computation + class _Tensor: + def __init__(self, data): + self._data = list(data) if hasattr(data, "__iter__") else [data] + + def float(self): + return self + + def to(self, *args, **kwargs): + return self + + def __truediv__(self, other): + return self + + def __eq__(self, other): + class _Mask: + def nonzero(self, as_tuple=False): + return ([],) + + return _Mask() + + def __getitem__(self, idx): + return self + + def __len__(self): + return len(self._data) + + def arange(*args, dtype=None): + start, stop, step = 0, args[0], 1 + if len(args) == 3: + start, stop, step = args + return _Tensor(range(start, stop, step)) + + torch.arange = arange + torch.no_grad = MagicMock(return_value=MagicMock(__enter__=MagicMock(return_value=None), __exit__=MagicMock(return_value=False))) + + def softmax(tensor, dim=-1): + return _Tensor([0.3, 0.7]) + + torch.softmax = softmax + + nn = types.ModuleType("torch.nn") + torch.nn = nn + + return torch + + +def _make_transformers_stub(model_mock): + """Minimal transformers stub that returns `model_mock` from from_pretrained.""" + tf = types.ModuleType("transformers") + + tokenizer_mock = MagicMock() + tokenizer_mock.eos_token_id = 0 + tokenizer_mock.convert_tokens_to_ids.return_value = 151643 # + tokenizer_mock.apply_chat_template.return_value = "" + enc = MagicMock() + enc.__getitem__ = lambda self, k: MagicMock(to=lambda d: MagicMock()) + enc.get.return_value = None + tokenizer_mock.return_value = enc + tokenizer_mock.__call__ = lambda *a, **kw: enc + + config_mock = MagicMock() + config_mock.pad_token_id = None + + tf.AutoTokenizer = MagicMock() + tf.AutoTokenizer.from_pretrained.return_value = tokenizer_mock + tf.AutoModel = MagicMock() + tf.AutoModel.from_pretrained.return_value = model_mock + tf.AutoConfig = MagicMock() + tf.AutoConfig.from_pretrained.return_value = config_mock + + return tf, tokenizer_mock + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestImportGuard: + def test_raises_import_error_without_torch(self): + with patch.dict(sys.modules, {"torch": None, "transformers": None}): + # Force re-import + if "its_hub.integration.transformers_prm" in sys.modules: + del sys.modules["its_hub.integration.transformers_prm"] + from its_hub.integration.transformers_prm import TransformersProcessRewardModel + + with pytest.raises(ImportError, match="transformers and torch"): + TransformersProcessRewardModel() + + +class TestRepairRotaryEmbeddings: + def test_repairs_zeroed_inv_freq(self): + """_repair_rotary_embeddings should register fresh inv_freq on RotaryEmbedding modules.""" + import torch + + # Build a fake RotaryEmbedding module whose inv_freq is zeroed + class FakeRotaryEmbedding: + __name__ = "Qwen2RotaryEmbedding" + + def __init__(self): + self.base = 10000.0 + self.dim = 8 + self.inv_freq = torch.zeros(4) # zeroed — simulates the bug + self.max_seq_len_cached = 2048 + self._registered = {} + + def register_buffer(self, name, val, persistent=True): + self._registered[name] = val + + class FakeModel: + def __init__(self): + self._rope = FakeRotaryEmbedding() + + def parameters(self): + p = MagicMock() + p.device = torch.device("cpu") + return iter([p]) + + def modules(self): + return [self, self._rope] + + # Patch out the full construction; only test _repair_rotary_embeddings + rope = FakeRotaryEmbedding() + fake_model = FakeModel() + fake_model._rope = rope + + from its_hub.integration.transformers_prm import TransformersProcessRewardModel + + prm = object.__new__(TransformersProcessRewardModel) + prm._model = fake_model + prm._device = "cpu" + prm._repair_rotary_embeddings() + + assert "inv_freq" in rope._registered + repaired = rope._registered["inv_freq"] + # Should be non-zero (recomputed from base/dim) + assert repaired.abs().sum().item() > 0 + + +class TestScoreInterface: + @pytest.fixture + def prm(self): + """Construct a TransformersProcessRewardModel with all heavy deps mocked.""" + import torch + + # Build a model mock whose forward() returns logits with high score on pos 1 + output_mock = MagicMock() + # logits: shape (1, 5, 2) — positions 2 and 4 are tokens + logits = torch.zeros(1, 5, 2) + logits[0, :, 1] = 0.9 # positive class dominates everywhere + output_mock.logits = logits + + model_mock = MagicMock() + model_mock.return_value = output_mock + model_mock.eval.return_value = model_mock + model_mock.to.return_value = model_mock + + def fake_modules(): + return [] # no RotaryEmbedding modules → repair is a no-op + + model_mock.modules = fake_modules + + def fake_parameters(): + p = MagicMock() + p.device = torch.device("cpu") + return iter([p]) + + model_mock.parameters = fake_parameters + + tf, tokenizer_mock = _make_transformers_stub(model_mock) + + # Make tokenizer() call return input_ids with at positions 2,4 + token_ids = torch.tensor([[1, 2, 151643, 3, 151643]]) + enc = MagicMock() + enc.__getitem__ = lambda s, k: token_ids if k == "input_ids" else MagicMock() + enc.get.return_value = None + # MagicMock: use return_value (not __call__ assignment) so call dispatch works + tokenizer_mock.return_value = enc + tokenizer_mock.side_effect = lambda *a, **kw: enc + tf.AutoTokenizer.from_pretrained.return_value = tokenizer_mock + + with patch.dict( + sys.modules, + { + "torch": torch, + "transformers": tf, + }, + ): + if "its_hub.integration.transformers_prm" in sys.modules: + del sys.modules["its_hub.integration.transformers_prm"] + from its_hub.integration.transformers_prm import TransformersProcessRewardModel + + prm_instance = object.__new__(TransformersProcessRewardModel) + prm_instance._model = model_mock + prm_instance._device = "cpu" + prm_instance._tokenizer = tokenizer_mock + prm_instance._extra0_token_id = 151643 + yield prm_instance + + def test_score_returns_list_per_step(self, prm): + """score() must return exactly len(steps) floats.""" + import torch + + scores = prm._score_steps("What is 2+2?", ["Step 1: add", "Step 2: answer"]) + assert isinstance(scores, list) + assert len(scores) == 2 + assert all(isinstance(s, float) for s in scores) + + def test_scores_in_unit_interval(self, prm): + """All returned scores must be in [0, 1].""" + import torch + + scores = prm._score_steps("Solve x=1", ["x equals 1", "therefore x is 1"]) + assert all(0.0 <= s <= 1.0 for s in scores), f"Out-of-range scores: {scores}" + + def test_trailing_separator_in_assistant_content(self, prm): + """The assistant_content must end with , giving N separators for N steps.""" + captured = {} + + original_apply = prm._tokenizer.apply_chat_template + + def capturing_apply(messages, **kw): + captured["messages"] = messages + return "" + + prm._tokenizer.apply_chat_template = capturing_apply + + import torch + + prm._score_steps("problem", ["step A", "step B", "step C"]) + + assistant_msg = captured["messages"][-1] + assert assistant_msg["role"] == "assistant" + content = assistant_msg["content"] + + sep = "" + assert content.endswith(sep), f"Content should end with {sep!r}, got: {content!r}" + assert content.count(sep) == 3, f"Expected 3 separators for 3 steps, got: {content.count(sep)}" + + +class TestExport: + def test_exported_from_integration_package(self): + """TransformersProcessRewardModel must be importable from its_hub.integration.""" + # This test verifies the __init__.py export without loading torch/transformers + import importlib + import types + + # Stub heavy deps so the module-level import succeeds + fake_torch = types.ModuleType("torch") + fake_torch.bfloat16 = "bfloat16" + fake_tf = types.ModuleType("transformers") + + with patch.dict(sys.modules, {"torch": fake_torch, "transformers": fake_tf}): + for mod in list(sys.modules): + if "its_hub.integration" in mod: + del sys.modules[mod] + from its_hub.integration import TransformersProcessRewardModel # noqa: F401 + + assert TransformersProcessRewardModel is not None