355 lines
12 KiB
Python
355 lines
12 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import copy
|
|
import json
|
|
import random
|
|
import time
|
|
import uuid
|
|
from functools import partial
|
|
from multiprocessing import Process, Queue
|
|
from time import sleep
|
|
|
|
import requests
|
|
from openai import OpenAI
|
|
from tqdm import tqdm
|
|
from slime.rollout.rm_hub import get_deepscaler_rule_based_reward
|
|
|
|
TASK_TYPE = "math"
|
|
|
|
SAMPLING_PARAMS = {
|
|
"top_p": 1,
|
|
}
|
|
|
|
|
|
def get_rule_based_math_reward(item):
|
|
messages = item["messages"]
|
|
label = item["label"]
|
|
assert messages[-1]["role"] == "assistant", "last message must be assistant, but got {}".format(
|
|
messages[-1]["role"]
|
|
)
|
|
|
|
response = messages[-1]["content"]
|
|
if response is None or len(response) == 0:
|
|
return 0
|
|
|
|
reward = get_deepscaler_rule_based_reward(response, label)
|
|
return reward
|
|
|
|
|
|
def query_single_turn(client, messages, sampling_params, tools=None):
|
|
base_payload = {
|
|
"messages": messages,
|
|
**sampling_params,
|
|
"model": "custom",
|
|
"stream": False,
|
|
"seed": random.randint(1, 10000000),
|
|
"tools": tools,
|
|
}
|
|
|
|
text = None
|
|
accumulated_tokens = 0
|
|
finish_reason = "stop"
|
|
|
|
for _attempt in range(6):
|
|
try:
|
|
# Create a fresh payload for each attempt
|
|
current_payload = copy.deepcopy(base_payload)
|
|
|
|
if text is not None:
|
|
# Update messages with current progress
|
|
current_messages = copy.deepcopy(messages)
|
|
current_messages.append({"role": "assistant", "content": text})
|
|
current_payload["messages"] = current_messages
|
|
|
|
# Adjust max_tokens based on accumulated tokens
|
|
if "max_tokens" in sampling_params:
|
|
current_payload["max_tokens"] = max(0, sampling_params["max_tokens"] - accumulated_tokens)
|
|
|
|
# Add continue flag for partial rollouts
|
|
current_payload["extra_body"] = {"continue_final_message": True}
|
|
if current_payload["max_tokens"] == 0:
|
|
break
|
|
response = client.chat.completions.create(**current_payload)
|
|
|
|
if len(response.choices) > 0:
|
|
finish_reason = response.choices[0].finish_reason
|
|
if finish_reason == "abort":
|
|
print(
|
|
f"query failed, reason: {response.choices[0].finish_reason}, currently generated: {response.usage.completion_tokens}"
|
|
)
|
|
|
|
accumulated_tokens += response.usage.completion_tokens
|
|
|
|
if text is None:
|
|
text = response.choices[0].message.content
|
|
else:
|
|
text += response.choices[0].message.content
|
|
|
|
sleep(10)
|
|
continue
|
|
if text is None:
|
|
text = response.choices[0].message.content
|
|
elif response.choices[0].message.content is not None:
|
|
text += response.choices[0].message.content
|
|
break
|
|
else:
|
|
print(f"Error in query, status code: {response.status_code}")
|
|
continue
|
|
except Exception as e:
|
|
print(f"query failed in single turn, error: {e}")
|
|
continue
|
|
|
|
# Update final messages
|
|
if len(messages) > 0 and messages[-1]["role"] == "assistant":
|
|
messages = messages[:-1]
|
|
messages.append({"role": "assistant", "content": text})
|
|
|
|
return messages, finish_reason
|
|
|
|
|
|
def worker_process(task_queue, done_queue, rollout_func, reward_func, client, sampling_params):
|
|
|
|
for line in iter(task_queue.get, "STOP"):
|
|
if isinstance(line, str):
|
|
item = json.loads(line)
|
|
else:
|
|
item = line
|
|
|
|
# try:
|
|
messages, finish_reason = rollout_func(client, item["prompt"], sampling_params)
|
|
|
|
item["uid"] = str(uuid.uuid4())
|
|
item["messages"] = messages
|
|
reward = reward_func(item)
|
|
item["rollout_index"] = 1
|
|
item["reward"] = reward
|
|
item["extra_info"] = {}
|
|
item.update(sampling_params)
|
|
item["timestamp"] = str(time.time())
|
|
item["round_number"] = len([_ for _ in item["messages"] if _["role"] == "assistant"])
|
|
item["finish_reason"] = finish_reason
|
|
|
|
output_item = {
|
|
"uid": item.pop("uid"),
|
|
"messages": messages,
|
|
"reward": reward,
|
|
"instance_id": item.pop("instance_id"),
|
|
"extra_info": item,
|
|
}
|
|
|
|
done_queue.put(output_item)
|
|
|
|
done_queue.put("COMPLETE")
|
|
|
|
|
|
class BaseGenerator:
|
|
def __init__(
|
|
self,
|
|
remote_engine_url,
|
|
remote_buffer_url,
|
|
num_repeat_per_sample=1,
|
|
queue_size=1000000,
|
|
num_process=10,
|
|
task_type="math",
|
|
max_tokens=4096,
|
|
num_repeats=10,
|
|
skip_instance_ids: list[str] | None = None,
|
|
):
|
|
self.queue_size = queue_size
|
|
self.num_process = num_process
|
|
self.remote_engine_url = remote_engine_url
|
|
self.remote_buffer_url = remote_buffer_url
|
|
self.num_repeat_per_sample = num_repeat_per_sample
|
|
self.task_type = task_type
|
|
self.max_tokens = max_tokens
|
|
self.num_repeats = num_repeats
|
|
# Ensure skip_instance_ids is a mutable list (copy to avoid modifying original)
|
|
self.skip_instance_ids = list(skip_instance_ids) if skip_instance_ids is not None else None
|
|
|
|
if self.skip_instance_ids is not None:
|
|
print(f"BaseGenerator initialized with {len(self.skip_instance_ids)} instance_ids to skip")
|
|
self.skip_instance_ids = self.skip_instance_ids * self.num_repeat_per_sample
|
|
|
|
if "/v1" in remote_engine_url:
|
|
self.client = OpenAI(api_key="test", base_url=remote_engine_url)
|
|
else:
|
|
remote_engine_url = remote_engine_url.strip("/") + "/v1"
|
|
self.client = OpenAI(api_key="test", base_url=remote_engine_url)
|
|
|
|
def send_data_to_buffer(self, data):
|
|
remote_buffer_url = self.remote_buffer_url.rstrip("/") + "/buffer/write"
|
|
|
|
for _ in range(2):
|
|
try:
|
|
response = requests.post(remote_buffer_url, json=data)
|
|
if response.status_code == 200:
|
|
break
|
|
else:
|
|
print(f"send data to buffer failed, status code: {response.status_code}")
|
|
continue
|
|
except Exception as e:
|
|
print(f"send data to buffer failed, error: {e}")
|
|
continue
|
|
|
|
def run(self, input_file, rollout_func, reward_func):
|
|
task_queue, done_queue = Queue(maxsize=self.queue_size), Queue(maxsize=self.queue_size)
|
|
|
|
def read_data_into_queue():
|
|
cnt = 0
|
|
items = []
|
|
skipped_count = 0
|
|
with open(input_file) as f:
|
|
for i, line in enumerate(f):
|
|
item = json.loads(line)
|
|
if "instance_id" not in item:
|
|
item["instance_id"] = i
|
|
items.append(item)
|
|
random.shuffle(items)
|
|
|
|
for _ in range(self.num_repeats):
|
|
|
|
for item in items:
|
|
for _ in range(self.num_repeat_per_sample):
|
|
item_repeat = copy.deepcopy(item)
|
|
|
|
if "uid" not in item_repeat:
|
|
item_repeat["uid"] = str(uuid.uuid4())
|
|
|
|
# Check if instance_id should be skipped
|
|
if self.skip_instance_ids is not None and item_repeat["instance_id"] in self.skip_instance_ids:
|
|
print(f"Skipping instance_id: {item_repeat['instance_id']}")
|
|
# Remove from skip list to handle potential duplicates in multiple epochs
|
|
self.skip_instance_ids.remove(item_repeat["instance_id"])
|
|
skipped_count += 1
|
|
continue
|
|
|
|
task_queue.put(item_repeat)
|
|
cnt += 1
|
|
time.sleep(300)
|
|
|
|
if skipped_count > 0:
|
|
remaining_skip_count = len(self.skip_instance_ids) if self.skip_instance_ids is not None else 0
|
|
print(
|
|
f"Rollout summary: skipped {skipped_count} instance_ids, {remaining_skip_count} still in skip list"
|
|
)
|
|
|
|
for _ in range(self.num_process):
|
|
task_queue.put("STOP")
|
|
|
|
processes = []
|
|
SAMPLING_PARAMS["max_tokens"] = self.max_tokens
|
|
|
|
for _ in range(self.num_process):
|
|
process = Process(
|
|
target=partial(worker_process, client=self.client, sampling_params=SAMPLING_PARAMS),
|
|
args=(task_queue, done_queue, rollout_func, reward_func),
|
|
)
|
|
process.start()
|
|
processes.append(process)
|
|
|
|
process = Process(target=read_data_into_queue)
|
|
process.start()
|
|
|
|
progress_bar = tqdm()
|
|
num_finished = 0
|
|
while num_finished < self.num_process:
|
|
item = done_queue.get()
|
|
if item == "COMPLETE":
|
|
num_finished += 1
|
|
else:
|
|
assert "reward" in item, f"reward not in item: {item}"
|
|
assert "instance_id" in item, f"instance_id not in item: {item}"
|
|
self.send_data_to_buffer(item)
|
|
progress_bar.update(1)
|
|
|
|
progress_bar.close()
|
|
|
|
return "finished"
|
|
|
|
def entry(self, input_file, rollout_func, reward_func, num_epoch=1):
|
|
for _ in range(num_epoch):
|
|
self.run(input_file, rollout_func, reward_func)
|
|
|
|
|
|
def run_rollout(data: dict):
|
|
|
|
print(f"Starting math rollout with data: {data}")
|
|
|
|
rollout_func = query_single_turn
|
|
reward_func = get_rule_based_math_reward
|
|
|
|
print("Waiting for 10 seconds for buffer server to start")
|
|
time.sleep(10)
|
|
global SAMPLING_PARAMS
|
|
for k, v in data["sampling_params"].items():
|
|
SAMPLING_PARAMS[k] = v
|
|
print(f"Set {k} to {v}", type(v))
|
|
|
|
generator = BaseGenerator(
|
|
data["remote_engine_url"],
|
|
data["remote_buffer_url"],
|
|
num_repeat_per_sample=int(data["num_repeat_per_sample"]),
|
|
queue_size=1000000,
|
|
max_tokens=int(data["sampling_params"]["max_tokens"]),
|
|
num_process=int(data.get("num_process", 100)),
|
|
task_type=data["task_type"],
|
|
skip_instance_ids=data.get("skip_instance_ids", None),
|
|
)
|
|
|
|
generator.entry(data["input_file"], rollout_func, reward_func, int(data.get("num_epoch", 1)))
|
|
|
|
|
|
def normalize_group_data(group, epsilon=1e-8, algo="grpo"):
|
|
print(f"Using math-specific normalization for group {group[0]}")
|
|
|
|
assert algo == "grpo", "Only 'grpo' is supported for now."
|
|
|
|
instance_id = group[0]
|
|
data = group[1]
|
|
rewards = [item["reward"] for item in data]
|
|
|
|
valid_rewards = [r for r in rewards if 1 >= r >= 0]
|
|
|
|
if set(valid_rewards) == {0}:
|
|
normalized_rewards = rewards
|
|
else:
|
|
mean_reward = sum(valid_rewards) / len(valid_rewards)
|
|
std_reward = (sum((r - mean_reward) ** 2 for r in valid_rewards) / len(valid_rewards)) ** 0.5
|
|
|
|
if std_reward < epsilon:
|
|
print(f"[Math Info] Zero variance in group {instance_id}, setting all to 0.")
|
|
normalized_rewards = [0.0 if 1 >= r >= 0 else r for r in rewards]
|
|
else:
|
|
normalized_rewards = [(r - mean_reward) / (std_reward + epsilon) if 1 >= r >= 0 else r for r in rewards]
|
|
|
|
for i, item in enumerate(data):
|
|
item["reward"] = normalized_rewards[i]
|
|
item["raw_reward"] = rewards[i]
|
|
|
|
return (instance_id, data)
|
|
|
|
|
|
def is_valid_group(group, min_valid_group_size, task_type="math"):
|
|
# Handle both tuple and list inputs
|
|
if isinstance(group, tuple):
|
|
instance_id, items = group
|
|
else:
|
|
items = group
|
|
|
|
# Count valid items (non-empty responses)
|
|
valid_indices = []
|
|
for i, item in enumerate(items):
|
|
if item["messages"][-1]["content"].strip():
|
|
valid_indices.append(i)
|
|
|
|
group_size = len(items)
|
|
valid_count = len(valid_indices)
|
|
|
|
# A group is finished if it has reached the target size
|
|
is_finished = group_size >= min_valid_group_size
|
|
|
|
is_valid = is_finished and valid_count >= min_valid_group_size
|
|
|
|
return is_valid
|