|
|
@@ -1,295 +1,179 @@
|
|
|
-# train_model_grpo_v1.py
|
|
|
-
|
|
|
-import os
|
|
|
-import torch
|
|
|
-import torch.distributed as dist
|
|
|
-from unsloth import FastLanguageModel
|
|
|
-from unsloth import is_bfloat16_supported
|
|
|
-from trl import GRPOConfig, GRPOTrainer
|
|
|
-from datasets import load_dataset
|
|
|
-from conf_train import Config, load_config # 导入配置文件
|
|
|
+#!/usr/bin/env python3
|
|
|
import re
|
|
|
-# from transformers import BertTokenizer, BertModel # 分词模型最大支持 512 个token
|
|
|
-from transformers import LongformerTokenizer, LongformerModel # 分词模型最大支持 4096 个token
|
|
|
-import numpy as np
|
|
|
-from datasets import load_dataset, Dataset
|
|
|
+import torch
|
|
|
+import os
|
|
|
import json
|
|
|
-
|
|
|
-class ModelTrainer:
|
|
|
- def __init__(self, config: Config):
|
|
|
- """
|
|
|
- 初始化 ModelTrainer 类，加载配置参数。
|
|
|
- :param config: 配置对象，包含模型训练所需的参数
|
|
|
- """
|
|
|
- self.config: Config = config
|
|
|
- self.model_name = config.model_name
|
|
|
- self.max_seq_length = config.max_seq_length
|
|
|
- self.dtype = torch.float16 if config.dtype == "float16" else torch.bfloat16
|
|
|
- self.load_in_4bit = config.load_in_4bit
|
|
|
- self.fast_inference = config.fast_inference
|
|
|
- self.lora_rank = config.lora_rank
|
|
|
- self.gpu_memory_utilization = config.gpu_memory_utilization
|
|
|
- # 初始化 BERT 模型和分词器
|
|
|
- self.tokenizer = LongformerTokenizer.from_pretrained(f'../models/allenai/longformer-base-4096')
|
|
|
- self.longformer_model = LongformerModel.from_pretrained(f'../models/allenai/longformer-base-4096')
|
|
|
-
|
|
|
- def load_model(self):
|
|
|
- """
|
|
|
- 加载预训练模型和分词器。
|
|
|
- :return: 返回加载的模型和分词器
|
|
|
- """
|
|
|
- model, tokenizer = FastLanguageModel.from_pretrained(
|
|
|
- model_name=self.model_name,
|
|
|
- max_seq_length=self.max_seq_length,
|
|
|
- load_in_4bit=self.load_in_4bit, # False for LoRA 16bit
|
|
|
- dtype=self.dtype,
|
|
|
- fast_inference=self.fast_inference,
|
|
|
- max_lora_rank=self.lora_rank,
|
|
|
- gpu_memory_utilization=self.gpu_memory_utilization,
|
|
|
- )
|
|
|
-
|
|
|
- model = model.to_empty(device='cuda')
|
|
|
-
|
|
|
- # 初始化模型的权重
|
|
|
- for param in model.parameters():
|
|
|
- if param.is_meta:
|
|
|
- param.data = torch.randn_like(param)
|
|
|
-
|
|
|
- # 添加 LoRA 适配器
|
|
|
- model = FastLanguageModel.get_peft_model(
|
|
|
- model,
|
|
|
- max_seq_length=self.max_seq_length,
|
|
|
- r=self.lora_rank, # Choose any number>0!suggested 8,16,32,64，128
|
|
|
- target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
|
|
|
- "gate_proj", "up_proj", "down_proj"], # Remove QKVO if out of memory
|
|
|
- lora_alpha=16,
|
|
|
- lora_dropout=0, #Supports any, but = 0 is optimized
|
|
|
- bias="none", # Supports any, but = "none" is optimized
|
|
|
- #[NEW]"unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
|
|
- use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context
|
|
|
- random_state=3407,
|
|
|
- use_rslora=False, # We support rank stabilized LoRA
|
|
|
- loftq_config=None, # And LoftQ
|
|
|
- )
|
|
|
-
|
|
|
- return model, tokenizer
|
|
|
-
|
|
|
- def load_data(self, train_data_path):
|
|
|
+from unsloth import FastLanguageModel, PatchFastRL, is_bfloat16_supported
|
|
|
+from datasets import Dataset
|
|
|
+from trl import GRPOConfig, GRPOTrainer
|
|
|
+from vllm import SamplingParams
|
|
|
+import trl.trainer.grpo_trainer
|
|
|
+from conf_train import load_config
|
|
|
+
|
|
|
+class GRPOTrainerWrapper:
|
|
|
+ """
|
|
|
+ Wrapper class for GRPO training with object-oriented design.
|
|
|
+ """
|
|
|
+
|
|
|
+ # Constants
|
|
|
+ SYSTEM_PROMPT = """
|
|
|
+ Respond in the following format:
|
|
|
+ <reasoning>
|
|
|
+ ...
|
|
|
+ </reasoning>
|
|
|
+ <answer>
|
|
|
+ ...
|
|
|
+ </answer>
|
|
|
+ """
|
|
|
+
|
|
|
+ XML_COT_FORMAT = """\
|
|
|
+ <reasoning>
|
|
|
+ {reasoning}
|
|
|
+ </reasoning>
|
|
|
+ <answer>
|
|
|
+ {answer}
|
|
|
+ </answer>
|
|
|
+ """
|
|
|
+
|
|
|
+ def __init__(self, config):
|
|
|
"""
|
|
|
- 加载训练数据集。
|
|
|
- :param train_data_path: 训练数据路径
|
|
|
- :return: 返回加载的训练数据集
|
|
|
+ Initialize the trainer with configuration.
|
|
|
"""
|
|
|
-
|
|
|
- data = []
|
|
|
- with open(train_data_path, 'r') as f:
|
|
|
- for line in f:
|
|
|
- data.append(json.loads(line))
|
|
|
+ self.config = config
|
|
|
+ self.model = None
|
|
|
+ self.tokenizer = None
|
|
|
+ self.trainer = None
|
|
|
|
|
|
- # 将列表转换为 HuggingFace Dataset 对象
|
|
|
- train_dataset = Dataset.from_list(data)
|
|
|
-
|
|
|
- return train_dataset
|
|
|
+ # Enable Unsloth's CLI training metrics visualization
|
|
|
+ os.environ["UNSLOTH_DISPLAY_METRICS"] = "true"
|
|
|
+
|
|
|
+ # Apply patch
|
|
|
+ PatchFastRL("GRPO", FastLanguageModel)
|
|
|
+
|
|
|
+ # Monkey patch the validation in the GRPOTrainer
|
|
|
+ self._patch_grpo_trainer()
|
|
|
+
|
|
|
+ def _patch_grpo_trainer(self):
|
|
|
+ """
|
|
|
+ Monkey patch the GRPOTrainer to bypass the divisibility check.
|
|
|
+ """
|
|
|
+ original_init = trl.trainer.grpo_trainer.GRPOTrainer.__init__
|
|
|
+
|
|
|
+ def patched_init(self, *args, **kwargs):
|
|
|
+ try:
|
|
|
+ original_init(self, *args, **kwargs)
|
|
|
+ except ValueError as e:
|
|
|
+ if "evenly divisible by the number of generations per prompt" in str(e):
|
|
|
+ print("Bypassing TRL's batch divisibility check...")
|
|
|
+ # Continue with initialization despite the error
|
|
|
+ self.args = kwargs.get("args")
|
|
|
+ self.model = kwargs.get("model")
|
|
|
+ self.processing_class = kwargs.get("processing_class")
|
|
|
+ self.reward_funcs = kwargs.get("reward_funcs")
|
|
|
+ self.train_dataset = kwargs.get("train_dataset")
|
|
|
+ # Set up necessary trainer components without the check
|
|
|
+ self._setup_trainer()
|
|
|
+ else:
|
|
|
+ raise e
|
|
|
|
|
|
- def train(self, model, tokenizer, train_dataset):
|
|
|
+ trl.trainer.grpo_trainer.GRPOTrainer.__init__ = patched_init
|
|
|
+
|
|
|
+ def load_model(self):
|
|
|
"""
|
|
|
- 训练模型。
|
|
|
- :param model: 预训练模型
|
|
|
- :param tokenizer: 分词器
|
|
|
- :param train_dataset: 训练数据集
|
|
|
- :return: 返回训练后的模型
|
|
|
+ Load the model and tokenizer.
|
|
|
"""
|
|
|
- print("is_bfloat16_supported()=", is_bfloat16_supported())
|
|
|
- print(f"Reserved memory: {torch.cuda.memory_reserved()}")
|
|
|
- print(f"Allocated memory: {torch.cuda.memory_allocated()}")
|
|
|
-
|
|
|
- train_loader = torch.utils.data.DataLoader(
|
|
|
- train_dataset, batch_size=1, shuffle=True, pin_memory=True
|
|
|
+ self.model, self.tokenizer = FastLanguageModel.from_pretrained(
|
|
|
+ model_name=self.config.model_name,
|
|
|
+ max_seq_length=self.config.max_seq_length,
|
|
|
+ load_in_4bit=self.config.load_in_4bit,
|
|
|
+ fast_inference=self.config.fast_inference,
|
|
|
+ max_lora_rank=self.config.lora_rank,
|
|
|
+ gpu_memory_utilization=self.config.gpu_memory_utilization,
|
|
|
)
|
|
|
|
|
|
- torch.cuda.empty_cache()
|
|
|
- print("self.config.learning_rate=", float(self.config.learning_rate))
|
|
|
- training_args = GRPOConfig(
|
|
|
- use_vllm=self.config.use_vllm,
|
|
|
- learning_rate=float(self.config.learning_rate),
|
|
|
- adam_beta1=self.config.adam_beta1,
|
|
|
- adam_beta2=self.config.adam_beta2,
|
|
|
- weight_decay=self.config.weight_decay,
|
|
|
- warmup_ratio=self.config.warmup_ratio,
|
|
|
- lr_scheduler_type=self.config.lr_scheduler_type,
|
|
|
- optim=self.config.optim,
|
|
|
- logging_steps=self.config.logging_steps,
|
|
|
- bf16=is_bfloat16_supported(),
|
|
|
- fp16=not is_bfloat16_supported(),
|
|
|
- per_device_train_batch_size=self.config.per_device_train_batch_size,
|
|
|
- gradient_accumulation_steps=self.config.gradient_accumulation_steps,
|
|
|
- num_generations=self.config.num_generations,
|
|
|
- max_prompt_length=self.config.max_prompt_length,
|
|
|
- max_completion_length=self.config.max_completion_length,
|
|
|
- num_train_epochs=self.config.num_train_epochs,
|
|
|
- max_steps=self.config.max_steps,
|
|
|
- save_steps=self.config.save_steps,
|
|
|
- max_grad_norm=self.config.max_grad_norm,
|
|
|
- report_to=self.config.report_to,
|
|
|
- output_dir=self.config.output_dir,
|
|
|
+ self.model = FastLanguageModel.get_peft_model(
|
|
|
+ self.model,
|
|
|
+ r=self.config.lora_rank,
|
|
|
+ target_modules=[
|
|
|
+ "q_proj", "k_proj", "v_proj", "o_proj",
|
|
|
+ "gate_proj", "up_proj", "down_proj",
|
|
|
+ ],
|
|
|
+ lora_alpha=self.config.lora_rank,
|
|
|
+ use_gradient_checkpointing="unsloth",
|
|
|
+ random_state=3407,
|
|
|
)
|
|
|
-
|
|
|
+
|
|
|
+ @staticmethod
|
|
|
+ def extract_xml_answer(text: str) -> str:
|
|
|
"""
|
|
|
- PyTorch 的分布式进程组已初始化，但并行模式不等于 “分布式并行模式（ParallelMode.DISTRIBUTED）”。
|
|
|
- 为了使用 PyTorch 的分布式数据并行（DDP），请使用 python -m torch.distributed.launch 来启动你的脚本。
|
|
|
+ Extract answer from XML formatted text.
|
|
|
"""
|
|
|
+ answer = text.split("<answer>")[-1]
|
|
|
+ answer = answer.split("</answer>")[0]
|
|
|
+ return answer.strip()
|
|
|
|
|
|
- trainer = GRPOTrainer(
|
|
|
- model=model,
|
|
|
- processing_class=tokenizer, # 用于处理输入文本的分词器(tokenizer)。它将文本转换为模型可以理解的数字格式。
|
|
|
- reward_funcs=[
|
|
|
- self.xmlcount_reward_func, # 某种特定的基于XML计数的奖励函数
|
|
|
- self.soft_format_reward_func, # 基于软格式的奖励函数。
|
|
|
- self.strict_format_reward_func, # 基于严格格式的奖励函数。
|
|
|
- self.int_reward_func, # 整数奖励函数。
|
|
|
- self.correctness_reward_func, # 基于输出正确性的奖励函数
|
|
|
- ###
|
|
|
- # self.semantic_correctness_reward_func, # 语义正确性奖励函数
|
|
|
- # self.reasoning_quality_reward_func, # 推理质量奖励函数
|
|
|
- # self.combined_reward_func, # combined_reward_func
|
|
|
- ], # 这是一个奖励函数的列表，决定了模型输出的好坏。在GRPO训练中，奖励函数通常用来评估模型输出的质量。
|
|
|
- args=training_args, # 定义的训练超参数。
|
|
|
- train_dataset=train_dataset, # 训练数据集,
|
|
|
- )
|
|
|
-
|
|
|
- trainer.train()
|
|
|
- return model
|
|
|
-
|
|
|
- def save_model(self, model, tokenizer, save_path):
|
|
|
+ @staticmethod
|
|
|
+ def extract_hash_answer(text: str) -> str | None:
|
|
|
"""
|
|
|
- 保存训练后的模型和分词器。
|
|
|
- :param model: 训练后的模型
|
|
|
- :param tokenizer: 分词器
|
|
|
- :param save_path: 保存路径
|
|
|
+ Extract answer from hash formatted text.
|
|
|
"""
|
|
|
+ if "####" not in text:
|
|
|
+ return None
|
|
|
+ return text.split("####")[1].strip()
|
|
|
+
|
|
|
+ def load_dataset(self) -> Dataset:
|
|
|
"""
|
|
|
- # Save to 8bit Q8_0
|
|
|
- if False: model.save_pretrained_gguf("model", tokenizer,)
|
|
|
- # Remember to go to https://huggingface.co/settings/tokens for a token!
|
|
|
- # And change hf to your username!
|
|
|
- if False: model.push_to_hub_gguf("hf/model", tokenizer, token = "")
|
|
|
-
|
|
|
- # Save to 16bit GGUF
|
|
|
- if False: model.save_pretrained_gguf("model", tokenizer, quantization_method = "f16")
|
|
|
- if False: model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "f16", token = "")
|
|
|
-
|
|
|
- # Save to q4_k_m GGUF
|
|
|
- if False: model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m")
|
|
|
- if False: model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "q4_k_m", token = "")
|
|
|
-
|
|
|
- ###
|
|
|
-
|
|
|
- model.save_pretrained_merged("model", tokenizer, save_method = "merged_16bit",) # save_method = "merged_4bit" Merge to 4bit ; save_method = "lora" Just LoRA adapters ;
|
|
|
+ Load and prepare the training dataset.
|
|
|
"""
|
|
|
-
|
|
|
- model.save_pretrained(save_path)
|
|
|
- tokenizer.save_pretrained(save_path)
|
|
|
+ # Read JSONL file
|
|
|
+ data = []
|
|
|
+ with open(self.config.train_data_path, 'r') as f:
|
|
|
+ for line in f:
|
|
|
+ data.append(json.loads(line))
|
|
|
|
|
|
- print(f"Model saved to {save_path}")
|
|
|
+ # Convert list to HuggingFace Dataset object
|
|
|
+ return Dataset.from_list(data)
|
|
|
|
|
|
- @staticmethod
|
|
|
- def cosine_similarity(vec1, vec2):
|
|
|
- """
|
|
|
- 计算两个向量的余弦相似度。
|
|
|
- :param vec1: 第一个向量，形状为 (1, 768)
|
|
|
- :param vec2: 第二个向量，形状为 (1, 768)
|
|
|
- :return: 余弦相似度
|
|
|
- """
|
|
|
- # 将 (1, 768) 的矩阵转换为 (768,) 的一维向量
|
|
|
- vec1 = vec1.squeeze() # 形状从 (1, 768) 变为 (768,)
|
|
|
- vec2 = vec2.squeeze() # 形状从 (1, 768) 变为 (768,)
|
|
|
- # print(f"vec1 shape: {vec1.shape}, vec2 shape: {vec2.shape}")
|
|
|
- # 计算余弦相似度
|
|
|
- return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
|
|
|
+ def correctness_reward_func(self, prompts, completions, answer, **kwargs) -> list[float]:
|
|
|
+ """
|
|
|
+ Reward function for answer correctness.
|
|
|
+ """
|
|
|
+ responses = [completion[0]['content'] for completion in completions]
|
|
|
+ q = prompts[0][-1]['content']
|
|
|
+ extracted_responses = [self.extract_xml_answer(r) for r in responses]
|
|
|
+ print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}",
|
|
|
+ f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
|
|
|
+ return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]
|
|
|
|
|
|
- def semantic_correctness_reward_func(self, prompts, completions, answer, **kwargs):
|
|
|
+ def int_reward_func(self, completions, **kwargs) -> list[float]:
|
|
|
"""
|
|
|
- 使用 Longformer 计算生成答案与标准答案的语义相似度。
|
|
|
+ Reward function for integer answers.
|
|
|
"""
|
|
|
responses = [completion[0]['content'] for completion in completions]
|
|
|
extracted_responses = [self.extract_xml_answer(r) for r in responses]
|
|
|
- scores = []
|
|
|
- for resp, ans in zip(extracted_responses, answer):
|
|
|
- # 截断文本，确保长度不超过 4096
|
|
|
- resp = self.tokenizer.decode(self.tokenizer.encode(resp, truncation=True, max_length=4096))
|
|
|
- ans = self.tokenizer.decode(self.tokenizer.encode(ans, truncation=True, max_length=4096))
|
|
|
- # 编码生成答案和标准答案
|
|
|
- inputs_resp = self.tokenizer(resp, return_tensors='pt', padding=True, truncation=True, max_length=4096)
|
|
|
- inputs_ans = self.tokenizer(ans, return_tensors='pt', padding=True, truncation=True, max_length=4096)
|
|
|
- with torch.no_grad():
|
|
|
- outputs_resp = self.longformer_model(**inputs_resp).last_hidden_state.mean(dim=1) # 形状为 (1, 768)
|
|
|
- outputs_ans = self.longformer_model(**inputs_ans).last_hidden_state.mean(dim=1) # 形状为 (1, 768)
|
|
|
- # 计算余弦相似度
|
|
|
- similarity = self.cosine_similarity(outputs_resp.numpy(), outputs_ans.numpy())
|
|
|
- scores.append(similarity)
|
|
|
- return scores
|
|
|
-
|
|
|
- def combined_reward_func(self, prompts, completions, answer, **kwargs):
|
|
|
- """
|
|
|
- 综合多个奖励函数，动态调整权重。
|
|
|
- :param prompts: 输入提示
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :param answer: 标准答案
|
|
|
- :return: 综合得分列表
|
|
|
- """
|
|
|
- # 计算各奖励函数的得分
|
|
|
- format_score = self.strict_format_reward_func(completions)
|
|
|
- semantic_score = self.semantic_correctness_reward_func(prompts, completions, answer)
|
|
|
- correctness_score = self.correctness_reward_func(prompts, completions, answer)
|
|
|
+ return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]
|
|
|
|
|
|
- # 动态调整权重
|
|
|
- combined_scores = []
|
|
|
- for fs, ss, cs in zip(format_score, semantic_score, correctness_score):
|
|
|
- if cs == 2.0: # 答案完全正确
|
|
|
- combined_scores.append(fs * 0.2 + ss * 0.3 + cs * 0.5)
|
|
|
- else: # 答案不完全正确
|
|
|
- combined_scores.append(fs * 0.4 + ss * 0.4 + cs * 0.2)
|
|
|
- return combined_scores
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def reasoning_quality_reward_func(completions, **kwargs):
|
|
|
+ def strict_format_reward_func(self, completions, **kwargs) -> list[float]:
|
|
|
"""
|
|
|
- 检查推理过程的质量。
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :return: 推理过程质量得分列表
|
|
|
+ Strict format reward function.
|
|
|
"""
|
|
|
+ pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"
|
|
|
responses = [completion[0]["content"] for completion in completions]
|
|
|
- scores = []
|
|
|
- for response in responses:
|
|
|
- reasoning_match = re.search(r"<reasoning>\n(.+?)\n</reasoning>", response, re.DOTALL)
|
|
|
- if reasoning_match:
|
|
|
- reasoning_content = reasoning_match.group(1).strip()
|
|
|
- # 简单检查推理内容是否包含关键词
|
|
|
- if "诊断" in reasoning_content and "原因" in reasoning_content:
|
|
|
- scores.append(1.0)
|
|
|
- else:
|
|
|
- scores.append(0.5)
|
|
|
- else:
|
|
|
- scores.append(0.0)
|
|
|
- return scores
|
|
|
+ matches = [re.match(pattern, r) for r in responses]
|
|
|
+ return [0.5 if match else 0.0 for match in matches]
|
|
|
|
|
|
- @staticmethod
|
|
|
- def extract_xml_answer(text: str) -> str:
|
|
|
+ def soft_format_reward_func(self, completions, **kwargs) -> list[float]:
|
|
|
"""
|
|
|
- 从文本中提取 XML 格式的答案。
|
|
|
- :param text: 包含 XML 格式的文本
|
|
|
- :return: 提取的答案
|
|
|
+ Soft format reward function.
|
|
|
"""
|
|
|
- answer = text.split("<answer>")[-1]
|
|
|
- answer = answer.split("</answer>")[0]
|
|
|
- return answer.strip()
|
|
|
+ pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
|
|
|
+ responses = [completion[0]["content"] for completion in completions]
|
|
|
+ matches = [re.match(pattern, r) for r in responses]
|
|
|
+ return [0.5 if match else 0.0 for match in matches]
|
|
|
|
|
|
- @staticmethod
|
|
|
- def count_xml(text) -> float:
|
|
|
+ def count_xml(self, text) -> float:
|
|
|
"""
|
|
|
- 计算 XML 标签的数量和完整性。
|
|
|
- :param text: 包含 XML 格式的文本
|
|
|
- :return: XML 标签的完整性得分
|
|
|
+ Count XML tags in text.
|
|
|
"""
|
|
|
count = 0.0
|
|
|
if text.count("<reasoning>\n") == 1:
|
|
|
@@ -298,134 +182,104 @@ class ModelTrainer:
|
|
|
count += 0.125
|
|
|
if text.count("\n<answer>\n") == 1:
|
|
|
count += 0.125
|
|
|
- count -= len(text.split("\n</answer>\n")[-1]) * 0.001
|
|
|
+ count -= len(text.split("\n</answer>\n")[-1])*0.001
|
|
|
if text.count("\n</answer>") == 1:
|
|
|
count += 0.125
|
|
|
- count -= (len(text.split("\n</answer>")[-1]) - 1) * 0.001
|
|
|
+ count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001
|
|
|
return count
|
|
|
|
|
|
- @staticmethod
|
|
|
- def xmlcount_reward_func(completions, **kwargs):
|
|
|
+ def xmlcount_reward_func(self, completions, **kwargs) -> list[float]:
|
|
|
"""
|
|
|
- 计算 XML 标签的完整性得分。
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :return: XML 标签的完整性得分列表
|
|
|
+ Reward function based on XML tag count.
|
|
|
"""
|
|
|
contents = [completion[0]["content"] for completion in completions]
|
|
|
- return [ModelTrainer.count_xml(c) for c in contents]
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def soft_format_reward_func(completions, **kwargs):
|
|
|
- """
|
|
|
- 检查补全内容是否符合软格式要求。
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :return: 符合软格式要求的得分列表
|
|
|
- """
|
|
|
- pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
|
|
|
- responses = [completion[0]["content"] for completion in completions]
|
|
|
- matches = [re.match(pattern, r) for r in responses]
|
|
|
- return [0.5 if match else 0.0 for match in matches]
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def strict_format_reward_func(completions, **kwargs):
|
|
|
+ return [self.count_xml(c) for c in contents]
|
|
|
+
|
|
|
+ def prepare_training_args(self) -> GRPOConfig:
|
|
|
"""
|
|
|
- 检查响应是否符合严格的 XML 格式要求，并确保标签内容非空。
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :return: 符合严格格式要求的得分列表
|
|
|
+ Prepare training arguments from config.
|
|
|
"""
|
|
|
- pattern = r"^<reasoning>\n(.+?)\n</reasoning>\n<answer>\n(.+?)\n</answer>\n$"
|
|
|
- responses = [completion[0]["content"] for completion in completions]
|
|
|
- scores = []
|
|
|
- for response in responses:
|
|
|
- match = re.match(pattern, response, re.DOTALL)
|
|
|
- if match:
|
|
|
- reasoning_content = match.group(1).strip()
|
|
|
- answer_content = match.group(2).strip()
|
|
|
- # 检查内容是否非空
|
|
|
- if reasoning_content and answer_content:
|
|
|
- scores.append(1.0) # 格式和内容均符合要求
|
|
|
- else:
|
|
|
- scores.append(0.5) # 格式符合但内容为空
|
|
|
- else:
|
|
|
- scores.append(0.0) # 格式不符合
|
|
|
- return scores
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def int_reward_func(completions, **kwargs):
|
|
|
+ return GRPOConfig(
|
|
|
+ use_vllm=self.config.use_vllm,
|
|
|
+ learning_rate=self.config.learning_rate,
|
|
|
+ adam_beta1=self.config.adam_beta1,
|
|
|
+ adam_beta2=self.config.adam_beta2,
|
|
|
+ weight_decay=self.config.weight_decay,
|
|
|
+ warmup_ratio=self.config.warmup_ratio,
|
|
|
+ lr_scheduler_type=self.config.lr_scheduler_type,
|
|
|
+ optim=self.config.optim,
|
|
|
+ logging_steps=self.config.logging_steps,
|
|
|
+ bf16=is_bfloat16_supported(),
|
|
|
+ fp16=not is_bfloat16_supported(),
|
|
|
+ per_device_train_batch_size=self.config.per_device_train_batch_size,
|
|
|
+ gradient_accumulation_steps=self.config.gradient_accumulation_steps,
|
|
|
+ num_generations=self.config.num_generations,
|
|
|
+ max_prompt_length=self.config.max_prompt_length,
|
|
|
+ max_completion_length=self.config.max_completion_length,
|
|
|
+ max_steps=self.config.max_steps,
|
|
|
+ save_steps=self.config.save_steps,
|
|
|
+ max_grad_norm=self.config.max_grad_norm,
|
|
|
+ report_to=self.config.report_to,
|
|
|
+ output_dir=self.config.output_dir,
|
|
|
+ save_total_limit=3,
|
|
|
+ log_level="info",
|
|
|
+ disable_tqdm=False,
|
|
|
+ evaluation_strategy="no",
|
|
|
+ )
|
|
|
+
|
|
|
+ def initialize_trainer(self, dataset):
|
|
|
"""
|
|
|
- 检查补全内容是否包含整数。
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :return: 包含整数的得分列表
|
|
|
+ Initialize the GRPO trainer.
|
|
|
"""
|
|
|
- responses = [completion[0]['content'] for completion in completions]
|
|
|
- extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
|
|
|
- return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]
|
|
|
-
|
|
|
- @staticmethod
|
|
|
- def correctness_reward_func(prompts, completions, answer, **kwargs):
|
|
|
+ training_args = self.prepare_training_args()
|
|
|
+
|
|
|
+ self.trainer = GRPOTrainer(
|
|
|
+ model=self.model,
|
|
|
+ processing_class=self.tokenizer,
|
|
|
+ reward_funcs=[
|
|
|
+ self.xmlcount_reward_func,
|
|
|
+ self.soft_format_reward_func,
|
|
|
+ self.strict_format_reward_func,
|
|
|
+ self.int_reward_func,
|
|
|
+ self.correctness_reward_func,
|
|
|
+ ],
|
|
|
+ args=training_args,
|
|
|
+ train_dataset=dataset,
|
|
|
+ )
|
|
|
+
|
|
|
+ def train(self):
|
|
|
"""
|
|
|
- 检查补全内容是否正确。
|
|
|
- :param prompts: 输入提示
|
|
|
- :param completions: 模型生成的补全内容
|
|
|
- :param answer: 正确答案
|
|
|
- :return: 补全内容正确的得分列表
|
|
|
+ Execute the training process.
|
|
|
"""
|
|
|
- print("completions : \n ",completions)
|
|
|
- responses = [completion[0]['content'] for completion in completions]
|
|
|
- q = prompts[0][-1]['content']
|
|
|
- extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
|
|
|
- # print(f"\n Response:\n {responses}",f"\n Extracted:\n {responses}")
|
|
|
- print('-' * 20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
|
|
|
- return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]
|
|
|
-
|
|
|
-if __name__ == "__main__":
|
|
|
- try:
|
|
|
- # 加载配置文件
|
|
|
- config = load_config(f"../conf/conf_train.yaml")
|
|
|
- print("train config:\n",config)
|
|
|
-
|
|
|
- # 设置环境变量
|
|
|
+ print("Starting GRPO training...")
|
|
|
+ print(f"per_device_train_batch_size = {self.config.per_device_train_batch_size}")
|
|
|
+ print(f"gradient_accumulation_steps = {self.config.gradient_accumulation_steps}")
|
|
|
+ print(f"num_generations = {self.config.num_generations}")
|
|
|
+ print(f"max_steps = {self.config.max_steps}")
|
|
|
+
|
|
|
+ self.trainer.train()
|
|
|
+
|
|
|
+ def save_model(self):
|
|
|
"""
|
|
|
- # 多机多卡
|
|
|
- # export RANK=0 # 第一台机器的 rank
|
|
|
- # export WORLD_SIZE=4 # 总共有 4 台机器
|
|
|
- # export MASTER_ADDR=<主节点 IP>
|
|
|
- # export MASTER_PORT=12345
|
|
|
-
|
|
|
- # 单机多卡
|
|
|
- os.environ['RANK'] = '0'
|
|
|
- os.environ['WORLD_SIZE'] = '1'
|
|
|
- os.environ['MASTER_ADDR'] = 'localhost'
|
|
|
- os.environ['MASTER_PORT'] = '12345'
|
|
|
-
|
|
|
- # 根据操作系统选择后端
|
|
|
- backend = 'gloo' if os.name == 'nt' else 'nccl'
|
|
|
- # 使用文件初始化方法 2025-3-11 成功验证支持windows
|
|
|
- init_method = f'env://' # env:// # 文件路径需要所有进程都能访问
|
|
|
- dist.init_process_group(backend=backend, init_method=init_method)
|
|
|
- print(f"Initialized distributed training with backend: {backend}")
|
|
|
+ Save the trained model.
|
|
|
"""
|
|
|
+ print(f"Saving model to {self.config.save_path}...")
|
|
|
+ os.makedirs(os.path.dirname(self.config.save_path), exist_ok=True)
|
|
|
+ self.model.save_pretrained(self.config.save_path)
|
|
|
+ self.tokenizer.save_pretrained(self.config.save_path)
|
|
|
+ print("Training complete!")
|
|
|
|
|
|
- # 初始化 ModelTrainer
|
|
|
- trainer = ModelTrainer(config)
|
|
|
-
|
|
|
- # 加载模型和分词器
|
|
|
- model, tokenizer = trainer.load_model()
|
|
|
-
|
|
|
- # 加载数据集
|
|
|
- train_dataset = trainer.load_data(config.train_data_path)
|
|
|
-
|
|
|
- # 训练模型
|
|
|
- model = trainer.train(model, tokenizer, train_dataset)
|
|
|
+def main():
|
|
|
+ # Load configuration
|
|
|
+ config = load_config()
|
|
|
+
|
|
|
+ # Initialize and run trainer
|
|
|
+ trainer = GRPOTrainerWrapper(config)
|
|
|
+ trainer.load_model()
|
|
|
+ dataset = trainer.load_dataset()
|
|
|
+ trainer.initialize_trainer(dataset)
|
|
|
+ trainer.train()
|
|
|
+ trainer.save_model()
|
|
|
|
|
|
- # 保存模型
|
|
|
- trainer.save_model(model, tokenizer, config.save_path)
|
|
|
-
|
|
|
- print("Training completed.")
|
|
|
- except Exception as e:
|
|
|
- print("exception \n ",e)
|
|
|
- finally:
|
|
|
- # # 确保进程组被销毁
|
|
|
- # if dist.is_initialized():
|
|
|
- # dist.destroy_process_group()
|
|
|
- print("end")
|
|
|
+if __name__ == "__main__":
|
|
|
+ main()
|
zhouyang.xie