[Bug]: Reward model usage

[Django-Jiang]

Your current environment

I am recently trying to use the new feature to serve reward model with vLLM, but I note that the sequential classifier type RM is not supported well (I am using 0.7.1),

I checked #8976 #10444 seems it is been resolved already, I am not sure if it is a bug or not

🐛 Describe the bug

for example (adapted from this )

from vllm import LLM

llm = LLM(model="Skywork/Skywork-Reward-Llama-3.1-8B-v0.2", task="reward")
(output,) = llm.encode("Hello, my name is")

data = output.outputs.data
print(f"Data: {data!r}")

And I got ValueError: Model architectures ['LlamaForSequenceClassification'] are not supported for now. Similarly, Gemma2ForSequenceClassification is not supported. I wonder if there will be support for these models

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[git-xp]

fwiw, on vllm-0.7.2, running the same snippet:

from vllm import LLM

llm = LLM(model="Skywork/Skywork-Reward-Llama-3.1-8B-v0.2", task="reward")
(output,) = llm.encode("Hello, my name is")

data = output.outputs.data
print(f"Data: {data!r}")

I got a 6x4096 tensor instead of a single reward score. I think llm.encode() is pooling the hidden states of the last layer instead of computing the reward score, even we set task="reward".
Here is the output:

Processed prompts: 100%|██████████| 1/1 [00:00<00:00, 42.76it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Data: tensor([[ 1.2734,  2.6875,  1.6484,  ..., -1.8984,  2.7500,  1.8984],
        [-0.3281, -1.9141, -0.1025,  ...,  0.6562, -4.1875, -0.8789],
        [ 0.8711, -3.2500, -0.4277,  ..., -2.2031, -2.4062,  1.2891],
        [ 2.0781, -4.8125, -1.1797,  ..., -0.0869, -2.7500,  2.9688],
        [-2.9375, -1.3828, -0.7852,  ...,  1.4062, -2.5469, -0.5547],
        [-1.5938, -1.1406, -3.9219,  ...,  0.9062, -1.1562, -0.6484]],
       dtype=torch.float32)

[DarkLight1337]

For this particular model, seems that they are using a classifier architecture. So, I suggest setting task="classify"

[joanvelja]

Reward models are often based off of AutoModelForSequenceClassification (always, iirc). Though setting task='classify' is not hot it would be intended to work, given that the output that matters is the logit value, not the probit value.

Would it be possible for .classify() to instead return logits and probits instead of just the probits?

[DarkLight1337]

You can customize the pooling config via --override-pooler-config as documented here: https://docs.vllm.ai/en/latest/models/pooling_models.html

[joanvelja]

It is still unclear to me how to get reward modelling scores this way. Classify is not set up when setting type == reward, and the pooling docs is quite messy. Sorry for maybe misreading.

[DarkLight1337]

You can find more details in the CLI help here: https://docs.vllm.ai/en/latest/serving/engine_args.html#vllm.engine.arg_utils-_engine_args_parser-modelconfig

[joanvelja]

The docs clearly claim that the encode method simply returns the hidden states, and fwiw redirecting to different parts in the docs without a clear logic isn't helpful.

[DarkLight1337]

Sorry I didn't have bandwidth to go into much detail previously so I only pointed you to the relevant materials.

If you want to get the logits then you can use task="classify" and set --override-pooler-config {"softmax": false} . Based on the table, this should give you pooling type "LAST", no normalization and no softmax which corresponds to logits.

[joanvelja]

No worries, I was preparing an issue + suggested PR on how to handle the situation. Was able to do exactly what you suggested right now. Appreciate the patience!

[Tendo33]

hello, so is the only option available now to use the Transformer approach?
Here is the official example for Skywork-Reward-Llama-3.1-8B-v0.2d:

import torch

from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load model and tokenizer
device = "cuda:0"
model_name = "Skywork/Skywork-Reward-Llama-3.1-8B-v0.2"
rm = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map=device,
    attn_implementation="flash_attention_2",
    num_labels=1,
)
rm_tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "Jane has 12 apples. She gives 4 apples to her friend Mark, then buys 1 more apple, and finally splits all her apples equally among herself and her 2 siblings. How many apples does each person get?"
response1 = "1. Jane starts with 12 apples and gives 4 to Mark. 12 - 4 = 8. Jane now has 8 apples.\n2. Jane buys 1 more apple. 8 + 1 = 9. Jane now has 9 apples.\n3. Jane splits the 9 apples equally among herself and her 2 siblings (3 people in total). 9 ÷ 3 = 3 apples each. Each person gets 3 apples."
response2 = "1. Jane starts with 12 apples and gives 4 to Mark. 12 - 4 = 8. Jane now has 8 apples.\n2. Jane buys 1 more apple. 8 + 1 = 9. Jane now has 9 apples.\n3. Jane splits the 9 apples equally among her 2 siblings (2 people in total). 9 ÷ 2 = 4.5 apples each. Each person gets 4 apples."

conv1 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response1}]
conv2 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response2}]

# Format and tokenize the conversations
# If you use `tokenize=False` with `apply_chat_template` and `tokenizer()` to tokenize the conversation,
# remeber to remove the duplicated BOS token.
conv1_tokenized = rm_tokenizer.apply_chat_template(conv1, tokenize=True, return_tensors="pt").to(device)
conv2_tokenized = rm_tokenizer.apply_chat_template(conv2, tokenize=True, return_tensors="pt").to(device)

# Get the reward scores
with torch.no_grad():
    score1 = rm(conv1_tokenized).logits[0][0].item()
    score2 = rm(conv2_tokenized).logits[0][0].item()
print(f"Score for response 1: {score1}")
print(f"Score for response 2: {score2}")

# Output:
# 27B: 
# Score for response 1: 0.5625
# Score for response 2: -8.5

# 8B:
# Score for response 1: 13.6875
# Score for response 2: -9.1875

When I use vllm (0.8.5.post1)

 LlamaForSequenceClassification has no vLLM implementation, falling back to Transformers implementation. Some features may not be supported and performance may not be optimal.

[DarkLight1337]

Yes, that's correct. We don't support LlamaForSequenceClassification directly.

[stan-sony]

@Django-Jiang @Tendo33 @joanvelja @git-xp

You can simply do this, and it outputs the exact values (at least for demo examples)

from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "/mnt/mt_intern/HF_Models/Skywork-Reward-Llama-3.1-8B"
rm_tokenizer = AutoTokenizer.from_pretrained(model_name)
prompt = "Jane has 12 apples. She gives 4 apples to her friend Mark, then buys 1 more apple, and finally splits all her apples equally among herself and her 2 siblings. How many apples does each person get?"
response1 = "1. Jane starts with 12 apples and gives 4 to Mark. 12 - 4 = 8. Jane now has 8 apples.\n2. Jane buys 1 more apple. 8 + 1 = 9. Jane now has 9 apples.\n3. Jane splits the 9 apples equally among herself and her 2 siblings (3 people in total). 9 ÷ 3 = 3 apples each. Each person gets 3 apples."
response2 = "1. Jane starts with 12 apples and gives 4 to Mark. 12 - 4 = 8. Jane now has 8 apples.\n2. Jane buys 1 more apple. 8 + 1 = 9. Jane now has 9 apples.\n3. Jane splits the 9 apples equally among her 2 siblings (2 people in total). 9 ÷ 2 = 4.5 apples each. Each person gets 4 apples."

conv1 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response1}]
conv2 = [{"role": "user", "content": prompt}, {"role": "assistant", "content": response2}]

# Format and tokenize the conversations
conv1_formatted = rm_tokenizer.apply_chat_template(conv1, tokenize=False)
conv2_formatted = rm_tokenizer.apply_chat_template(conv2, tokenize=False)

conv1_tokenized = rm_tokenizer(conv1_formatted)
conv2_tokenized = rm_tokenizer(conv2_formatted)

from vllm import LLM
llm = LLM(model="/mnt/mt_intern/HF_Models/Skywork-Reward-Llama-3.1-8B", 
          task="classify",
          override_pooler_config={"softmax": False}  # This disables softmax for raw reward scores
          )
out1 = llm.encode(prompt_token_ids=conv1_tokenized['input_ids'])
out2 = llm.encode(prompt_token_ids=conv2_tokenized['input_ids'])
print('conv1', out1[0].outputs.data)
print('conv2', out2[0].outputs.data)
#conv1 tensor([12.6250])
#conv2 tensor([-15.2500])