integration: add TransformersProcessRewardModel with RoPE buffer repair - #215
Open
lambdabaa wants to merge 1 commit into
Open
integration: add TransformersProcessRewardModel with RoPE buffer repair#215lambdabaa wants to merge 1 commit into
lambdabaa wants to merge 1 commit into
Conversation
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 <extra_0> 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 <aria@caa.columbia.edu>
Contributor
|
@lambdabaa could we please rebase to main/v1 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this adds
TransformersProcessRewardModel— a new integration class forclassifier-head PRMs (Qwen2.5-Math-PRM-7B and compatible models) using
standard
transformers+torch, with no dependency on a running vLLMserver.
Before this PR,
its_hub's only in-process PRM integration wasLocalVllmProcessRewardModel, which requires a running vLLM server.This adds a second path: any machine with
transformersandtorchinstalled (MPS, CUDA, or CPU) can now load and score trajectories with
a classifier-head PRM in a single in-process call.
Bug fixed: RoPE buffers zeroed by transformers ≥5 meta-tensor init
Without the repair, loading Qwen2.5-Math-PRM-7B via
AutoModel.from_pretrained(..., trust_remote_code=True)with transformers≥5 produces a constant score of ~0.50003 for every step regardless of
input — the model appears to work but all scores are meaningless.
Root cause: transformers 5.x uses meta-tensor initialisation during
from_pretrained. Non-persistent buffers (inv_freq,cos_cached,sin_cached) onQwen2RotaryEmbeddingare materialised as zeros ratherthan being computed, and
load_state_dictskips them because they arenon-persistent. The result is an identity-collapsed rotary embedding that
projects every token to roughly the same position.
_repair_rotary_embeddings()detects and fixes this post-load byrecomputing
inv_freqfrommodule.baseandmodule.dimfor everyRotaryEmbeddingsubmodule. No other changes to model weights are made.Other implementation notes
<extra_0>separator:"<extra_0>".join(steps) + "<extra_0>"gives exactly N separator tokens for N steps, matching the model's training
format.
bfloat16default: avoids NaN hidden states that fp16 produces on MPSfor Qwen2ForProcessRewardModel's deep backbone.
use_cache=False: skipsDynamicCache.from_legacy_cachechurn intransformers 5.x for inference-only scoring.
Tests
Six tests added covering:
ImportErrorguard whentorch/transformersare absent_repair_rotary_embeddings()recomputes zeroedinv_freqbuffersscore()returnslen(steps)floats in[0, 1]<extra_0>separator is present (N steps → N separators)its_hub.integration