[Usage]: How to infer the reward model of nvidia/Llama-3.1-Nemotron-70B-Reward-HF?

[xs1997zju]

Your current environment

The output of `python collect_env.py`

python3.9

How would you like to use vllm

I want to run inference of a [specific model](put link here). I don't know how to integrate it with vllm.

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.

[xs1997zju]

here is my infer scripts:
`from vllm import LLM, SamplingParams
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = 'nvidia/Llama-3.1-Nemotron-70B-Reward-HF'

model = LLM(model=model_path,
tensor_parallel_size=8,
dtype="float16",
enable_chunked_prefill=False,
task="embed",
)

tokenizer = AutoTokenizer.from_pretrained(model_path)

prompt = "What is 1+1?"

good_response = "1+1=2"

bad_response = "1+1=3"

messages = [
{'role': "user", "content": prompt},
{'role': "assistant", "content": good_response}
]

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

output = model.encode(prompt)

output[0].output.data.size(): [8192]`

but i want get the reward score, which should be: reward for good_response = -3.28125, according to the HF example

[DarkLight1337]

See #8967 (comment)

[DarkLight1337]

Please try out #11469. You should be able to run the model just by setting --task reward.

[nzw0301]

Hi, in my understanding, nvidia/Llama-3.1-Nemotron-70B-Reward-HF is implemented as a causal LM, and its reward computation uses lm_lead Thus as_reward_model via --task reward does not work for it, since lm_head is aways removed

vllm/vllm/model_executor/models/adapters.py

Lines 56 to 59 in 9ddac56

	# These are not used in pooling models
	for attr in ("lm_head", "logits_processor"):
	if hasattr(self, attr):
	delattr(self, attr)

. I suppose one workaround is to convert this causal model into a reward model explicitly.

[DarkLight1337]

Looking deeper into this, that alone isn't quite enough to solve the problem because lm_head is skipped in pooling task even if we load it in the model. Returning logits is not really supported right now, perhaps you can try to follow #11397

[ControllableGeneration]

Try this code snippet to get logits:
logits = model.lm_head.linear_method.apply(model.lm_head, hidden_states)

I have verified that the logits here is semi-identical to those obtained by huggingface transformers