[Model] Support math-shepherd-mistral-7b-prm model

[Went-Liang]

FILL IN THE PR DESCRIPTION HERE

Support peiyi9979/math-shepherd-mistral-7b-prm as embedding model.

As mentioned by 9314, the Process-Supervised Reward Model, which provides reward scores for intermediate steps generated by LLMs, can offer more fine-grained optimization for Reinforcement Learning (RL). This will help the community reproduce the OpenAI O1 model. PR 9424 allows any model that adds a pooler method to be used as an embedding model.

Therefore, this PR adds a pooler method to LlamaForCausalLM, introduces a pooling-type named "STEP" and adds a PoolerConfig class to facilitate users to configure the pooler method. In STEP mode, users can use the peiyi9979/math-shepherd-mistral-7b-prm model by setting the pooling-step-tag-id and pooling-returned-token-ids variables. pooling-returned-token-ids represents a list of indices for the vocabulary dimensions to be extracted, such as the token IDs of good_token and bad_token in the math-shepherd-mistral-7b-prm model. When pooling-step-tag-id is not None, it indicates that the score corresponding to the pooling-step-tag-id in the generated sentence should be returned. Otherwise, it returns the scores for all tokens.

The model can be served with:

python -m vllm.entrypoints.openai.api_server \
	--model peiyi9979/math-shepherd-mistral-7b-prm \
	--trust-remote-code \
	--served-model-name math-shepherd-mistral-7b-prm \
	--port 8080 \
	--tensor-parallel-size 1 \
	--enforce-eager \
        --task embedding \
	--pooling-type STEP \
	--pooling-step-tag-id 12902 \
	--pooling-returned-token-ids 648 387

And a test correspond to the example in the huggingface model page is:

import torch
from openai import OpenAI

question = """Janet\u2019s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?"""
output1 = """Step 1: Janet's ducks lay 16 eggs per day. ки\nStep 2: She eats three for breakfast every morning, so she has 16 - 3 = 13 eggs left. ки\nStep 3: She bakes muffins for her friends every day with four eggs, so she has 13 - 4 = 9 eggs left. ки\nStep 4: She sells the remainder at the farmers' market daily for $2 per fresh duck egg, so she makes 9 * $2 = $18 every day at the farmers' market. The answer is: 18 ки""" # 18 is right
output2 = """Step 1: Janet's ducks lay 16 eggs per day. ки\nStep 2: She eats three for breakfast every morning, so she has 16 - 3 = 13 eggs left. ки\nStep 3: She bakes muffins for her friends every day with four eggs, so she has 13 - 4 = 9 eggs left. ки\nStep 4: She sells the remainder at the farmers' market daily for $2 per fresh duck egg, so she makes 9 * $2 = $17 every day at the farmers' market. The answer is: 17 ки""" # 17 is wrong
full_prompt = [f"{question} {output}" for output in [output1, output2]]


# Modify OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8080/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

models = client.models.list()
model = models.data[0].id

responses = client.embeddings.create(
    input=full_prompt,
    model=model,
)

for data in responses.data:
    print(torch.tensor(data.embedding).view(-1, 2)[:, 0])
# tensor([0.9956, 0.9956, 0.9985, 0.9956])
# tensor([0.9956, 0.9956, 0.9985, 0.0240])

Of course, you can also use it directly like this:

from vllm import LLM

llm = LLM(
    model='peiyi9979/math-shepherd-mistral-7b-prm', 
    tensor_parallel_size=1, 
    task="embedding",
    pooling_type="STEP",
    pooling_step_tag_id=12902,
    pooling_returned_token_ids=[648, 387],
)

question = """Janet\u2019s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?"""
output1 = """Step 1: Janet's ducks lay 16 eggs per day. ки\nStep 2: She eats three for breakfast every morning, so she has 16 - 3 = 13 eggs left. ки\nStep 3: She bakes muffins for her friends every day with four eggs, so she has 13 - 4 = 9 eggs left. ки\nStep 4: She sells the remainder at the farmers' market daily for $2 per fresh duck egg, so she makes 9 * $2 = $18 every day at the farmers' market. The answer is: 18 ки""" # 18 is right
output2 = """Step 1: Janet's ducks lay 16 eggs per day. ки\nStep 2: She eats three for breakfast every morning, so she has 16 - 3 = 13 eggs left. ки\nStep 3: She bakes muffins for her friends every day with four eggs, so she has 13 - 4 = 9 eggs left. ки\nStep 4: She sells the remainder at the farmers' market daily for $2 per fresh duck egg, so she makes 9 * $2 = $17 every day at the farmers' market. The answer is: 17 ки""" # 17 is wrong

good_token = '+'
bad_token = '-'
step_tag = 'ки'

tokenizer = llm.get_tokenizer()
[good_token_id, bad_token_id, step_tag_id] = tokenizer.encode(f"{good_token} {bad_token} {step_tag}")[1:] # [648, 387, 12902]


full_prompt = [f"{question} {output}" for output in [output1, output2]]
outputs = llm.encode(full_prompt)

Thank you for your time on reviewing this PR :)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[DarkLight1337]

Can we add a factory method to Pooler to automatically merge the config with model specific defaults?

[DarkLight1337]

e.g. we should be able to write

self._pooler = Pooler.from_config_with_defaults(
    pooler_config,
    # These values are overridden if they are set inside the config
    pooling_type=PoolingType.CLS,
    normalize=True,
)

[Went-Liang]

Please check it

[DarkLight1337]

Looks good now, thanks for your effort and patience!

[DarkLight1337]

Now you just have to get the tests to pass.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. @Went-Liang please rebase it. https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Went-Liang]

Now you just have to get the tests to pass.

Excuse me, the test produced the following error (as shown in the image). This doesn't seem to be caused by my code changes. Could you please advise on how to handle this? @DarkLight1337

[DarkLight1337]

I have retried the failing test, see if it passes this time.

[DarkLight1337]

Looks like this issue comes from main branch, I have asked those with permissions to force-merge this.

[Went-Liang]

Looks like this issue comes from main branch, I have asked those with permissions to force-merge this.

Thanks so much !!!