Selaa lähdekoodia

换用github jwjohns/unsloth-GRPO-qwen2.5 验证GRPO训练模型

zhouyang.xie 2 kuukautta sitten
vanhempi
commit
e0482cfcbb
1 muutettua tiedostoa jossa 218 lisäystä ja 36 poistoa
  1. 218 36
      src/train_model_grpo_v1.2.py

+ 218 - 36
src/train_model_grpo_v1.2.py

@@ -9,8 +9,11 @@ from trl import GRPOConfig, GRPOTrainer
 from datasets import load_dataset
 from conf_train import Config, load_config  # 导入配置文件
 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 json
 
 class ModelTrainer:
     def __init__(self, config: Config):
@@ -26,6 +29,9 @@ class ModelTrainer:
         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):
         """
@@ -35,7 +41,7 @@ class ModelTrainer:
         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
+            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,
@@ -53,12 +59,13 @@ class ModelTrainer:
         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
+            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
+                          "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
+            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
@@ -73,11 +80,15 @@ class ModelTrainer:
         :param train_data_path: 训练数据路径
         :return: 返回加载的训练数据集
         """
+
+        data = []
         with open(train_data_path, 'r') as f:
-            train_dataset = load_dataset("json", data_files={"train": train_data_path}, split="train")
-            print("train_dataset -->\n", train_dataset)
-            # 打印第一条数据,检查格式是否正确
-            print("First example in train_dataset:", train_dataset[0])
+            for line in f:
+                data.append(json.loads(line))
+        
+        # 将列表转换为 HuggingFace Dataset 对象
+        data = Dataset.from_list(data)
+
         return train_dataset
 
     def train(self, model, tokenizer, train_dataset):
@@ -123,18 +134,27 @@ class ModelTrainer:
             output_dir=self.config.output_dir,
         )
 
+        """
+        PyTorch 的分布式进程组已初始化,但并行模式不等于 “分布式并行模式(ParallelMode.DISTRIBUTED)”。
+        为了使用 PyTorch 的分布式数据并行(DDP),请使用 python -m torch.distributed.launch 来启动你的脚本。
+        """
+
         trainer = GRPOTrainer(
             model=model,
-            processing_class=tokenizer,  # 用于处理输入文本的分词器(tokenizer)
+            processing_class=tokenizer, # 用于处理输入文本的分词器(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=train_dataset,  # 训练数据集
+                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()
@@ -147,18 +167,130 @@ class ModelTrainer:
         :param tokenizer: 分词器
         :param save_path: 保存路径
         """
+        """
+        # 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 ;
+        """
+
         model.save_pretrained(save_path)
         tokenizer.save_pretrained(save_path)
+        
         print(f"Model saved to {save_path}")
+    
+    @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 semantic_correctness_reward_func(self, prompts, completions, answer, **kwargs):
+        """
+        使用 Longformer 计算生成答案与标准答案的语义相似度。
+        """
+        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)
+
+        # 动态调整权重
+        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):
+        """
+        检查推理过程的质量。
+        :param completions: 模型生成的补全内容
+        :return: 推理过程质量得分列表
+        """
+        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
 
     @staticmethod
     def extract_xml_answer(text: str) -> str:
+        """
+        从文本中提取 XML 格式的答案。
+        :param text: 包含 XML 格式的文本
+        :return: 提取的答案
+        """
         answer = text.split("<answer>")[-1]
         answer = answer.split("</answer>")[0]
         return answer.strip()
-    
+
     @staticmethod
     def count_xml(text) -> float:
+        """
+        计算 XML 标签的数量和完整性。
+        :param text: 包含 XML 格式的文本
+        :return: XML 标签的完整性得分
+        """
         count = 0.0
         if text.count("<reasoning>\n") == 1:
             count += 0.125
@@ -166,17 +298,18 @@ 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):
         """
-        Reward function that counts XML tags in the completion.
+        计算 XML 标签的完整性得分。
+        :param completions: 模型生成的补全内容
+        :return: XML 标签的完整性得分列表
         """
         contents = [completion[0]["content"] for completion in completions]
         return [ModelTrainer.count_xml(c) for c in contents]
@@ -184,7 +317,9 @@ class ModelTrainer:
     @staticmethod
     def soft_format_reward_func(completions, **kwargs):
         """
-        Reward function that checks if the completion has a specific format.
+        检查补全内容是否符合软格式要求。
+        :param completions: 模型生成的补全内容
+        :return: 符合软格式要求的得分列表
         """
         pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
         responses = [completion[0]["content"] for completion in completions]
@@ -194,17 +329,33 @@ class ModelTrainer:
     @staticmethod
     def strict_format_reward_func(completions, **kwargs):
         """
-        Reward function that checks if the completion has a specific format.
+        检查响应是否符合严格的 XML 格式要求,并确保标签内容非空。
+        :param completions: 模型生成的补全内容
+        :return: 符合严格格式要求的得分列表
         """
-        pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"
+        pattern = r"^<reasoning>\n(.+?)\n</reasoning>\n<answer>\n(.+?)\n</answer>\n$"
         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]
+        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):
         """
-        Reward function that checks if the completion contains an integer.
+        检查补全内容是否包含整数。
+        :param completions: 模型生成的补全内容
+        :return: 包含整数的得分列表
         """
         responses = [completion[0]['content'] for completion in completions]
         extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
@@ -213,19 +364,47 @@ class ModelTrainer:
     @staticmethod
     def correctness_reward_func(prompts, completions, answer, **kwargs):
         """
-        Reward function that checks if the completion matches the correct answer.
+        检查补全内容是否正确。
+        :param prompts: 输入提示
+        :param completions: 模型生成的补全内容
+        :param answer: 正确答案
+        :return: 补全内容正确的得分列表
         """
+        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('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
+        # 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("train config:\n",config)
+
+        # 设置环境变量
+        """
+            # 多机多卡
+            # 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}")
+        """
 
         # 初始化 ModelTrainer
         trainer = ModelTrainer(config)
@@ -241,9 +420,12 @@ if __name__ == "__main__":
 
         # 保存模型
         trainer.save_model(model, tokenizer, config.save_path)
-
+        
         print("Training completed.")
     except Exception as e:
-        print("exception \n ", e)
+        print("exception \n ",e)
     finally:
-        print("end")
+        # # 确保进程组被销毁
+        # if dist.is_initialized():
+        #     dist.destroy_process_group()
+        print("end")