upload training script and reward server script

reward_server/launch_reward.sh

@@ -0,0 +1,38 @@
+set -x
+
+MODEL_PATH=$1
+ANSWER_PATH=$2
+METRIC=$3
+PORT=8800
+
+export VLLM_ENGINE_ITERATION_TIMEOUT_S=60
+nohup vllm serve ${MODEL_PATH} \
+--trust-remote-code \
+--served-model-name server_model \
+--max-num-seqs 256 \
+--max-model-len 4096 \
+--port 8000 \
+> vllm_server.log &
+
+# sleep 60
+
+if [[ "${METRIC}" == "prob" ]]; then
+    nohup python model_server.py \
+--tokenizer_path ${MODEL_PATH} \
+--answer_path ${ANSWER_PATH} \
+--normalize_reward \
+--port ${PORT} \
+--prob_reward \
+--vllm_url "http://localhost:8000/v1" \
+--vllm_model server_model \
+> reward_server.log &
+else
+    nohup python model_server.py \
+--tokenizer_path ${MODEL_PATH} \
+--answer_path ${ANSWER_PATH} \
+--normalize_reward \
+--port ${PORT} \
+--vllm_url "http://localhost:8000/v1" \
+--vllm_model server_model \
+> reward_server.log &
+fi

reward_server/model_server.py

@@ -0,0 +1,252 @@
+import argparse
+import re
+import torch
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from transformers import AutoTokenizer
+import asyncio
+from collections import defaultdict
+import json
+from openai import AsyncOpenAI
+import time
+import math
+# Set OpenAI's API key and API base to use vLLM's API server.
+
+# for free-form including multiple-choice
+PROMPT_critic_updated = '''
+Given a problem, determine whether the final answer in the provided (incomplete) solution process matches the reference answer.  
+The reference answer may be one single option character (e.g., A, B, C, D), a numerical value, an expression, or a list of answers if multiple questions are involved.  
+**The reference answer may be in Chinese or another language, but your evaluation should be language-agnostic.**  
+
+Your task:  
+- Compare the final output of the solution process with the reference answer.  
+- If they **match exactly**, output **YES**.  
+- If they **do not match**, output **NO**.  
+- If the solution process is unclear, incomplete, or ambiguous, assume it is incorrect and output **NO**.  
+
+Your output must be strictly **'YES'** or **'NO'**, with no additional words, punctuation, or explanation.  
+
+---
+
+**Question:**  
+{question}  
+
+**Solution Process (Final Step Only):**  
+{response}  
+
+**Reference Answer:**  
+{reference}  
+
+**Output:**  
+'''
+
+
+
+def parse_im_sections(text):
+    # Match all sections between <|im_start|> and <|im_end|>
+    sections = re.findall(r"<\|im_start\|>(.*?)<\|im_end\|>", text, re.DOTALL)
+    parsed = {}
+    for section in sections:
+        try:
+            # Split the role and content
+            role, content = section.split("\n", 1)
+            parsed[role.strip()] = content.strip()
+        except ValueError:
+            print(f"Skipping malformed section: {section}")
+    return parsed
+
+def extract_last_non_empty_line(text, role="assistant"):
+    # Extract the last non-empty line from assistant's content
+    pattern = fr"<\|im_start\|>{role}(.*?)(?:<\|im_start\|>|<\|endoftext\|>|<\|eot_id\|>|$)"
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        content = match.group(1).strip()
+        # Get the last non-empty line
+        lines = [line for line in content.splitlines() if line.strip()]
+        if lines:
+            last_non_empty_line=lines[-1]
+        else:
+            return ""
+        return last_non_empty_line
+    return ""
+
+
+def reward_normalization(rewards):
+    if len(rewards) == 1:
+        return [0.0]
+    rewards = torch.tensor(rewards, dtype=torch.float64)
+    if rewards.std() == 0:
+        normalized_rewards = torch.zeros_like(rewards)
+    else:
+        normalized_rewards = (rewards - rewards.mean()) / rewards.std()
+
+    return normalized_rewards.tolist()
+
+
+def strip_sequence(text, pad_token, eos_token):
+    pad_token_escaped = re.escape(pad_token)
+    eos_token_escaped = re.escape(eos_token)
+
+    pattern = f"^({eos_token_escaped}|{pad_token_escaped})+"
+    text = re.sub(pattern, "", text)
+
+    pattern = f"({eos_token_escaped}|{pad_token_escaped})+$"
+    text = re.sub(pattern, "", text)
+    return text
+
+
+def group_reward_normalization(rewards, n_samples_per_prompt=4):
+    rewards = torch.tensor(rewards, dtype=torch.float64)
+    rewards = rewards.reshape(-1, n_samples_per_prompt)
+
+    mean = rewards.mean(dim=-1, keepdim=True)
+    std = rewards.std(dim=-1, keepdim=True)
+
+    normalized_rewards = torch.where(std == 0, torch.zeros_like(rewards), (rewards - mean) / std)
+
+    return normalized_rewards.flatten().tolist()
+
+
+class RewardModelProxy:
+    def __init__(self, args):
+        self.tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+        self.normalize_reward = args.normalize_reward
+        self.group_normalize_reward = args.group_normalize_reward
+        self.qa_dict = defaultdict(str)
+        self.load_dict(args.answer_path)
+        self.temperature = 0
+        self.stop=[self.tokenizer.eos_token,"<|im_end|>"]
+        self.max_tokens=1
+        self.prob_reward=args.prob_reward
+        self.log_path=args.log_path
+        self.vllm_model=args.vllm_model
+
+    def load_dict(self, path):
+        # Initialize self.qa_dict
+        with open(path, "r", encoding="utf-8") as file:
+            data = json.load(file)
+            for unit in data:
+                question = unit["query"][1]["content"]
+                label = unit["label"]
+                self.qa_dict[question] = label
+
+        if self.qa_dict:
+            sample_question, sample_label = next(iter(self.qa_dict.items()))
+            print("Sample Question:", sample_question)
+            print("Sample Label:", sample_label)
+        else:
+            print("qa_dict is empty.")
+
+
+    async def process_sample(self,query):
+        query = strip_sequence(query, self.tokenizer.pad_token, self.tokenizer.eos_token)+ self.tokenizer.eos_token
+        question = parse_im_sections(query)["user"]
+        answer = extract_last_non_empty_line(query, role="assistant")
+        if not answer.strip():
+            return 0.0
+        else:
+            prompt_question = PROMPT_critic_updated.format(question=question, reference=self.qa_dict[question], response=answer)
+            return await self.get_reward_from_vllm(prompt_question)
+
+    async def get_reward_from_vllm(self, query):
+        """Retrieve model judgment reward (with probability analysis)"""
+        max_retries = 10
+        delay=10
+        for attempt in range(max_retries):
+            try:
+                response = await client.chat.completions.create(
+                    model=self.vllm_model,
+                    messages=[
+                        {"role": "system", "content": "You are a helpful assistant."},
+                        {"role": "user", "content": query},
+                    ],
+                    temperature=self.temperature,
+                    max_tokens=self.max_tokens,
+                    stop=self.stop,
+                    logprobs=True,
+                    top_logprobs=10  # Get top 10 token probabilities
+                )
+                return self.calculate_reward_from_logprobs(response)
+                
+            except Exception as e:
+                print(f"Attempt {attempt+1} failed: {str(e)}, retrying in {delay} seconds...")
+                await asyncio.sleep(delay)
+        print(f"Failed after {max_retries} retries, query content: {query[:200]}...")
+        return 0.0  # Return baseline value on failure
+
+    def calculate_reward_from_logprobs(self, response):
+        """Calculate normalized reward based on log probabilities"""
+        # Extract probabilities of all possible tokens
+        logprobs = response.choices[0].logprobs.content[0].top_logprobs
+        token_probs = {token.token: math.exp(token.logprob) for token in logprobs}
+        
+        # Combine probabilities of YES/NO (case-insensitive)
+        yes_prob = sum(prob for token, prob in token_probs.items() if token.lower().strip()=="yes")
+        no_prob = sum(prob for token, prob in token_probs.items()if token.lower().strip()=="no")
+        total = yes_prob + no_prob
+        if total == 0:
+            return 0.0  # Return baseline value when no valid judgment
+        if self.prob_reward:
+            print(yes_prob/total)
+            return yes_prob / total  # Normalized probability
+        return 1.0 if yes_prob > no_prob else 0.0  # Hard judgment mode
+
+    async def get_reward(self, queries):
+        print("Processing queries[0]: {}".format(queries[0]))
+        tasks = [self.process_sample(query) for query in queries]
+        scores = await asyncio.gather(*tasks)
+        print("Generated scores: {}".format(scores))
+        if self.log_path:
+            with open(self.log_path, 'a', encoding='utf-8') as f:
+                unit = {
+                    "query_list": queries if isinstance(queries, list) else [],
+                    "hard_score_list": scores if isinstance(scores, list) else []
+                }
+                json.dump(unit, f, ensure_ascii=False)
+                f.write('\n')
+        if self.normalize_reward:
+            return reward_normalization(scores)
+        elif self.group_normalize_reward:
+            return group_reward_normalization(scores)
+        else:
+            return scores
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Reward Model
+    parser.add_argument("--tokenizer_path", type=str, default=None)
+    parser.add_argument("--answer_path", type=str, default=None)
+    parser.add_argument("--prob_reward", action="store_true", default=False)
+    parser.add_argument("--normalize_reward", action="store_true", default=False, help="Enable Reward Normazation")
+    parser.add_argument("--group_normalize_reward", action="store_true", default=False, help="Enable Group Reward Normazation")
+    parser.add_argument("--port", type=int, default=5000, help="Port number for the server")
+    parser.add_argument("--host", type=str, default="0.0.0.0", help="IP for the server")
+    parser.add_argument("--log_path", type=str, default=None)
+    parser.add_argument("--vllm_url", type=str, default=None)
+    parser.add_argument("--vllm_model", type=str, default=None)
+    args = parser.parse_args()
+    openai_api_key = "EMPTY"
+    openai_api_base = args.vllm_url
+
+    client = AsyncOpenAI(
+        api_key=openai_api_key,
+        base_url=openai_api_base,
+    )
+
+    # Server setup
+    reward_model = RewardModelProxy(args)
+    app = FastAPI()
+
+
+    @app.post("/get_reward")
+    async def get_reward(request: Request):
+        data = await request.json()
+        queries = data.get("query")
+        rewards = await reward_model.get_reward(queries)
+        result = {"rewards": rewards}
+        print(f"Sent JSON response: {result}")
+        return JSONResponse(result)
+
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

train.sh

@@ -0,0 +1,36 @@
+set -x
+
+EXPERIMENT_NAME=$1   #  for example, "sft_reward_training"
+PRETRAIN_PATH=$2     # path_to_Qwen2.5-7B-Instruct
+TRAIN_DATA_PATH=$3   # path_to_training_data (https://huggingface.co/datasets/sarosavo/Master-RM)
+
+working_dir=$(pwd)
+LOG_PATH=${working_dir}/${EXPERIMENT_NAME}/train.log
+SAVE_PATH=${working_dir}/${EXPERIMENT_NAME}/checkpoint
+mkdir -p ${SAVE_PATH}
+
+
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+
+deepspeed --module openrlhf.cli.train_sft \
+   --max_len 4096 \
+   --dataset $TRAIN_DATA_PATH \
+   --input_key query \
+   --output_key output \
+   --apply_chat_template \
+   --train_batch_size 128 \
+   --micro_train_batch_size 4 \
+   --pretrain $PRETRAIN_PATH \
+   --save_path $SAVE_PATH \
+   --save_steps -1 \
+   --logging_steps 1 \
+   --eval_steps -1 \
+   --zero_stage 3 \
+   --max_epochs 1 \
+   --bf16 \
+   --flash_attn \
+   --learning_rate 5e-6 \
+   --packing_samples \
+   2>&1 | tee ${LOG_PATH}
+