This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,58 @@
|
||||
# 8 * 80G
|
||||
# docs: https://swift.readthedocs.io/en/latest/Instruction/GRPO/AdvancedResearch/deepeyes.html
|
||||
|
||||
# First: Deploy Qwen2.5-VL-72B-Instruct for verify
|
||||
# CUDA_VISIBLE_DEVICES=4,5,6,7 \
|
||||
# swift deploy \
|
||||
# --model Qwen/Qwen2.5-VL-72B-Instruct \
|
||||
# --vllm_tensor_parallel_size 4
|
||||
|
||||
# Second: Run swift rollout to deploy rollout server
|
||||
# MAX_PIXELS=602112 \
|
||||
# CUDA_VISIBLE_DEVICES=3 \
|
||||
# swift rollout \
|
||||
# --model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
# --vllm_use_async_engine true \
|
||||
# --external_plugins examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py \
|
||||
# --multi_turn_scheduler deepeyes_scheduler \
|
||||
# --vllm_max_model_len 8192 \
|
||||
# --vllm_gpu_memory_utilization 0.8 \
|
||||
# --max_turns 5
|
||||
|
||||
# Third: Run swift rlhf to train GRPO model
|
||||
MAX_PIXELS=602112 \
|
||||
CUDA_VISIBLE_DEVICES=0,1,2 \
|
||||
NPROC_PER_NODE=3 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--dataset "path/to/data_0.1.2_visual_toolbox_v2.parquet"\
|
||||
"path/to/data_v0.8_visual_toolbox_v2.parquet"\
|
||||
"path/to/data_thinklite_reasoning_acc.parquet" \
|
||||
--load_from_cache_file true \
|
||||
--use_vllm true \
|
||||
--vllm_mode server \
|
||||
--vllm_server_host 127.0.0.1 \
|
||||
--vllm_server_port 8001 \
|
||||
--offload_optimizer true \
|
||||
--offload_model true \
|
||||
--sleep_level 1 \
|
||||
--external_plugins examples/train/grpo/plugin/deepeyes/deepeyes_plugin.py \
|
||||
--reward_funcs deepeyes_reward \
|
||||
--tuner_type full \
|
||||
--torch_dtype bfloat16 \
|
||||
--max_completion_length 2048 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--deepspeed zero3 \
|
||||
--learning_rate 1e-6 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--logging_steps 5 \
|
||||
--warmup_ratio 0.05 \
|
||||
--dataloader_num_workers 4 \
|
||||
--dataset_num_proc 4 \
|
||||
--num_generations 12 \
|
||||
--beta 0 \
|
||||
--temperature 0.9 \
|
||||
--report_to tensorboard swanlab \
|
||||
--log_completions true
|
||||
@@ -0,0 +1,443 @@
|
||||
# some code borrowed from https://github.com/Visual-Agent/DeepEyes/blob/main
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from math import ceil, floor
|
||||
from openai import OpenAI
|
||||
from PIL import Image
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from swift.rewards.orm import ORM, orms
|
||||
from swift.rollout.multi_turn import MultiTurnScheduler, multi_turns
|
||||
|
||||
try:
|
||||
from math_verify import parse, verify
|
||||
except ImportError as e:
|
||||
raise ImportError('please install math_verify by `pip install math_verify`') from e
|
||||
"""
|
||||
3 dataset file
|
||||
1. data_v0.8_visual_toolbox_v2.parquet: data_source == 'chart' (vl_agent.compute_score)
|
||||
2. data_0.1.2_visual_toolbox_v2.parquet : data_source == 'vstar' (vl_agent.compute_score)
|
||||
3. data_thinklite_reasoning_acc.parquet: data_source == 'thinklite_eureka' (vl_agent.compute_score_math)
|
||||
|
||||
tool:
|
||||
image_zoom_in_tool: zoom in the image, return a cropped image
|
||||
"""
|
||||
|
||||
MATH_VERIFY_PROMPT = """# CONTEXT #
|
||||
I am a teacher, and I have some high-level math problems. I am tasked with evaluating the correctness of a student's answer.
|
||||
Below, I am provided with a problem and a reference answer. Additionally, a student's answer is provided. My job is to assess whether the student's answer captures the same meaning as the reference answer, even when expressed with different wording or format.
|
||||
|
||||
# OBJECTIVE #
|
||||
I need you to judge whether the student's answer is correct given the ground truth answer.
|
||||
|
||||
Your tasks include:
|
||||
1. Identify Mathematical or Notational Equivalence: Pay special attention to any LaTeX expressions in both answers. Confirm that the mathematical relationships, variables, and operations conveyed are equivalent.
|
||||
|
||||
# TONE #
|
||||
Professional, scientific.
|
||||
|
||||
# RESPONSE: MARKDOWN REPORT #
|
||||
## Equivalence Judgement
|
||||
[Whether the student's answer share the same meaning with the reference answer. (TRUE or FALSE)]
|
||||
|
||||
# ATTENTION #
|
||||
- The reference answer is ALWAYS correct. You should carefully judge whether the student gives the same answer as reference answer.
|
||||
- The Equivalence Judgement is only TRUE or FALSE. The answer is FALSE even if the student's final answer almost correct with a minor mistakes.
|
||||
- Don't give extra explanation.
|
||||
|
||||
**Question**:
|
||||
{query}
|
||||
|
||||
**Reference Answer**
|
||||
{gold_ans}
|
||||
|
||||
## Student Final Answer
|
||||
{pred_ans}"""# noqa
|
||||
|
||||
|
||||
def extract_answer(action_string: str) -> Dict[str, any]:
|
||||
answer = re.findall(r'<answer>(.*?)</answer>', action_string, re.DOTALL)
|
||||
return answer[-1] if answer else None
|
||||
|
||||
|
||||
def extract_action(action_string: str) -> Dict[str, Any]:
|
||||
tool_call_match = re.findall(r'<tool_call>(.*?)</tool_call>', action_string, re.DOTALL)
|
||||
return tool_call_match[-1] if tool_call_match else None
|
||||
|
||||
|
||||
def get_chat_template():
|
||||
chat_template = """
|
||||
Below are two answers to a question. Question is [Question], [Standard Answer] is the standard answer to the question, and [Model_answer] is the answer extracted from a model's output to this question. Determine whether these two answers are consistent.
|
||||
Note that [Model Answer] is consistent with [Standard Answer] whenever they are essentially the same. If the meaning is expressed in the same way, it is considered consistent, for example, 'pink' and 'it is pink'.
|
||||
If they are consistent, Judement is 1; if they are different, Judement is 0. Just output Judement and don't output anything else.\n\n
|
||||
"""# noqa
|
||||
return chat_template
|
||||
|
||||
|
||||
def get_gpt4_score_ICE():
|
||||
example_1 = """
|
||||
[Question]: Is the countertop tan or blue?
|
||||
[Standard Answer]: The countertop is tan.
|
||||
[Model_answer] : tan
|
||||
Judgement: 1
|
||||
""" # noqa
|
||||
|
||||
example_2 = """
|
||||
[Question]: On which side of the picture is the barrier?
|
||||
[Standard Answer]: The barrier is on the left side of the picture.
|
||||
[Model_answer] : left
|
||||
Judgement: 1
|
||||
""" # noqa
|
||||
|
||||
example_3 = """
|
||||
[Question]: Is the kite brown and large?
|
||||
[Standard Answer]: Yes, the kite is brown and large.
|
||||
[Model_answer] : Yes
|
||||
Judgement: 1
|
||||
""" # noqa
|
||||
|
||||
example_4 = """
|
||||
[Question]: Are the spots on a giraffe?
|
||||
[Standard Answer]: No, the spots are on a banana.
|
||||
[Model_answer] : no
|
||||
Judgement: 1
|
||||
""" # noqa
|
||||
|
||||
example_5 = """
|
||||
[Question]: Who is wearing pants?
|
||||
[Standard Answer]: The boy is wearing pants.
|
||||
[Model_answer] : The person in the picture is wearing pants.
|
||||
Judgement: 1
|
||||
""" # noqa
|
||||
|
||||
example_6 = """
|
||||
[Question]: Is the man phone both blue and closed?
|
||||
[Standard Answer]: Yes, the man phone is both blue and closed.
|
||||
[Model_answer] : No.
|
||||
Judgement: 0
|
||||
""" # noqa
|
||||
|
||||
example_7 = """
|
||||
[Question]: What color is the towel in the center of the picture?
|
||||
[Standard Answer]: The towel in the center of the picture is blue.
|
||||
[Model_answer] : The towel in the center of the picture is pink.
|
||||
Judgement: 0
|
||||
""" # noqa
|
||||
|
||||
return [example_1, example_2, example_3, example_4, example_5, example_6, example_7]
|
||||
|
||||
|
||||
def get_prompt(predict_str, ground_truth, question):
|
||||
examples = get_gpt4_score_ICE()
|
||||
chat_template = get_chat_template()
|
||||
demo_prompt = chat_template
|
||||
for example in examples:
|
||||
demo_prompt += example + '\n\n'
|
||||
test_prompt = f"""
|
||||
[Question]: {question}
|
||||
[Standard Answer]: {ground_truth}
|
||||
[Model_answer] : {predict_str}
|
||||
Judgement:"""
|
||||
full_prompt = f'{demo_prompt}{test_prompt}'
|
||||
|
||||
return full_prompt
|
||||
|
||||
|
||||
def load_pil_image(img):
|
||||
try:
|
||||
if isinstance(img, Image.Image):
|
||||
return img
|
||||
|
||||
elif isinstance(img, Dict):
|
||||
return Image.open(io.BytesIO(img['bytes']))
|
||||
|
||||
elif isinstance(img, str):
|
||||
if os.path.exists(img):
|
||||
return Image.open(img)
|
||||
|
||||
if ',' in img:
|
||||
img_data = img.split(',')[1]
|
||||
else:
|
||||
img_data = img
|
||||
img_bytes = base64.b64decode(img_data)
|
||||
return Image.open(io.BytesIO(img_bytes))
|
||||
|
||||
elif isinstance(img, bytes):
|
||||
return Image.open(io.BytesIO(img))
|
||||
|
||||
elif hasattr(img, 'read'):
|
||||
return Image.open(img)
|
||||
else:
|
||||
return img
|
||||
|
||||
except Exception:
|
||||
return img
|
||||
|
||||
|
||||
def rule_math_verify(ground_truth, model_answer):
|
||||
gold = parse(ground_truth)
|
||||
answer = parse(model_answer)
|
||||
return verify(gold, answer)
|
||||
|
||||
|
||||
class DeepEyesReward(ORM):
|
||||
|
||||
def __init__(self, args, **kwargs):
|
||||
super().__init__(args)
|
||||
try:
|
||||
self.client = OpenAI(
|
||||
api_key='EMPTY',
|
||||
base_url='http://127.0.0.1:8000/v1',
|
||||
)
|
||||
self.verify_model_name = self.client.models.list().data[0].id
|
||||
except Exception as e:
|
||||
raise RuntimeError('Failed to connect to the model service. Please deploy the model '
|
||||
"using 'swift deploy' or 'vllm serve'.") from e
|
||||
|
||||
def __call__(self, completions, reward_model, extra_info, data_source, **kwargs) -> List[float]:
|
||||
# reference: https://github.com/Visual-Agent/DeepEyes/blob/main/verl/utils/reward_score/vl_agent.py
|
||||
# NOTE: reward_model is a column name from the dataset, which contains the ground truth answer
|
||||
rewards = []
|
||||
messages = kwargs.get('messages')
|
||||
for completion, solution, info, source, message in zip(completions, reward_model, extra_info, data_source,
|
||||
messages):
|
||||
sol = solution['ground_truth']
|
||||
info['messages'] = message
|
||||
if source in ['vstar', 'chart']:
|
||||
rewards.append(self.compute_score(completion, sol, info))
|
||||
elif source in ['thinklite_eureka']:
|
||||
rewards.append(self.compute_score_math(completion, sol, info))
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return rewards
|
||||
|
||||
def compute_score(self, predict_str: str, ground_truth: str, extra_info) -> float:
|
||||
is_format_error = False
|
||||
# predict_str = "<think>" + predict_str
|
||||
count_think_1 = predict_str.count('<think>')
|
||||
count_think_2 = predict_str.count('</think>')
|
||||
if count_think_1 != count_think_2:
|
||||
is_format_error = True
|
||||
count_tool_1 = predict_str.count('<tool_call>')
|
||||
count_tool_2 = predict_str.count('</tool_call>')
|
||||
if count_tool_1 != count_tool_2:
|
||||
is_format_error = True
|
||||
|
||||
predict_no_think = predict_str.split('</think>')[-1].strip()
|
||||
count_answer_1 = predict_no_think.count('<answer>')
|
||||
count_answer_2 = predict_no_think.count('</answer>')
|
||||
if count_answer_1 != count_answer_2:
|
||||
is_format_error = True
|
||||
|
||||
answer_text = predict_str.split('<answer>')[-1].split('</answer>')[0].strip()
|
||||
|
||||
question_text = extra_info['question']
|
||||
full_prompt = get_prompt(answer_text, ground_truth, question_text)
|
||||
|
||||
chat_response = self.client.chat.completions.create(
|
||||
model=self.verify_model_name,
|
||||
messages=[
|
||||
{
|
||||
'role': 'system',
|
||||
'content': 'You are a helpful assistant.'
|
||||
},
|
||||
{
|
||||
'role': 'user',
|
||||
'content': full_prompt
|
||||
},
|
||||
],
|
||||
seed=random.randint(0, 1000000),
|
||||
temperature=0.3,
|
||||
)
|
||||
response = chat_response.choices[0].message.content.strip()
|
||||
if 'Judgement:' in response:
|
||||
response = response.split('Judgement:')[-1].strip()
|
||||
if '1' in response:
|
||||
acc_reward = 1.0
|
||||
elif '0' in response:
|
||||
acc_reward = 0.0
|
||||
else:
|
||||
acc_reward = 0.0
|
||||
else:
|
||||
if response == '1':
|
||||
acc_reward = 1.0
|
||||
elif response == '0':
|
||||
acc_reward = 0.0
|
||||
else:
|
||||
acc_reward = 0.0
|
||||
|
||||
# Penalize for model trying to predict longer answer to hack llm-as-judge
|
||||
if len(answer_text) >= 1000:
|
||||
acc_reward = 0.0
|
||||
is_format_error = True
|
||||
|
||||
num_image = 0
|
||||
for message in extra_info['messages']:
|
||||
if message['role'] == 'user' and '<image>' in message['content']:
|
||||
num_image += 1
|
||||
# More than one image indicates a successful tool call.
|
||||
tool_reward = 1.0 if num_image > 1 and acc_reward > 0.5 else 0.0
|
||||
format_reward = -1.0 if is_format_error else 0.0
|
||||
|
||||
return 0.8 * acc_reward + 0.2 * format_reward + 1.2 * tool_reward
|
||||
|
||||
def compute_score_math(self, predict_str: str, ground_truth: str, extra_info=None) -> float:
|
||||
is_format_error = False
|
||||
# predict_str = "<think>" + predict_str
|
||||
count_think_1 = predict_str.count('<think>')
|
||||
count_think_2 = predict_str.count('</think>')
|
||||
if count_think_1 != count_think_2:
|
||||
is_format_error = True
|
||||
|
||||
model_answer = ''
|
||||
predict_no_think = predict_str.split('</think>')[-1].strip()
|
||||
answer_pattern = r'\\boxed{([^}]+)}'
|
||||
answer_list = re.findall(answer_pattern, predict_no_think, flags=re.DOTALL)
|
||||
if len(answer_list) == 0:
|
||||
acc_reward = 0.0
|
||||
is_format_error = True
|
||||
else:
|
||||
if len(answer_list) > 1:
|
||||
is_format_error = True
|
||||
|
||||
model_answer = answer_list[-1]
|
||||
if rule_math_verify(ground_truth, model_answer):
|
||||
acc_reward = 1.0
|
||||
else:
|
||||
acc_reward = 0
|
||||
full_prompt = MATH_VERIFY_PROMPT.format(
|
||||
query=extra_info['question'],
|
||||
gold_ans=ground_truth,
|
||||
pred_ans=model_answer,
|
||||
)
|
||||
response = ''
|
||||
for _ in range(8):
|
||||
try:
|
||||
chat_response = self.client.chat.completions.create(
|
||||
model=self.verify_model_name,
|
||||
messages=[
|
||||
{
|
||||
'role': 'user',
|
||||
'content': full_prompt
|
||||
},
|
||||
],
|
||||
seed=random.randint(0, 1000000),
|
||||
temperature=0.0,
|
||||
)
|
||||
response = chat_response.choices[0].message.content.strip()
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
judgement = response.split('## Equivalence Judgement')[-1].lower()
|
||||
if 'true' in judgement and 'false' not in judgement:
|
||||
acc_reward = 1.0
|
||||
|
||||
format_reward = -1.0 if is_format_error else 0.0
|
||||
return 1.2 * acc_reward + 0.4 * format_reward
|
||||
|
||||
|
||||
orms['deepeyes_reward'] = DeepEyesReward
|
||||
|
||||
|
||||
class VisualToolBoxScheduler(MultiTurnScheduler):
|
||||
user_prompt = ('\nThink first, call **image_zoom_in_tool** if needed, then answer. '
|
||||
'Format strictly as: <think>...</think> <tool_call>...</tool_call> (if tools needed)'
|
||||
' <answer>...</answer> ')
|
||||
|
||||
def __init__(self, infer_engine=None, max_turns=None, *args, **kwargs):
|
||||
super().__init__(infer_engine, max_turns, *args, **kwargs)
|
||||
|
||||
def check_finished(self, infer_request, response_choice, current_turn):
|
||||
should_stop = super().check_finished(infer_request, response_choice, current_turn)
|
||||
if should_stop:
|
||||
return True
|
||||
|
||||
last_completion = infer_request.messages[-1]['content']
|
||||
|
||||
action = extract_action(last_completion)
|
||||
# if the last completion is a tool call, do not finished yet
|
||||
if action:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def step(self, infer_request, response_choice, current_turn):
|
||||
from qwen_vl_utils import fetch_image
|
||||
completion = response_choice.message.content
|
||||
action = extract_action(completion)
|
||||
cropped_img = None
|
||||
extra_info = {}
|
||||
try:
|
||||
tool_call = json.loads(action.strip())
|
||||
tool_name = tool_call['name']
|
||||
if tool_name != 'image_zoom_in_tool':
|
||||
raise ValueError(f'Unknown tool name: {tool_name}')
|
||||
args = tool_call['arguments']
|
||||
bbox = args['bbox_2d']
|
||||
# NOTE: this function is only compatible with the QwenVL series models
|
||||
# If you use another MLLM, please adjust the fetch_image function accordingly
|
||||
# ensure the returned img is of type PIL.Image.Image and
|
||||
# has been processed to a maximum size of max_pixels
|
||||
img = fetch_image({'image': load_pil_image(infer_request.images[0])})
|
||||
|
||||
origin_height = img.height
|
||||
origin_width = img.width
|
||||
bbox = self.maybe_resize_bbox(bbox=bbox, origin_width=origin_width, origin_height=origin_height)
|
||||
# for invalid bbox, the exception will be catched in except block
|
||||
cropped_img = img.crop(bbox)
|
||||
query = '<tool_response>' + '<image>' + self.user_prompt + '</tool_response>'
|
||||
except Exception as e:
|
||||
error_msg = f'Invalid tool call format: {action.strip()}. Error: {e}'
|
||||
query = f'Error: {str(error_msg)}'
|
||||
|
||||
infer_request.messages.append({'role': 'user', 'content': query})
|
||||
if cropped_img:
|
||||
infer_request.images.append(cropped_img)
|
||||
# override the images
|
||||
extra_info['images'] = infer_request.images
|
||||
|
||||
# Return dictionary format according to new MultiTurnScheduler interface
|
||||
return {'infer_request': infer_request, 'rollout_infos': extra_info}
|
||||
|
||||
def validate_bbox(self, left, top, right, bottom):
|
||||
assert left < right and bottom > top, f'invalid shape for {left=}, {top=}, {right=}, {bottom=}'
|
||||
height = bottom - top
|
||||
width = right - left
|
||||
assert max(height, width) / min(height,
|
||||
width) <= 100, f'aspect ratio error: {left=}, {top=}, {right=}, {bottom=}'
|
||||
assert min(height, width) > 30, f'{height=}, {width=} is too small'
|
||||
return True
|
||||
|
||||
def maybe_resize_bbox(self, bbox, origin_width, origin_height):
|
||||
left, top, right, bottom = bbox
|
||||
|
||||
left = max(0, left)
|
||||
top = max(0, top)
|
||||
right = min(origin_width, right)
|
||||
bottom = min(origin_height, bottom)
|
||||
self.validate_bbox(left, top, right, bottom)
|
||||
|
||||
height = bottom - top
|
||||
width = right - left
|
||||
if height < 28 or width < 28:
|
||||
center_x = (left + right) / 2.0
|
||||
center_y = (top + bottom) / 2.0
|
||||
ratio = 28 / min(height, width)
|
||||
new_half_height = ceil(height * ratio * 0.5)
|
||||
new_half_width = ceil(width * ratio * 0.5)
|
||||
new_left = floor(center_x - new_half_width)
|
||||
new_right = ceil(center_x + new_half_width)
|
||||
new_top = floor(center_y - new_half_height)
|
||||
new_bottom = ceil(center_y + new_half_height)
|
||||
self.validate_bbox(new_left, new_top, new_right, new_bottom)
|
||||
return [new_left, new_top, new_right, new_bottom]
|
||||
return [left, top, right, bottom]
|
||||
|
||||
|
||||
multi_turns['deepeyes_scheduler'] = VisualToolBoxScheduler
|
||||
@@ -0,0 +1,43 @@
|
||||
SYSTEM_PROMPT="""You are a helpful math assistant. Solve the problem step by step and put your final answer within \\boxed{}."""
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
NPROC_PER_NODE=4 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen3.5-2B \
|
||||
--external_plugins examples/train/grpo/plugin/gsm8k/gsm8k_plugin.py \
|
||||
--reward_funcs gsm8k_accuracy gsm8k_format \
|
||||
--columns '{"answer": "solution"}' \
|
||||
--enable_thinking false \
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate \
|
||||
--vllm_gpu_memory_utilization 0.4 \
|
||||
--vllm_tensor_parallel_size 1 \
|
||||
--vllm_max_model_len 10240 \
|
||||
--sleep_level 1 \
|
||||
--tuner_type full \
|
||||
--torch_dtype bfloat16 \
|
||||
--dataset 'modelscope/gsm8k' \
|
||||
--load_from_cache_file true \
|
||||
--max_length 2048 \
|
||||
--max_completion_length 8192 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 4 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--learning_rate 1e-6 \
|
||||
--lr_scheduler_type cosine \
|
||||
--save_steps 10 \
|
||||
--save_total_limit 100 \
|
||||
--logging_steps 1 \
|
||||
--warmup_ratio 0.0 \
|
||||
--dataloader_num_workers 4 \
|
||||
--num_generations 8 \
|
||||
--temperature 1.0 \
|
||||
--system "$SYSTEM_PROMPT" \
|
||||
--deepspeed zero2 \
|
||||
--log_completions true \
|
||||
--report_to tensorboard swanlab \
|
||||
--max_grad_norm 1.0 \
|
||||
--epsilon 0.2 \
|
||||
--epsilon_high 0.28 \
|
||||
--scale_rewards none
|
||||
@@ -0,0 +1,49 @@
|
||||
import re
|
||||
from typing import List
|
||||
|
||||
from swift.rewards import ORM, orms
|
||||
|
||||
|
||||
class GSM8KAccuracy(ORM):
|
||||
|
||||
@staticmethod
|
||||
def extract_answer(text: str) -> str:
|
||||
"""Extract the last #### number from text."""
|
||||
text = text[-500:] if len(text) > 500 else text
|
||||
# Prefer \boxed{} format
|
||||
boxed = re.findall(r'\\boxed\{([^}]+)\}', text)
|
||||
if boxed:
|
||||
return boxed[-1].replace(',', '').replace(' ', '').strip()
|
||||
# Fallback to #### format
|
||||
matches = re.findall(r'####\s*([\-\d,\.\s]+)', text)
|
||||
if matches:
|
||||
return matches[-1].replace(',', '').replace(' ', '').strip()
|
||||
return ''
|
||||
|
||||
def __call__(self, completions, solution, **kwargs) -> List[float]:
|
||||
rewards = []
|
||||
for completion, gt_answer in zip(completions, solution):
|
||||
gt_num = self.extract_answer(gt_answer)
|
||||
pred_num = self.extract_answer(completion)
|
||||
correct = False
|
||||
if pred_num and gt_num:
|
||||
try:
|
||||
correct = abs(float(pred_num) - float(gt_num)) < 1e-5
|
||||
except (ValueError, OverflowError):
|
||||
correct = pred_num == gt_num
|
||||
rewards.append(1.0 if correct else 0.0)
|
||||
return rewards
|
||||
|
||||
|
||||
class GSM8KFormat(ORM):
|
||||
|
||||
def __call__(self, completions, **kwargs) -> List[float]:
|
||||
rewards = []
|
||||
for completion in completions:
|
||||
has_answer = bool(re.search(r'\\boxed\{[^}]+\}', completion) or re.search(r'####\s*[\-\d,\.]+', completion))
|
||||
rewards.append(1.0 if has_answer else 0.0)
|
||||
return rewards
|
||||
|
||||
|
||||
orms['gsm8k_accuracy'] = GSM8KAccuracy
|
||||
orms['gsm8k_format'] = GSM8KFormat
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
# ============================================================
|
||||
# Swift GRPO training with OpenEnv TextArena Sudoku
|
||||
#
|
||||
# Prerequisites:
|
||||
# 1. Start Sudoku server (separate terminal):
|
||||
# TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 \
|
||||
# python examples/train/grpo/plugin/openenv/start_sudoku_server.py
|
||||
#
|
||||
# 2. This script uses colocate mode:
|
||||
# - vLLM and training share the same GPUs
|
||||
# - No separate rollout server needed
|
||||
#
|
||||
# Environment: TextArena Sudoku (local server, port 8000)
|
||||
# Model: Qwen3.5-4B (enable_thinking=false)
|
||||
# Scheduler: SudokuScheduler (multi-turn, content diff tracking)
|
||||
# Multi-turn: max_turns=20 (20 moves per game)
|
||||
# Rewards: 5-component (empty_cell/valid_move/repetition/progress/correct)
|
||||
# Hints: Board parsing + guaranteed moves + candidates
|
||||
#
|
||||
# ============================================================
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
NPROC_PER_NODE=4 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen3.5-4B \
|
||||
--dataset examples/train/grpo/plugin/openenv/sudoku.jsonl#1000 \
|
||||
--external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
--enable_thinking false \
|
||||
--torch_dtype bfloat16 \
|
||||
--max_completion_length 256 \
|
||||
--max_length 8192 \
|
||||
--learning_rate 5e-6 \
|
||||
--num_train_epochs 3 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--num_generations 4 \
|
||||
--generation_batch_size 4 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--temperature 1 \
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate \
|
||||
--vllm_max_model_len 12288 \
|
||||
--vllm_gpu_memory_utilization 0.35 \
|
||||
--gradient_checkpointing true \
|
||||
--use_gym_env true \
|
||||
--multi_turn_scheduler sudoku_scheduler \
|
||||
--max_turns 20 \
|
||||
--save_strategy steps \
|
||||
--save_steps 50 \
|
||||
--logging_steps 1 \
|
||||
--log_completions true \
|
||||
--report_to tensorboard swanlab
|
||||
@@ -0,0 +1,65 @@
|
||||
# ============================================================
|
||||
# Swift GRPO training with OpenEnv TextArena Sudoku (Server Mode)
|
||||
#
|
||||
# Prerequisites:
|
||||
# 1. Start Sudoku server (separate terminal):
|
||||
# TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 \
|
||||
# python examples/train/grpo/plugin/openenv/start_sudoku_server.py
|
||||
#
|
||||
# 2. Start vLLM rollout server (separate terminal):
|
||||
# CUDA_VISIBLE_DEVICES=0 \
|
||||
# swift rollout \
|
||||
# --model Qwen/Qwen3.5-4B \
|
||||
# --external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
# --enable_thinking false \
|
||||
# --max_length 8192 \
|
||||
# --vllm_max_model_len 12288 \
|
||||
# --vllm_gpu_memory_utilization 0.9 \
|
||||
# --use_gym_env true \
|
||||
# --multi_turn_scheduler sudoku_scheduler \
|
||||
# --max_turns 20
|
||||
#
|
||||
# 3. This script starts training in server mode:
|
||||
# - vLLM rollout server handles multi-turn + env interaction
|
||||
# - Training process sends generation requests to rollout server
|
||||
# - --multi_turn_scheduler / --max_turns go to BOTH rollout and rlhf
|
||||
#
|
||||
# Environment: TextArena Sudoku (local server, port 8000)
|
||||
# Model: Qwen3.5-4B (enable_thinking=false)
|
||||
# Scheduler: SudokuScheduler (multi-turn, content diff tracking)
|
||||
# Multi-turn: max_turns=20 (20 moves per game)
|
||||
# Rewards: 5-component (empty_cell/valid_move/repetition/progress/correct)
|
||||
# Hints: Board parsing + guaranteed moves + candidates
|
||||
#
|
||||
# ============================================================
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1,2,3 \
|
||||
NPROC_PER_NODE=3 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen3.5-4B \
|
||||
--dataset examples/train/grpo/plugin/openenv/sudoku.jsonl \
|
||||
--external_plugins examples/train/grpo/plugin/openenv/sudoku_scheduler.py \
|
||||
--enable_thinking false \
|
||||
--torch_dtype bfloat16 \
|
||||
--max_completion_length 256 \
|
||||
--max_length 8192 \
|
||||
--learning_rate 5e-6 \
|
||||
--num_train_epochs 3 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--num_generations 6 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--temperature 1 \
|
||||
--use_vllm true \
|
||||
--vllm_mode server \
|
||||
--vllm_server_host 127.0.0.1 \
|
||||
--vllm_server_port 8001 \
|
||||
--gradient_checkpointing true \
|
||||
--use_gym_env true \
|
||||
--multi_turn_scheduler sudoku_scheduler \
|
||||
--max_turns 20 \
|
||||
--save_strategy steps \
|
||||
--save_steps 50 \
|
||||
--logging_steps 1 \
|
||||
--log_completions true \
|
||||
--report_to tensorboard swanlab
|
||||
@@ -0,0 +1,47 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Start the TextArena Sudoku server with configurable concurrent sessions.
|
||||
|
||||
The default OpenEnv server only allows 1 concurrent session because
|
||||
TextArenaEnvironment is not marked as SUPPORTS_CONCURRENT_SESSIONS.
|
||||
Since each WebSocket session creates an independent game instance,
|
||||
it is safe to enable concurrent sessions.
|
||||
|
||||
Usage:
|
||||
TEXTARENA_ENV_ID=Sudoku-v0 python start_sudoku_server.py
|
||||
TEXTARENA_ENV_ID=Sudoku-v0 MAX_CONCURRENT_ENVS=8 python start_sudoku_server.py
|
||||
"""
|
||||
import os
|
||||
import uvicorn
|
||||
from openenv.core.env_server.http_server import create_app
|
||||
from textarena_env.server.app import (TextArenaAction, TextArenaObservation, build_textarena_gradio_app,
|
||||
create_textarena_environment)
|
||||
from textarena_env.server.environment import TextArenaEnvironment
|
||||
|
||||
# Read config from environment
|
||||
# Note: TEXTARENA_ENV_ID is read by create_textarena_environment factory,
|
||||
# not by this script directly.
|
||||
max_concurrent_envs = int(os.getenv('MAX_CONCURRENT_ENVS', '8'))
|
||||
host = os.getenv('HOST', '0.0.0.0')
|
||||
port = int(os.getenv('PORT', '8000'))
|
||||
|
||||
# Mark TextArenaEnvironment as supporting concurrent sessions.
|
||||
# Each WebSocket session creates an independent game instance via the factory,
|
||||
# so concurrent sessions are safe.
|
||||
|
||||
TextArenaEnvironment.SUPPORTS_CONCURRENT_SESSIONS = True
|
||||
|
||||
# Build the app with custom max_concurrent_envs
|
||||
|
||||
app = create_app(
|
||||
create_textarena_environment,
|
||||
TextArenaAction,
|
||||
TextArenaObservation,
|
||||
env_name='textarena_env',
|
||||
max_concurrent_envs=max_concurrent_envs,
|
||||
gradio_builder=build_textarena_gradio_app,
|
||||
)
|
||||
|
||||
if __name__ == '__main__':
|
||||
env_id = os.getenv('TEXTARENA_ENV_ID', 'Sudoku-v0')
|
||||
print(f'Starting server: env={env_id}, max_concurrent_envs={max_concurrent_envs}')
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
@@ -0,0 +1,10 @@
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
{"messages": [{"role": "user", "content": "Play"}], "env_config": {"name": "openenv", "base_url": "http://127.0.0.1:8000", "reset_kwargs": {}}}
|
||||
@@ -0,0 +1,416 @@
|
||||
"""Sudoku scheduler for OpenEnv TextArena Sudoku environment.
|
||||
|
||||
Reference: TRL openenv_sudoku_grpo.ipynb
|
||||
Key features:
|
||||
1. Multiple reward functions (empty_cell, valid_move, repetition, progress, correct)
|
||||
2. Hints system: parse board, provide guaranteed moves and candidates
|
||||
3. Board state tracking with content diff for bounded context
|
||||
"""
|
||||
import asyncio
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from swift.rollout.multi_turn import OpenEnvScheduler, multi_turns
|
||||
|
||||
SUDOKU_SYSTEM_PROMPT = """You are an expert Sudoku player with deep knowledge of logical deduction strategies.
|
||||
|
||||
## GAME RULES
|
||||
1. The puzzle is a 9x9 grid divided into nine 3x3 subgrids (boxes)
|
||||
2. Some cells are pre-filled with numbers 1-9
|
||||
3. Fill empty cells ('.') with numbers 1-9
|
||||
4. Each row, column, and 3x3 box must contain 1-9 without repetition
|
||||
5. Cannot overwrite pre-filled cells
|
||||
6. Invalid moves result in penalties
|
||||
|
||||
## HOW TO PLAY
|
||||
Output your move in this format: [row col number]
|
||||
Example: [3 5 7] means place 7 at row 3, column 5.
|
||||
You may reason before your move, but always end with [row col number].
|
||||
|
||||
## STRATEGIC APPROACH
|
||||
- Naked Singles: If a cell has only one possible candidate, fill it immediately.
|
||||
- Hidden Singles: If a number can only go in one cell within a row/column/box, place it there.
|
||||
- Scanning: Look at each row, column, and box to find where numbers can go.
|
||||
|
||||
## COMMON PITFALLS
|
||||
- Don't guess randomly - Sudoku is pure logic
|
||||
- Don't overwrite pre-filled cells
|
||||
- Don't repeat a move that was already made
|
||||
- Coordinates are 1-indexed (1-9)
|
||||
|
||||
## BOARD READING
|
||||
- Rows labeled R1-R9 (top to bottom)
|
||||
- Columns labeled C1-C9 (left to right)
|
||||
- Empty cells shown as '.'"""
|
||||
|
||||
|
||||
def _is_valid_board_state(board_str: str) -> bool:
|
||||
return 'R1' in board_str and 'R9' in board_str and '|' in board_str
|
||||
|
||||
|
||||
def _parse_board(board_str: str) -> list:
|
||||
grid = [[0] * 9 for _ in range(9)]
|
||||
if not _is_valid_board_state(board_str):
|
||||
return grid
|
||||
for line in board_str.split('\n'):
|
||||
line_stripped = line.strip()
|
||||
if line_stripped and line_stripped[0] == 'R' and len(line_stripped) > 1 and line_stripped[1].isdigit():
|
||||
row = int(line_stripped[1]) - 1
|
||||
col = 0
|
||||
for char in line_stripped[2:]:
|
||||
if col >= 9:
|
||||
break
|
||||
if char == '.':
|
||||
grid[row][col] = 0
|
||||
col += 1
|
||||
elif char.isdigit():
|
||||
grid[row][col] = int(char)
|
||||
col += 1
|
||||
return grid
|
||||
|
||||
|
||||
def _count_filled_cells(board_str: str) -> int:
|
||||
grid = _parse_board(board_str)
|
||||
return sum(1 for row in grid for cell in row if cell != 0)
|
||||
|
||||
|
||||
def _get_valid_numbers(grid: list, row: int, col: int) -> set:
|
||||
if grid[row][col] != 0:
|
||||
return set()
|
||||
used = set()
|
||||
for c in range(9):
|
||||
if grid[row][c] != 0:
|
||||
used.add(grid[row][c])
|
||||
for r in range(9):
|
||||
if grid[r][col] != 0:
|
||||
used.add(grid[r][col])
|
||||
box_row, box_col = 3 * (row // 3), 3 * (col // 3)
|
||||
for r in range(box_row, box_row + 3):
|
||||
for c in range(box_col, box_col + 3):
|
||||
if grid[r][c] != 0:
|
||||
used.add(grid[r][c])
|
||||
return set(range(1, 10)) - used
|
||||
|
||||
|
||||
def _extract_empty_cells_with_candidates(board_str: str, sort_by_difficulty: bool = True):
|
||||
grid = _parse_board(board_str)
|
||||
cells = []
|
||||
for row in range(9):
|
||||
for col in range(9):
|
||||
if grid[row][col] == 0:
|
||||
candidates = _get_valid_numbers(grid, row, col)
|
||||
cells.append((row + 1, col + 1, candidates))
|
||||
if sort_by_difficulty:
|
||||
cells.sort(key=lambda x: len(x[2]))
|
||||
return cells
|
||||
|
||||
|
||||
def _extract_empty_cells(board_str: str) -> list:
|
||||
"""Return list of (row, col) tuples for empty cells, 0-indexed."""
|
||||
grid = _parse_board(board_str)
|
||||
return [(r, c) for r in range(9) for c in range(9) if grid[r][c] == 0]
|
||||
|
||||
|
||||
def _extract_board_only(text: str) -> str:
|
||||
if not text:
|
||||
return ''
|
||||
lines = text.split('\n')
|
||||
board_lines = []
|
||||
in_board = False
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped.startswith('C1') or (stripped and stripped[0] == 'R' and len(stripped) > 1
|
||||
and stripped[1].isdigit()):
|
||||
in_board = True
|
||||
if in_board and (stripped.startswith('-') or stripped.startswith('R') or stripped.startswith('C1')):
|
||||
board_lines.append(line)
|
||||
elif (in_board and stripped and not stripped.startswith('-')
|
||||
and not (stripped[0] == 'R' and len(stripped) > 1 and stripped[1].isdigit())):
|
||||
break
|
||||
return '\n'.join(board_lines) if board_lines else ''
|
||||
|
||||
|
||||
def _make_hints(board_str: str, successful_moves: list, failed_moves: list, difficulty: str = 'easy') -> str:
|
||||
parts = []
|
||||
all_tried = successful_moves + failed_moves
|
||||
if all_tried:
|
||||
parts.append(f"\nMOVES ALREADY TRIED (do not repeat): {', '.join(all_tried[:10])}")
|
||||
if not board_str or not _is_valid_board_state(board_str):
|
||||
return '\n'.join(parts)
|
||||
|
||||
cells = _extract_empty_cells_with_candidates(board_str, sort_by_difficulty=True)
|
||||
if cells:
|
||||
guaranteed = []
|
||||
other = []
|
||||
for r, c, candidates in cells[:10]:
|
||||
if len(candidates) == 1:
|
||||
guaranteed.append(f'[{r} {c} {list(candidates)[0]}]')
|
||||
elif len(candidates) <= 3:
|
||||
nums = ','.join(str(n) for n in sorted(candidates))
|
||||
other.append(f'({r},{c})->{nums}')
|
||||
if guaranteed:
|
||||
parts.append(f"\nGUARANTEED MOVES (only one option): {', '.join(guaranteed[:5])}")
|
||||
if other:
|
||||
parts.append(f"Other options: {' | '.join(other[:5])}")
|
||||
|
||||
return '\n'.join(parts)
|
||||
|
||||
|
||||
class SudokuScheduler(OpenEnvScheduler):
|
||||
"""Sudoku scheduler with multi-reward and hints system.
|
||||
|
||||
Tracks 5 reward components per trajectory:
|
||||
- empty_cell_reward: Did the model target empty cells? (+1/-1)
|
||||
- valid_move_reward: Were moves accepted by env? (1.0/-0.5/0.0)
|
||||
- repetition_reward: Penalty for repeating moves (exponential)
|
||||
- progress_reward: How much of the puzzle was filled (0-1)
|
||||
- correct_reward: Environment's reward (0 or 1)
|
||||
|
||||
Combined reward = sum of all components.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._last_content_len: Dict[str, int] = {}
|
||||
# Per-uuid state tracking
|
||||
self._board_states: Dict[str, str] = {}
|
||||
self._move_counts: Dict[str, defaultdict] = {}
|
||||
self._successful_moves: Dict[str, list] = {}
|
||||
self._failed_moves: Dict[str, list] = {}
|
||||
self._valid_move_scores: Dict[str, list] = {}
|
||||
self._empty_cell_scores: Dict[str, list] = {}
|
||||
self._correct_scores: Dict[str, list] = {}
|
||||
self._repetition_scores: Dict[str, list] = {}
|
||||
self._initial_filled: Dict[str, int] = {}
|
||||
self._max_filled: Dict[str, int] = {}
|
||||
|
||||
async def on_trajectory_start(self, requests):
|
||||
"""Initialize env, parse board, compute hints."""
|
||||
semaphore = asyncio.Semaphore(getattr(self, 'max_concurrent_envs', 4))
|
||||
|
||||
async def _init_single(req):
|
||||
async with semaphore:
|
||||
uuid = req.uuid
|
||||
if uuid in self._envs:
|
||||
await self._close_and_remove(uuid)
|
||||
|
||||
row_env_config = (req.data_dict or {}).get('env_config', {}) if hasattr(req, 'data_dict') else {}
|
||||
env_config = {**getattr(self, 'env_config_defaults', {}), **row_env_config}
|
||||
wrapper = self._create_env(env_config)
|
||||
|
||||
obs, metadata = await asyncio.to_thread(wrapper.reset)
|
||||
system_message = env_config.get('system_message', SUDOKU_SYSTEM_PROMPT)
|
||||
|
||||
content = self._extract_content(obs)
|
||||
self._last_content_len[uuid] = len(content)
|
||||
|
||||
# Parse initial board state
|
||||
board = _extract_board_only(content) if _is_valid_board_state(content) else content
|
||||
self._board_states[uuid] = content if _is_valid_board_state(content) else ''
|
||||
initial_filled = _count_filled_cells(self._board_states[uuid]) if self._board_states[uuid] else 0
|
||||
|
||||
# Initialize tracking state
|
||||
self._move_counts[uuid] = defaultdict(int)
|
||||
self._successful_moves[uuid] = []
|
||||
self._failed_moves[uuid] = []
|
||||
self._valid_move_scores[uuid] = []
|
||||
self._empty_cell_scores[uuid] = []
|
||||
self._correct_scores[uuid] = []
|
||||
self._repetition_scores[uuid] = []
|
||||
self._initial_filled[uuid] = initial_filled
|
||||
self._max_filled[uuid] = initial_filled
|
||||
|
||||
# Build initial message with board + hints
|
||||
hints = _make_hints(self._board_states[uuid], [], [])
|
||||
user_content = f'{board}{hints}' if board else content
|
||||
|
||||
from swift.rollout.multi_turn import Messages
|
||||
messages = []
|
||||
if system_message:
|
||||
messages.append({'role': 'system', 'content': system_message})
|
||||
messages.append({'role': 'user', 'content': user_content})
|
||||
req.messages = messages
|
||||
|
||||
self._envs[uuid] = wrapper
|
||||
self._total_rewards[uuid] = 0.0
|
||||
self._step_rewards[uuid] = []
|
||||
self._pending_obs[uuid] = None
|
||||
|
||||
await asyncio.gather(*[_init_single(req) for req in requests])
|
||||
|
||||
async def _close_and_remove(self, uuid):
|
||||
"""Override to clean up all tracking state."""
|
||||
await super()._close_and_remove(uuid)
|
||||
self._last_content_len.pop(uuid, None)
|
||||
self._board_states.pop(uuid, None)
|
||||
self._move_counts.pop(uuid, None)
|
||||
self._successful_moves.pop(uuid, None)
|
||||
self._failed_moves.pop(uuid, None)
|
||||
self._valid_move_scores.pop(uuid, None)
|
||||
self._empty_cell_scores.pop(uuid, None)
|
||||
self._correct_scores.pop(uuid, None)
|
||||
self._repetition_scores.pop(uuid, None)
|
||||
self._initial_filled.pop(uuid, None)
|
||||
self._max_filled.pop(uuid, None)
|
||||
|
||||
def _extract_content(self, observation: Any) -> str:
|
||||
if isinstance(observation, dict):
|
||||
messages = observation.get('messages', [])
|
||||
if messages:
|
||||
return messages[0].get('content', '')
|
||||
prompt = observation.get('prompt', '')
|
||||
if prompt:
|
||||
return prompt
|
||||
return str(observation)
|
||||
|
||||
async def on_turn_end(self, infer_request, response_choice, current_turn):
|
||||
"""Parse move, step env, compute multi-reward, generate hints."""
|
||||
uuid = infer_request.uuid
|
||||
wrapper = self._envs.get(uuid)
|
||||
if wrapper is None:
|
||||
return {'done': True, 'rollout_infos': {}}
|
||||
|
||||
action_text = response_choice.message.content
|
||||
action_dict = self.parse_action(action_text)
|
||||
if action_dict is None:
|
||||
# Parse failed: end trajectory with penalty instead of polluting env
|
||||
self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) - 1.0
|
||||
self._step_rewards.setdefault(uuid, []).append(-1.0)
|
||||
await self._close_and_remove(uuid)
|
||||
return {
|
||||
'done': True,
|
||||
'rollout_infos': {
|
||||
'total_reward': self._total_rewards[uuid],
|
||||
'step_rewards': list(self._step_rewards.get(uuid, [])),
|
||||
'gym_done': True,
|
||||
}
|
||||
}
|
||||
move = action_dict.get('message', '')
|
||||
|
||||
# Step environment
|
||||
obs, env_reward, done, metadata = await asyncio.to_thread(wrapper.step, action_dict)
|
||||
correct_score = float(env_reward or 0.0)
|
||||
|
||||
# Extract new content (diff from last seen)
|
||||
full_content = self._extract_content(obs)
|
||||
last_len = self._last_content_len.get(uuid, 0)
|
||||
new_content = full_content[last_len:] if len(full_content) > last_len else full_content
|
||||
self._last_content_len[uuid] = len(full_content)
|
||||
|
||||
# Check if env says invalid
|
||||
new_content_lower = new_content.lower()
|
||||
env_says_invalid = any(kw in new_content_lower
|
||||
for kw in ['invalid', 'error', 'cannot', 'already', 'violation', 'lost'])
|
||||
|
||||
# Check if move targets an empty cell
|
||||
if self._board_states.get(uuid):
|
||||
empty_cells = _extract_empty_cells(self._board_states[uuid])
|
||||
# Convert move coords (1-indexed from model) to 0-indexed for comparison
|
||||
move_nums = re.findall(r'\d+', move)
|
||||
targets_empty = tuple(int(x) - 1 for x in move_nums[:2]) in empty_cells if len(move_nums) >= 3 else True
|
||||
else:
|
||||
targets_empty = True
|
||||
|
||||
# Empty cell reward: +1 if targeted empty, -1 if tried to overwrite
|
||||
empty_cell_score = 1.0 if targets_empty else -1.0
|
||||
|
||||
# Repetition tracking
|
||||
is_new_move = self._move_counts[uuid][move] == 0
|
||||
repetition_count = self._move_counts[uuid][move]
|
||||
self._move_counts[uuid][move] += 1
|
||||
repetition_score = -min(2**repetition_count, 10.0) if repetition_count > 0 else 0.0
|
||||
|
||||
# Valid move score
|
||||
is_valid = not env_says_invalid and targets_empty
|
||||
if is_valid and is_new_move:
|
||||
valid_move_score = 1.0
|
||||
self._successful_moves[uuid].append(move)
|
||||
elif 'please resubmit' in new_content_lower or 'avoid penalties' in new_content_lower:
|
||||
valid_move_score = -0.5
|
||||
self._failed_moves[uuid].append(move)
|
||||
else:
|
||||
valid_move_score = 0.0
|
||||
if not is_valid:
|
||||
self._failed_moves[uuid].append(move)
|
||||
|
||||
# Update board state if valid and new content has board
|
||||
if is_valid and _is_valid_board_state(new_content):
|
||||
self._board_states[uuid] = new_content
|
||||
current_filled = _count_filled_cells(new_content)
|
||||
if current_filled > self._max_filled[uuid]:
|
||||
self._max_filled[uuid] = current_filled
|
||||
|
||||
# Progress reward
|
||||
remaining = 81 - self._initial_filled[uuid]
|
||||
if remaining > 0:
|
||||
progress_score = (self._max_filled[uuid] - self._initial_filled[uuid]) / remaining
|
||||
else:
|
||||
progress_score = 1.0
|
||||
|
||||
# Track all scores
|
||||
self._valid_move_scores[uuid].append(valid_move_score)
|
||||
self._empty_cell_scores[uuid].append(empty_cell_score)
|
||||
self._correct_scores[uuid].append(correct_score)
|
||||
self._repetition_scores[uuid].append(repetition_score)
|
||||
|
||||
combined_reward = (
|
||||
sum(self._empty_cell_scores[uuid]) / max(len(self._empty_cell_scores[uuid]), 1)
|
||||
+ sum(self._valid_move_scores[uuid]) / max(len(self._valid_move_scores[uuid]), 1)
|
||||
+ sum(self._repetition_scores[uuid]) / max(len(self._repetition_scores[uuid]), 1) + progress_score
|
||||
+ correct_score)
|
||||
|
||||
self._total_rewards[uuid] = combined_reward
|
||||
self._step_rewards.setdefault(uuid, []).append(combined_reward)
|
||||
|
||||
# Build next observation with board + hints
|
||||
if not done:
|
||||
board_str = self._board_states.get(uuid, '')
|
||||
board = _extract_board_only(board_str) if board_str else ''
|
||||
hints = _make_hints(
|
||||
board_str,
|
||||
self._successful_moves[uuid],
|
||||
self._failed_moves[uuid],
|
||||
)
|
||||
step_num = len(self._successful_moves[uuid])
|
||||
next_obs = f'Step {step_num}. Progress: {step_num} cells filled.\n\nBoard:\n{board}{hints}'
|
||||
else:
|
||||
next_obs = None
|
||||
|
||||
self._pending_obs[uuid] = next_obs
|
||||
|
||||
rollout_infos = {
|
||||
'total_reward': self._total_rewards[uuid],
|
||||
'step_rewards': list(self._step_rewards.get(uuid, [])),
|
||||
'gym_done': done,
|
||||
'empty_cell_reward': sum(self._empty_cell_scores[uuid]) / max(len(self._empty_cell_scores[uuid]), 1),
|
||||
'valid_move_reward': sum(self._valid_move_scores[uuid]) / max(len(self._valid_move_scores[uuid]), 1),
|
||||
'repetition_reward': sum(self._repetition_scores[uuid]) / max(len(self._repetition_scores[uuid]), 1),
|
||||
'progress_reward': progress_score,
|
||||
'correct_reward': correct_score,
|
||||
}
|
||||
if done:
|
||||
await self._close_and_remove(uuid)
|
||||
|
||||
return {'done': done, 'rollout_infos': rollout_infos}
|
||||
|
||||
def parse_action(self, text: str) -> Optional[Dict[str, Any]]:
|
||||
"""Extract [row col number] from model output. Returns None if parse fails."""
|
||||
match = re.search(r'\[\s*(\d+)\s+(\d+)\s+(\d+)\s*\]', text)
|
||||
if match:
|
||||
row, col, num = match.groups()
|
||||
return {'message': f'[{row} {col} {num}]'}
|
||||
|
||||
numbers = re.findall(r'\d+', text)
|
||||
if len(numbers) >= 3:
|
||||
return {'message': f'[{numbers[0]} {numbers[1]} {numbers[2]}]'}
|
||||
|
||||
return None
|
||||
|
||||
def format_observation(self, observation: Any) -> Union[str, List[Dict]]:
|
||||
return self._extract_content(observation)
|
||||
|
||||
|
||||
# Register scheduler so --external_plugins can load it
|
||||
|
||||
multi_turns['sudoku_scheduler'] = SudokuScheduler
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,49 @@
|
||||
# register customized plugins in plugin.py file
|
||||
|
||||
PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
NPROC_PER_NODE=4 \
|
||||
MAX_PIXELS=602112 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen2.5-VL-3B-Instruct \
|
||||
--external_plugins examples/train/grpo/plugin/plugin.py \
|
||||
--reward_funcs external_r1v_acc format \
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate \
|
||||
--vllm_gpu_memory_utilization 0.6 \
|
||||
--vllm_tensor_parallel_size 1 \
|
||||
--vllm_max_model_len 16384 \
|
||||
--tuner_type full \
|
||||
--torch_dtype bfloat16 \
|
||||
--dataset 'AI-ModelScope/clevr_cogen_a_train' \
|
||||
--overlong_filter false \
|
||||
--importance_sampling_level token \
|
||||
--epsilon 0.2 \
|
||||
--epsilon_high 0.28 \
|
||||
--max_completion_length 8192 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 4 \
|
||||
--learning_rate 1e-6 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--steps_per_generation 4 \
|
||||
--eval_steps 1000 \
|
||||
--save_steps 1000 \
|
||||
--save_total_limit 10 \
|
||||
--sleep_level 1 \
|
||||
--offload_model true \
|
||||
--offload_optimizer true \
|
||||
--logging_steps 1 \
|
||||
--dataloader_num_workers 4 \
|
||||
--num_generations 8 \
|
||||
--temperature 1.0 \
|
||||
--system 'examples/train/grpo/prompt.txt' \
|
||||
--deepspeed zero1 \
|
||||
--log_completions true \
|
||||
--report_to tensorboard swanlab \
|
||||
--num_iterations 1 \
|
||||
--async_generate false \
|
||||
--beta 0.001 \
|
||||
--loss_type grpo \
|
||||
--vllm_enable_lora false \
|
||||
--advantage_estimator grpo
|
||||
@@ -0,0 +1,23 @@
|
||||
# see rm_plugin example in swift/rewards/rm_plugin.py
|
||||
# register customized plugin in external_plugins file
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
|
||||
NPROC_PER_NODE=8 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen2.5-7B \
|
||||
--dataset AI-MO/NuminaMath-TIR#5000 \
|
||||
--load_from_cache_file true \
|
||||
--use_vllm true \
|
||||
--vllm_mode colocate \
|
||||
--vllm_gpu_memory_utilization 0.5 \
|
||||
--external_plugins examples/train/grpo/plugin/plugin.py \
|
||||
--reward_funcs format \
|
||||
--reward_model Qwen/Qwen2.5-3B-Instruct Shanghai_AI_Laboratory/internlm2-7b-reward \
|
||||
--reward_model_plugin genrm my_rmplugin \
|
||||
--reward_weights 0.1 1 1 \
|
||||
--sleep_level 1 \
|
||||
--offload_model true \
|
||||
--offload_optimizer true \
|
||||
--log_completions true \
|
||||
--deepspeed zero2
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
# This script require main branch ms-swift
|
||||
# This script is intended solely as a Tool Calling training example
|
||||
# The calculator tool implemented here can perform only basic arithmetic operations and may not be able to solve all math problems in the dataset.
|
||||
# Before running this script, please run the following `swift rollout` script first
|
||||
|
||||
# CUDA_VISIBLE_DEVICES=0 \
|
||||
# swift rollout \
|
||||
# --model Qwen/Qwen2.5-7B-Instruct \
|
||||
# --vllm_use_async_engine true \
|
||||
# --external_plugins examples/train/grpo/plugin/plugin.py \
|
||||
# --multi_turn_scheduler tool_call_scheduler \
|
||||
# --vllm_max_model_len 8192 \
|
||||
# --vllm_gpu_memory_utilization 0.8 \
|
||||
# --max_turns 5
|
||||
|
||||
SYSTEM_PROMPT='
|
||||
Answer the following questions as best you can. You have access to the following tools:
|
||||
|
||||
calculator
|
||||
Purpose: Perform basic arithmetic (+, -, *, /, parentheses) and return the result as text.
|
||||
Input (single string): the math expression to evaluate, e.g. "2*(3+4)".
|
||||
Only digits, spaces, and the characters +-*/(). are allowed.
|
||||
|
||||
Use the following format:
|
||||
|
||||
Question: the input question you must answer
|
||||
Thought: you should always think about what to do
|
||||
Action: the action to take, should be one of [calculator]
|
||||
Action Input: the input to the action
|
||||
Observation: the result of the action
|
||||
... (this Thought/Action/Action Input/Observation can be repeated zero or more times)
|
||||
Thought: I now know the final answer
|
||||
Final Answer: the final answer to the original input question, the answer should be written as \(\boxed{<answer>}\), e.g. \(\boxed{10}\)
|
||||
|
||||
Begin!
|
||||
'
|
||||
|
||||
CUDA_VISIBLE_DEVICES=1,2,3 \
|
||||
NPROC_PER_NODE=3 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen2.5-7B-Instruct \
|
||||
--reward_funcs accuracy \
|
||||
--tuner_type full \
|
||||
--torch_dtype bfloat16 \
|
||||
--use_vllm true \
|
||||
--vllm_mode server \
|
||||
--vllm_server_host 127.0.0.1 \
|
||||
--vllm_server_port 8000 \
|
||||
--dataset 'AI-MO/NuminaMath-TIR#1000' \
|
||||
--load_from_cache_file true \
|
||||
--max_completion_length 2048 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--learning_rate 1e-5 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--eval_steps 100 \
|
||||
--save_steps 100 \
|
||||
--save_total_limit 2 \
|
||||
--logging_steps 5 \
|
||||
--output_dir output \
|
||||
--warmup_ratio 0.05 \
|
||||
--dataloader_num_workers 4 \
|
||||
--dataset_num_proc 4 \
|
||||
--num_generations 4 \
|
||||
--temperature 0.9 \
|
||||
--system "$SYSTEM_PROMPT" \
|
||||
--log_completions true \
|
||||
--deepspeed zero3 \
|
||||
--stop_words "Observation:" \
|
||||
--report_to swanlab tensorboard
|
||||
@@ -0,0 +1,214 @@
|
||||
import re
|
||||
import torch
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from modelscope.preprocessors.templates.utils import Messages
|
||||
from typing import List
|
||||
|
||||
from swift.infer_engine.protocol import ChatCompletionResponseChoice
|
||||
|
||||
|
||||
class SampleStatus(Enum):
|
||||
INITIAL = 'initial'
|
||||
TO_INFER = 'to_infer'
|
||||
FINISH_NEXT_INFER = 'finish_next_infer'
|
||||
FINISHED = 'finished'
|
||||
ROLLBACK = 'rollback'
|
||||
|
||||
|
||||
class FinishedReason(Enum):
|
||||
ANSWER = 'finished_with_answer'
|
||||
MAX_INFER_STEP = 'finished_with_max_infer_steps'
|
||||
UNFINISHED = 'unfinished'
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataSampleTree:
|
||||
"""
|
||||
Attributes:
|
||||
tree_idx (str):
|
||||
for example 0/1-2/2-3/4-0, root_node = 0, next node = 1-2 infer batch 1 and index 2 sample
|
||||
|
||||
last_response (ChatCompletionResponseChoice):
|
||||
vllm previous round output
|
||||
"""
|
||||
tree_idx: str
|
||||
request_id: str
|
||||
|
||||
messages: Messages
|
||||
logprobs: List[List[float]] = field(default_factory=list)
|
||||
|
||||
all_response_ids: List[List[int]] = field(default_factory=list)
|
||||
last_response: ChatCompletionResponseChoice = None
|
||||
|
||||
token_count_per_step: List[int] = field(default_factory=list)
|
||||
|
||||
status: SampleStatus = SampleStatus.INITIAL
|
||||
finished_reason: FinishedReason = FinishedReason.UNFINISHED
|
||||
|
||||
@property
|
||||
def root_node(self):
|
||||
return int(self.tree_idx.split('/')[0])
|
||||
|
||||
@property
|
||||
def depth(self):
|
||||
return len(self.tree_idx.split('/')) - 1
|
||||
|
||||
@property
|
||||
def response_num(self):
|
||||
return len(self.all_response_ids)
|
||||
|
||||
def response_truncate(self, truncate_len: int):
|
||||
"""
|
||||
Before rollback, truncate the response.
|
||||
"""
|
||||
|
||||
if truncate_len < 1:
|
||||
return
|
||||
|
||||
self.logprobs = self.logprobs[:-truncate_len]
|
||||
self.all_response_ids = self.all_response_ids[:-truncate_len]
|
||||
self.messages = self.messages[:-(truncate_len * 2 - 1)]
|
||||
self.last_response = None
|
||||
|
||||
def extend_response(self, choice: ChatCompletionResponseChoice):
|
||||
self.extend_response_text(choice.message.content)
|
||||
self.extend_logprobs([item['logprob'] for item in choice.logprobs['content']])
|
||||
|
||||
self.all_response_ids.append(choice.token_ids)
|
||||
self.token_count_per_step.append(len(choice.token_ids))
|
||||
|
||||
choice.logprobs = None
|
||||
self.last_response = deepcopy(choice)
|
||||
|
||||
def extend_response_text(self, response_text: str):
|
||||
self.messages.append({'role': 'assistant', 'content': response_text})
|
||||
|
||||
def extend_logprobs(self, logprobs: List[float]):
|
||||
self.logprobs.append(logprobs)
|
||||
|
||||
|
||||
def _repeat_list_interleave(any_list, repeat_times):
|
||||
# return [item for sublist in [[item] * repeat_times for item in any_list] for item in sublist]
|
||||
return [deepcopy(item) for sublist in [[item] * repeat_times for item in any_list] for item in sublist]
|
||||
|
||||
|
||||
def _increment_tree_idx_depth(
|
||||
samples: list[DataSampleTree],
|
||||
next_infer_step: int,
|
||||
) -> list[DataSampleTree]:
|
||||
for infer_batch_idx, sample in enumerate(samples):
|
||||
sample.tree_idx = sample.tree_idx + '/' + f'{next_infer_step}-{infer_batch_idx}'
|
||||
return samples
|
||||
|
||||
|
||||
def extract_last_boxed(text):
|
||||
pattern = r'\\boxed\{((?:[^{}]|\{(?:[^{}]|\{[^{}]*\})*\})*)\}'
|
||||
|
||||
matches = list(re.finditer(pattern, text))
|
||||
if matches:
|
||||
return matches[-1].group(0)
|
||||
return None
|
||||
|
||||
|
||||
class AbstractDivergence:
|
||||
|
||||
@classmethod
|
||||
def calc_weights(cls, root_idx, samples_to_go_deeper, **kwargs) -> List[float]:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def allocate_with_weights(cls, weights, budget, max_divergence) -> List[int]:
|
||||
n = len(weights)
|
||||
alloc = [0] * n
|
||||
|
||||
w = [float(wi) if wi is not None and wi > 0 else 0.0 for wi in weights]
|
||||
total_w = sum(w)
|
||||
if total_w == 0:
|
||||
return alloc
|
||||
|
||||
# first round of allocation by weight ratio
|
||||
ideals = [(w[i] / total_w) * budget if w[i] > 0 else 0.0 for i in range(n)]
|
||||
for i in range(n):
|
||||
if w[i] <= 0:
|
||||
continue
|
||||
f = int(ideals[i])
|
||||
alloc[i] = min(f, max_divergence)
|
||||
|
||||
# second round of allocation by greedy allocation
|
||||
remain = budget - sum(alloc)
|
||||
if remain <= 0:
|
||||
return alloc
|
||||
|
||||
# weights desc, index asc
|
||||
remainders = [(ideals[i] - int(ideals[i]), i) for i in range(n) if w[i] > 0 and alloc[i] < max_divergence]
|
||||
remainders.sort(key=lambda x: (-x[0], x[1]))
|
||||
|
||||
idx = 0
|
||||
while remain > 0 and remainders:
|
||||
frac, i = remainders[idx % len(remainders)]
|
||||
if alloc[i] < max_divergence:
|
||||
alloc[i] += 1
|
||||
remain -= 1
|
||||
|
||||
if alloc[i] >= max_divergence:
|
||||
remainders = [r for r in remainders if r[1] != i]
|
||||
idx = 0
|
||||
continue
|
||||
idx += 1
|
||||
|
||||
return alloc
|
||||
|
||||
@classmethod
|
||||
def apply(cls, root_idx, samples_to_go_deeper, divergence_budget, max_divergence, **kwargs) -> List[DataSampleTree]:
|
||||
"""
|
||||
Args:
|
||||
root_idx: current root node idx
|
||||
samples_to_go_deeper: go deeper samples which root_node = root_idx
|
||||
divergence_budget: total divergence
|
||||
max_divergence: each sample max divergence
|
||||
"""
|
||||
weights = cls.calc_weights(root_idx, samples_to_go_deeper, **kwargs)
|
||||
allocate_divergence = cls.allocate_with_weights(weights, divergence_budget, max_divergence)
|
||||
|
||||
divergence_samples = []
|
||||
for sample, divergence in zip(samples_to_go_deeper, allocate_divergence):
|
||||
for _ in range(divergence):
|
||||
divergence_samples.append(deepcopy(sample))
|
||||
|
||||
return divergence_samples
|
||||
|
||||
|
||||
class LogProbDivergence(AbstractDivergence):
|
||||
|
||||
@classmethod
|
||||
def calc_weights(cls, root_idx, samples_to_go_deeper, **kwargs) -> List[float]:
|
||||
"""
|
||||
In this strategy, weight is proportional to entropy
|
||||
"""
|
||||
entropies = []
|
||||
for sample in samples_to_go_deeper:
|
||||
log_probs = torch.tensor(sample.logprobs[-1])
|
||||
|
||||
probs = torch.exp(log_probs)
|
||||
entropy = -torch.sum(probs * log_probs)
|
||||
entropies.append(entropy)
|
||||
|
||||
entropies_tensor = torch.stack(entropies)
|
||||
weights = torch.softmax(entropies_tensor, dim=0)
|
||||
|
||||
return weights.tolist()
|
||||
|
||||
|
||||
class AvgDivergence(AbstractDivergence):
|
||||
|
||||
@classmethod
|
||||
def calc_weights(cls, root_idx, samples_to_go_deeper, **kwargs) -> List[float]:
|
||||
avg = torch.ones(len(samples_to_go_deeper))
|
||||
weights = torch.softmax(avg, dim=0)
|
||||
|
||||
return weights.tolist()
|
||||
|
||||
|
||||
DivergenceStrategyMapping = {'logprobs': LogProbDivergence, 'average': AvgDivergence}
|
||||
@@ -0,0 +1,50 @@
|
||||
# This script is a example for multi-turn training with tree-rollout.
|
||||
# Regarding parameter configuration, currently tree_rollout, acting as the inference side, cannot receive relevant training parameters. Please note the following:
|
||||
# 1.Ensure that max_tree_width in tree_rollout is equal to num_generations.
|
||||
# 2.If DP (Data Parallelism) is enabled during the rollout stage, ensures that data within the same group is allocated to the same inference device.
|
||||
# For example: If generation_batch_size(per_device_batch_size * gradient_accumulation_steps * num_processes) = 32 and num_generations = 8,
|
||||
# then the rollout DP num should equal 4/2/1.
|
||||
# For more details on tool invocation, dialogue termination criteria, and other logic, please refer to the TreeRolloutScheduler implementation.
|
||||
|
||||
# First: Run swift rollout to deploy rollout server
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
swift rollout \
|
||||
--model Qwen/Qwen2.5-0.5B \
|
||||
--vllm_use_async_engine true \
|
||||
--external_plugins examples/train/grpo/plugin/treepo/tree_rollout_plugin.py \
|
||||
--multi_turn_scheduler tree_rollout_scheduler \
|
||||
--max_turns 6
|
||||
|
||||
|
||||
# Second: Run swift rlhf to train GRPO model
|
||||
CUDA_VISIBLE_DEVICES=1 \
|
||||
swift rlhf \
|
||||
--rlhf_type grpo \
|
||||
--model Qwen/Qwen2.5-0.5B \
|
||||
--reward_funcs accuracy \
|
||||
--use_vllm true \
|
||||
--vllm_mode server \
|
||||
--vllm_server_host 127.0.0.1 \
|
||||
--vllm_server_port 8000 \
|
||||
--tuner_type full \
|
||||
--torch_dtype bfloat16 \
|
||||
--external_plugins examples/train/grpo/plugin/treepo/tree_rollout_plugin.py \
|
||||
--dataset AI-MO/NuminaMath-TIR#1000 \
|
||||
--split_dataset_ratio 0 \
|
||||
--max_completion_length 2048 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 2 \
|
||||
--learning_rate 1e-6 \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--save_total_limit 2 \
|
||||
--logging_steps 1 \
|
||||
--warmup_ratio 0.05 \
|
||||
--dataloader_num_workers 4 \
|
||||
--dataset_num_proc 4 \
|
||||
--num_generations 8 \
|
||||
--temperature 1.0 \
|
||||
--top_p 0.9 \
|
||||
--top_k 50 \
|
||||
--log_completions true \
|
||||
--num_iterations 1 \
|
||||
--beta 0.04
|
||||
@@ -0,0 +1,254 @@
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from concurrent.futures import ALL_COMPLETED, ThreadPoolExecutor, wait
|
||||
from copy import deepcopy
|
||||
from tree_rollout import (DataSampleTree, DivergenceStrategyMapping, FinishedReason, SampleStatus,
|
||||
_increment_tree_idx_depth, _repeat_list_interleave, extract_last_boxed)
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from swift.infer_engine import RequestConfig
|
||||
from swift.infer_engine.protocol import ChatCompletionResponse, RolloutInferRequest, RolloutOutput
|
||||
from swift.rewards import MultiTurnScheduler, multi_turns
|
||||
|
||||
|
||||
class TreeRolloutScheduler(MultiTurnScheduler):
|
||||
"""
|
||||
Base class for multi-turn tree-rollout scheduling.
|
||||
|
||||
Provides default implementation for multi-turn conversation management.
|
||||
|
||||
CUSTOMIZATION:
|
||||
Implement the required `step()` method and optionally override `check_finished()`
|
||||
- Uses TreeRolloutScheduler's run() method infrastructure
|
||||
- Only need to implement turn transition logic in step()
|
||||
- Optionally customize termination conditions
|
||||
|
||||
Attributes:
|
||||
max_tree_width (int):
|
||||
For GRPO, it must be equal to num_generations.
|
||||
max_tree_depth (int):
|
||||
Controls the maximum number of reasoning turns for a single prompt.
|
||||
root_divergence (int):
|
||||
Number of branches generated in the first-round inference at the root node.
|
||||
max_divergence (int):
|
||||
Maximum number of branches allowed for each node.
|
||||
divergence_strategy (str):
|
||||
Strategy for selecting branch nodes; defaults to logprobs.
|
||||
"""
|
||||
|
||||
def __init__(self, infer_engine=None, max_turns=None, *args, **kwargs):
|
||||
super().__init__(infer_engine, max_turns, *args, **kwargs)
|
||||
self.max_tree_width = 8
|
||||
self.max_tree_depth = max_turns | 6
|
||||
self.max_divergence = 2
|
||||
self.divergence_strategy = 'logprobs'
|
||||
self.root_divergence = 1
|
||||
|
||||
self.executor = ThreadPoolExecutor(max_workers=self.max_tree_width)
|
||||
|
||||
async def async_infer(self,
|
||||
infer_requests: List[Union['RolloutInferRequest', Dict[str, Any]]],
|
||||
request_config: 'RequestConfig',
|
||||
*,
|
||||
use_tqdm: Optional[bool] = None,
|
||||
**kwargs) -> List['RolloutOutput']:
|
||||
# dedup_requests_by_messages
|
||||
processed_request = []
|
||||
seen = set()
|
||||
uuids = []
|
||||
|
||||
for item in infer_requests:
|
||||
if isinstance(item, dict):
|
||||
req = RolloutInferRequest(**item)
|
||||
else:
|
||||
req = item
|
||||
|
||||
msg_key = json.dumps(req.messages, sort_keys=True)
|
||||
uuids.append(req.uuid)
|
||||
|
||||
if msg_key not in seen:
|
||||
seen.add(msg_key)
|
||||
processed_request.append(req)
|
||||
|
||||
request_config.logprobs = True
|
||||
|
||||
outputs = await super().async_infer(processed_request, request_config, use_tqdm=use_tqdm, **kwargs)
|
||||
|
||||
assert len(outputs) == len(uuids), '[Tree Rollout] Please check the max_tree_width is equal to num_generations.'
|
||||
|
||||
for idx, output in enumerate(outputs):
|
||||
output.response.id = uuids[idx]
|
||||
|
||||
return outputs
|
||||
|
||||
async def run(self, infer_request: Union[List[RolloutInferRequest], RolloutInferRequest],
|
||||
request_config: 'RequestConfig', **kwargs) -> List['RolloutOutput']:
|
||||
if isinstance(infer_request, RolloutInferRequest):
|
||||
infer_request = [infer_request]
|
||||
else:
|
||||
infer_request = list(infer_request)
|
||||
|
||||
request_config.logprobs = True
|
||||
|
||||
finished_rollout_by_root: Dict[int, List[RolloutOutput]] = {i: [] for i in range(len(infer_request))}
|
||||
finished_samples: Dict[int, List[DataSampleTree]] = {i: [] for i in range(len(infer_request))}
|
||||
|
||||
samples_to_infer = []
|
||||
|
||||
for root_idx in range(len(infer_request)):
|
||||
samples_to_infer.append(
|
||||
DataSampleTree(
|
||||
tree_idx=str(root_idx),
|
||||
request_id=infer_request[root_idx].uuid,
|
||||
messages=infer_request[root_idx].messages,
|
||||
status=SampleStatus.TO_INFER))
|
||||
|
||||
# first step
|
||||
next_infer_step = 1
|
||||
samples_to_infer = _repeat_list_interleave(samples_to_infer, self.root_divergence)
|
||||
samples_to_infer = _increment_tree_idx_depth(samples_to_infer, next_infer_step)
|
||||
|
||||
while len(samples_to_infer) > 0:
|
||||
# resolve the error: Request id xxx already running
|
||||
vllm_inputs = [
|
||||
RolloutInferRequest(messages=sample.messages, uuid=f'{sample.request_id}-{sample.tree_idx}')
|
||||
for sample in samples_to_infer
|
||||
]
|
||||
|
||||
# Get model response
|
||||
tasks = [self.infer_engine.infer_async(request, request_config, **kwargs) for request in vllm_inputs]
|
||||
outputs: List[ChatCompletionResponse] = await asyncio.gather(*tasks)
|
||||
|
||||
assert len(vllm_inputs) == len(
|
||||
outputs), f'outputs length {len(outputs)} != inputs length {len(vllm_inputs)}'
|
||||
|
||||
samples_last_step = deepcopy(samples_to_infer)
|
||||
samples_to_infer = []
|
||||
|
||||
for idx, (sample, output) in enumerate(zip(samples_last_step, outputs)):
|
||||
assert len(output.choices) == 1, 'vllm should only generate one output'
|
||||
self.check_finished(sample, output)
|
||||
|
||||
# bind the output and request
|
||||
output.id = sample.request_id
|
||||
choice = output.choices[0]
|
||||
child_sample = deepcopy(sample)
|
||||
child_sample.extend_response(choice)
|
||||
|
||||
if child_sample.status == SampleStatus.FINISHED:
|
||||
finished_samples[child_sample.root_node].append(child_sample)
|
||||
finished_rollout_by_root[child_sample.root_node].append(
|
||||
RolloutOutput(
|
||||
response=output,
|
||||
messages=deepcopy(child_sample.messages),
|
||||
response_token_ids=deepcopy(child_sample.all_response_ids),
|
||||
# If we use intermediate reasoning results when computing the reward,
|
||||
# but loss_mask is not explicitly set,
|
||||
# only the loss of the final round of reasoning will be computed.
|
||||
response_loss_mask=[[1] * len(response_ids)
|
||||
for response_ids in child_sample.all_response_ids],
|
||||
rollout_infos={'num_turns': next_infer_step},
|
||||
))
|
||||
else:
|
||||
samples_to_infer.append(child_sample)
|
||||
|
||||
# if we have budget, do divergence
|
||||
if len(samples_to_infer) > 0 and self.max_divergence > 1:
|
||||
for root_idx in finished_samples.keys():
|
||||
root_to_infer_samples = [sample for sample in samples_to_infer if sample.root_node == root_idx]
|
||||
root_finished_samples = finished_samples[root_idx]
|
||||
|
||||
budget = self.max_tree_width - len(root_finished_samples) - len(root_to_infer_samples)
|
||||
|
||||
if budget > 0 and len(root_to_infer_samples) > 0:
|
||||
divergence_executor = DivergenceStrategyMapping[self.divergence_strategy]
|
||||
if not divergence_executor:
|
||||
raise ValueError(
|
||||
f"[Tree Rollout] The divergence strategy: {self.divergence_strategy} doesn't exist.")
|
||||
|
||||
divergence_samples = divergence_executor.apply(root_idx, root_to_infer_samples, budget,
|
||||
self.max_divergence - 1)
|
||||
samples_to_infer.extend(divergence_samples)
|
||||
|
||||
# before end loop, if finished_count < max_tree_width, rollback
|
||||
if len(samples_to_infer) == 0 and any(count < self.max_tree_width
|
||||
for count in [len(value) for value in finished_samples.values()]):
|
||||
samples_to_infer = self.roll_back_to_divergence(finished_samples)
|
||||
|
||||
# tools call etc
|
||||
futures = [self.executor.submit(self.step, sample) for sample in samples_to_infer]
|
||||
wait(futures, return_when=ALL_COMPLETED)
|
||||
|
||||
next_infer_step += 1
|
||||
samples_to_infer = _increment_tree_idx_depth(samples_to_infer, next_infer_step)
|
||||
|
||||
# flatten finished outputs
|
||||
return [traj for lst in finished_rollout_by_root.values() for traj in lst]
|
||||
|
||||
def step(self, sample: DataSampleTree, **kwargs):
|
||||
"""
|
||||
You need to rewrite or modify this method to customize the next round of prompts, such as tools call.
|
||||
"""
|
||||
|
||||
# Special handling has already been done in the rollback.
|
||||
if sample.status == SampleStatus.ROLLBACK:
|
||||
sample.status = SampleStatus.TO_INFER
|
||||
return
|
||||
elif sample.status == SampleStatus.FINISH_NEXT_INFER:
|
||||
prompt = 'In this round of responses, you must generate an answer.'
|
||||
else:
|
||||
prompt = 'The answer is not correct, It seems You made a mistake, you need to recheck very carefully.'
|
||||
|
||||
sample.messages.append({'role': 'user', 'content': prompt})
|
||||
|
||||
def check_finished(self, sample: DataSampleTree, output: ChatCompletionResponse, **kwargs) -> bool:
|
||||
"""
|
||||
Rewrite this method to add custom check logic
|
||||
"""
|
||||
|
||||
boxed_answer = extract_last_boxed(output.choices[0].message.content)
|
||||
|
||||
if boxed_answer is not None:
|
||||
sample.status = SampleStatus.FINISHED
|
||||
sample.finished_reason = FinishedReason.ANSWER
|
||||
|
||||
elif sample.status == SampleStatus.FINISH_NEXT_INFER:
|
||||
sample.status = SampleStatus.FINISHED
|
||||
sample.finished_reason = FinishedReason.MAX_INFER_STEP
|
||||
|
||||
elif sample.depth >= self.max_tree_depth - 1:
|
||||
sample.status = SampleStatus.FINISH_NEXT_INFER
|
||||
|
||||
return sample.status == SampleStatus.FINISHED
|
||||
|
||||
def roll_back_to_divergence(
|
||||
self,
|
||||
finished_samples: Dict[int, List[DataSampleTree]],
|
||||
) -> List[DataSampleTree]:
|
||||
"""
|
||||
All nodes have completed inference, but there is still budget available, rollback.
|
||||
"""
|
||||
|
||||
sample_to_infer = []
|
||||
for root_idx, sample_list in finished_samples.items():
|
||||
if len(sample_list) >= self.max_tree_width:
|
||||
continue
|
||||
|
||||
diff_count = self.max_tree_width - len(sample_list)
|
||||
result = random.sample(sample_list, min(diff_count, len(sample_list)))
|
||||
|
||||
result_copy = deepcopy(result)
|
||||
|
||||
# Randomly rollback several inference iterations; The rollback strategy can be optimized subsequently.
|
||||
for sample in result_copy:
|
||||
sample.status = SampleStatus.ROLLBACK
|
||||
truncate_len = sample.response_num
|
||||
sample.response_truncate(random.randint(1, truncate_len))
|
||||
|
||||
sample_to_infer.extend(result_copy)
|
||||
|
||||
return sample_to_infer
|
||||
|
||||
|
||||
multi_turns['tree_rollout_scheduler'] = TreeRolloutScheduler
|
||||
Reference in New Issue
Block a user