Llama3.3-70B-Instruct-uPRM-T80-adapters

Llama3.3-70B-Instruct-uPRM-T80-adapters

This is a process reward model (PRM) built on meta-llama/Llama-3.3-70B-Instruct. It assigns a two-class score to each reasoning step:

• class 0: + (the step is correct)
• class 1: - (the step is incorrect)

The repository contains a PEFT LoRA adapter, a trained special-token embedding, a process-reward head, tokenizer files, and custom Transformers loading code. It does not contain a second copy of the 70B base weights.

Built with Llama. Use of the base model and this derivative is subject to the Llama 3.3 Community License and Acceptable Use Policy. See the base model card for details.

Requirements

• Access to the gated meta-llama/Llama-3.3-70B-Instruct repository
• Authentication with a Hugging Face token that has access to the base model
• trust_remote_code=True
• A 140 GB-class GPU for the documented single-GPU path (designed for one NVIDIA H200)

The publishing environment used:

torch==2.8.0
transformers==4.57.1
accelerate==1.12.0
peft==0.17.1

FlashAttention 2 is recommended. Authenticate before loading, for example with hf auth login or by securely setting HF_TOKEN in the environment. Do not put access tokens directly in scripts.

Inference

Each reasoning step is represented as a user turn. The assistant content is the special marker <|*|>. The model produces logits for every input token, but only the logits at marker positions should be interpreted as step scores.

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer


MODEL_NAME = "agadetskii/Llama3.3-70B-Instruct-uPRM-T80-adapters"
SPECIAL_TOKEN = "<|*|>"
POSITIVE_CLASS = 0

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Single-H200 loading path. Do not call model.to("cuda") after loading with device_map.
model = AutoModel.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map={"": 0},
    attn_implementation="flash_attention_2",
)
model.eval()

# Put inputs on the device that owns the input embedding. This also works if a future setup uses
# a multi-device map where the first and last transformer layers are on different devices.
input_device = model.get_input_embeddings().weight.device

SYSTEM_PROMPT = """You are a strict mathematical reasoning judge.

Your task is to evaluate one individual reasoning step of a math problem at a time.

- If the step is mathematically correct, respond with `+`.
- If the step is mathematically incorrect or logically flawed, respond with `-`.
- Do not provide any explanation, comment, or feedback — only respond with `+` or `-`, and nothing else.
- Each input is either a single reasoning step or a new problem followed by its first reasoning step. In both cases, evaluate only the validity of the reasoning step.
- For each new problem, once you determine that a step is incorrect, you must consider all subsequent steps for that problem to also be incorrect, and respond with `-` for them as well.

Your response must only be one of these two symbols: `+` or `-`.
"""

problem = (
    "Sue's neighbors put 18 pink flamingos in her yard. On Saturday they took back one third, "
    "painted those white, and returned them. On Sunday they added another 18 pink flamingos. "
    "At noon on Sunday, how many more pink flamingos were there than white flamingos?"
)

steps = [
    "On Friday there were 18 pink flamingos.",
    "One third of 18 is 6, so after repainting there were 12 pink and 6 white flamingos.",
    "Adding 18 pink flamingos on Sunday gives 30 pink and 6 white flamingos.",
    "Therefore, there were 30 - 6 = 24 more pink flamingos than white flamingos.",
]

messages = [{"role": "system", "content": SYSTEM_PROMPT}]
for step_index, step in enumerate(steps):
    user_content = f"{problem} {step}" if step_index == 0 else step
    messages.append({"role": "user", "content": user_content})
    messages.append({"role": "assistant", "content": SPECIAL_TOKEN})

rendered = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=False,
)

# The Llama chat template already inserts BOS. add_special_tokens=False prevents a second BOS.
inputs = tokenizer(
    rendered,
    return_tensors="pt",
    add_special_tokens=False,
).to(input_device)

marker_id = tokenizer.convert_tokens_to_ids(SPECIAL_TOKEN)
assert tokenizer.encode(SPECIAL_TOKEN, add_special_tokens=False) == [marker_id]
marker_mask = inputs["input_ids"].eq(marker_id)
assert int(marker_mask.sum()) == len(steps)

with torch.inference_mode():
    outputs = model(**inputs)

probabilities = F.softmax(outputs.logits.float(), dim=-1)
step_rewards = probabilities[0, marker_mask[0], POSITIVE_CLASS].cpu().tolist()

for step_index, reward in enumerate(step_rewards, start=1):
    print(f"Step {step_index}: P(correct) = {reward:.6f}")

step_rewards contains one P(correct) value per reasoning step, in the original order. Higher values mean the model assigns more probability to the + class. These values should be treated as model scores rather than universally calibrated probabilities.

Important Llama tokenization detail

Do not use the following marker lookup:

tokenizer.encode("<|*|>")[0]

For this Llama tokenizer, the default encoding prepends BOS and returns [128000, 128256], so index 0 is BOS rather than the PRM marker. Use convert_tokens_to_ids or encode with add_special_tokens=False, as in the inference example.

Likewise, always use add_special_tokens=False when tokenizing a string already produced by apply_chat_template(..., tokenize=False).

Model structure

• Base model: Llama 3.3 70B Instruct
• Adapter: LoRA, rank 64, alpha 32, all linear transformer projections
• PRM special token: <|*|> (token ID 128256)
• Reward head: two-class MLP over the final hidden state
• Published model is scoring-only: the causal-language-model output head is deliberately removed during loading to fit inference on one H200
• The training-time critic is not included

Limitations

• The model is intended for mathematical reasoning-step assessment.
• A high score does not prove that a step is mathematically correct.
• Scores can depend on prompt wording, decomposition into steps, and preceding context.
• Single-GPU inference is very memory intensive. Start with a small batch size.
• The custom multi-GPU/device-map path has not been validated as extensively as the single-H200 path.
• This model cannot generate text because its language-model output head is intentionally omitted.

License and attribution

This model is a derivative of Llama 3.3. Llama 3.3 is licensed under the Llama 3.3 Community License, Copyright © Meta Platforms, Inc. All Rights Reserved. Users are responsible for complying with the applicable license and Acceptable Use Policy described on the base model page.