Skip to content

integration: add TransformersProcessRewardModel with RoPE buffer repair - #215

Open
lambdabaa wants to merge 1 commit into
Red-Hat-AI-Innovation-Team:v0from
lambdabaa:fix/transformers-prm-rope-repair
Open

integration: add TransformersProcessRewardModel with RoPE buffer repair#215
lambdabaa wants to merge 1 commit into
Red-Hat-AI-Innovation-Team:v0from
lambdabaa:fix/transformers-prm-rope-repair

Conversation

@lambdabaa

@lambdabaa lambdabaa commented Apr 26, 2026

Copy link
Copy Markdown

What this adds

TransformersProcessRewardModel — a new integration class for
classifier-head PRMs (Qwen2.5-Math-PRM-7B and compatible models) using
standard transformers + torch, with no dependency on a running vLLM
server.

Before this PR, its_hub's only in-process PRM integration was
LocalVllmProcessRewardModel, which requires a running vLLM server.
This adds a second path: any machine with transformers and torch
installed (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) on Qwen2RotaryEmbedding are materialised as zeros rather
than being computed, and load_state_dict skips them because they are
non-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 by
recomputing inv_freq from module.base and module.dim for every
RotaryEmbedding submodule. No other changes to model weights are made.

Other implementation notes

  • Trailing <extra_0> separator: "<extra_0>".join(steps) + "<extra_0>"
    gives exactly N separator tokens for N steps, matching the model's training
    format.
  • bfloat16 default: avoids NaN hidden states that fp16 produces on MPS
    for Qwen2ForProcessRewardModel's deep backbone.
  • use_cache=False: skips DynamicCache.from_legacy_cache churn in
    transformers 5.x for inference-only scoring.

Tests

Six tests added covering:

  • ImportError guard when torch/transformers are absent
  • _repair_rotary_embeddings() recomputes zeroed inv_freq buffers
  • score() returns len(steps) floats in [0, 1]
  • Trailing <extra_0> separator is present (N steps → N separators)
  • Class is importable from its_hub.integration

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>
@beatsmonster

Copy link
Copy Markdown
Contributor

@lambdabaa could we please rebase to main/v1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants