train_model_grpo_v1.1.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422
  1. # train_model_grpo_v1.py
  2. import os
  3. import torch
  4. import torch.distributed as dist
  5. from unsloth import FastLanguageModel
  6. from unsloth import is_bfloat16_supported
  7. from trl import GRPOConfig, GRPOTrainer
  8. from datasets import load_dataset
  9. from conf_train import Config, load_config # 导入配置文件
  10. import re
  11. # from transformers import BertTokenizer, BertModel # 分词模型最大支持 512 个token
  12. from transformers import LongformerTokenizer, LongformerModel # 分词模型最大支持 4096 个token
  13. import numpy as np
  14. class ModelTrainer:
  15. def __init__(self, config: Config):
  16. """
  17. 初始化 ModelTrainer 类,加载配置参数。
  18. :param config: 配置对象,包含模型训练所需的参数
  19. """
  20. self.config: Config = config
  21. self.model_name = config.model_name
  22. self.max_seq_length = config.max_seq_length
  23. self.dtype = torch.float16 if config.dtype == "float16" else torch.bfloat16
  24. self.load_in_4bit = config.load_in_4bit
  25. self.fast_inference = config.fast_inference
  26. self.lora_rank = config.lora_rank
  27. self.gpu_memory_utilization = config.gpu_memory_utilization
  28. # 初始化 BERT 模型和分词器
  29. self.tokenizer = LongformerTokenizer.from_pretrained(f'../models/allenai/longformer-base-4096')
  30. self.longformer_model = LongformerModel.from_pretrained(f'../models/allenai/longformer-base-4096')
  31. def load_model(self):
  32. """
  33. 加载预训练模型和分词器。
  34. :return: 返回加载的模型和分词器
  35. """
  36. model, tokenizer = FastLanguageModel.from_pretrained(
  37. model_name=self.model_name,
  38. max_seq_length=self.max_seq_length,
  39. load_in_4bit=self.load_in_4bit, # False for LoRA 16bit
  40. dtype=self.dtype,
  41. fast_inference=self.fast_inference,
  42. max_lora_rank=self.lora_rank,
  43. gpu_memory_utilization=self.gpu_memory_utilization,
  44. )
  45. model = model.to_empty(device='cuda')
  46. # 初始化模型的权重
  47. for param in model.parameters():
  48. if param.is_meta:
  49. param.data = torch.randn_like(param)
  50. # 添加 LoRA 适配器
  51. model = FastLanguageModel.get_peft_model(
  52. model,
  53. max_seq_length=self.max_seq_length,
  54. r=self.lora_rank, # Choose any number>0!suggested 8,16,32,64,128
  55. target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
  56. "gate_proj", "up_proj", "down_proj"], # Remove QKVO if out of memory
  57. lora_alpha=16,
  58. lora_dropout=0, #Supports any, but = 0 is optimized
  59. bias="none", # Supports any, but = "none" is optimized
  60. #[NEW]"unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
  61. use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context
  62. random_state=3407,
  63. use_rslora=False, # We support rank stabilized LoRA
  64. loftq_config=None, # And LoftQ
  65. )
  66. return model, tokenizer
  67. def load_data(self, train_data_path):
  68. """
  69. 加载训练数据集。
  70. :param train_data_path: 训练数据路径
  71. :return: 返回加载的训练数据集
  72. """
  73. with open(train_data_path, 'r') as f:
  74. train_dataset = load_dataset("json", data_files={"train": train_data_path}, split="train")
  75. print("train_dataset",train_dataset)
  76. return train_dataset
  77. def train(self, model, tokenizer, train_dataset):
  78. """
  79. 训练模型。
  80. :param model: 预训练模型
  81. :param tokenizer: 分词器
  82. :param train_dataset: 训练数据集
  83. :return: 返回训练后的模型
  84. """
  85. print("is_bfloat16_supported()=", is_bfloat16_supported())
  86. print(f"Reserved memory: {torch.cuda.memory_reserved()}")
  87. print(f"Allocated memory: {torch.cuda.memory_allocated()}")
  88. train_loader = torch.utils.data.DataLoader(
  89. train_dataset, batch_size=1, shuffle=True, pin_memory=True
  90. )
  91. torch.cuda.empty_cache()
  92. print("self.config.learning_rate=", float(self.config.learning_rate))
  93. training_args = GRPOConfig(
  94. use_vllm=self.config.use_vllm,
  95. learning_rate=float(self.config.learning_rate),
  96. adam_beta1=self.config.adam_beta1,
  97. adam_beta2=self.config.adam_beta2,
  98. weight_decay=self.config.weight_decay,
  99. warmup_ratio=self.config.warmup_ratio,
  100. lr_scheduler_type=self.config.lr_scheduler_type,
  101. optim=self.config.optim,
  102. logging_steps=self.config.logging_steps,
  103. bf16=is_bfloat16_supported(),
  104. fp16=not is_bfloat16_supported(),
  105. per_device_train_batch_size=self.config.per_device_train_batch_size,
  106. gradient_accumulation_steps=self.config.gradient_accumulation_steps,
  107. num_generations=self.config.num_generations,
  108. max_prompt_length=self.config.max_prompt_length,
  109. max_completion_length=self.config.max_completion_length,
  110. num_train_epochs=self.config.num_train_epochs,
  111. max_steps=self.config.max_steps,
  112. save_steps=self.config.save_steps,
  113. max_grad_norm=self.config.max_grad_norm,
  114. report_to=self.config.report_to,
  115. output_dir=self.config.output_dir,
  116. )
  117. """
  118. PyTorch 的分布式进程组已初始化,但并行模式不等于 “分布式并行模式(ParallelMode.DISTRIBUTED)”。
  119. 为了使用 PyTorch 的分布式数据并行(DDP),请使用 python -m torch.distributed.launch 来启动你的脚本。
  120. """
  121. trainer = GRPOTrainer(
  122. model=model,
  123. processing_class=tokenizer, # 用于处理输入文本的分词器(tokenizer)。它将文本转换为模型可以理解的数字格式。
  124. reward_funcs=[
  125. self.xmlcount_reward_func, # 某种特定的基于XML计数的奖励函数
  126. self.soft_format_reward_func, # 基于软格式的奖励函数。
  127. self.strict_format_reward_func, # 基于严格格式的奖励函数。
  128. self.int_reward_func, # 整数奖励函数。
  129. self.correctness_reward_func, # 基于输出正确性的奖励函数
  130. self.semantic_correctness_reward_func, # 语义正确性奖励函数
  131. self.reasoning_quality_reward_func, # 推理质量奖励函数
  132. self.combined_reward_func, # combined_reward_func
  133. ], # 这是一个奖励函数的列表,决定了模型输出的好坏。在GRPO训练中,奖励函数通常用来评估模型输出的质量。
  134. args=training_args, # 定义的训练超参数。
  135. train_dataset=train_dataset, # 训练数据集,
  136. )
  137. trainer.train()
  138. return model
  139. def save_model(self, model, tokenizer, save_path):
  140. """
  141. 保存训练后的模型和分词器。
  142. :param model: 训练后的模型
  143. :param tokenizer: 分词器
  144. :param save_path: 保存路径
  145. """
  146. """
  147. # Save to 8bit Q8_0
  148. if False: model.save_pretrained_gguf("model", tokenizer,)
  149. # Remember to go to https://huggingface.co/settings/tokens for a token!
  150. # And change hf to your username!
  151. if False: model.push_to_hub_gguf("hf/model", tokenizer, token = "")
  152. # Save to 16bit GGUF
  153. if False: model.save_pretrained_gguf("model", tokenizer, quantization_method = "f16")
  154. if False: model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "f16", token = "")
  155. # Save to q4_k_m GGUF
  156. if False: model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m")
  157. if False: model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "q4_k_m", token = "")
  158. ###
  159. model.save_pretrained_merged("model", tokenizer, save_method = "merged_16bit",) # save_method = "merged_4bit" Merge to 4bit ; save_method = "lora" Just LoRA adapters ;
  160. """
  161. model.save_pretrained(save_path)
  162. tokenizer.save_pretrained(save_path)
  163. print(f"Model saved to {save_path}")
  164. @staticmethod
  165. def cosine_similarity(vec1, vec2):
  166. """
  167. 计算两个向量的余弦相似度。
  168. :param vec1: 第一个向量,形状为 (1, 768)
  169. :param vec2: 第二个向量,形状为 (1, 768)
  170. :return: 余弦相似度
  171. """
  172. # 将 (1, 768) 的矩阵转换为 (768,) 的一维向量
  173. vec1 = vec1.squeeze() # 形状从 (1, 768) 变为 (768,)
  174. vec2 = vec2.squeeze() # 形状从 (1, 768) 变为 (768,)
  175. # print(f"vec1 shape: {vec1.shape}, vec2 shape: {vec2.shape}")
  176. # 计算余弦相似度
  177. return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
  178. def semantic_correctness_reward_func(self, prompts, completions, answer, **kwargs):
  179. """
  180. 使用 Longformer 计算生成答案与标准答案的语义相似度。
  181. """
  182. responses = [completion[0]['content'] for completion in completions]
  183. extracted_responses = [self.extract_xml_answer(r) for r in responses]
  184. scores = []
  185. for resp, ans in zip(extracted_responses, answer):
  186. # 截断文本,确保长度不超过 4096
  187. resp = self.tokenizer.decode(self.tokenizer.encode(resp, truncation=True, max_length=4096))
  188. ans = self.tokenizer.decode(self.tokenizer.encode(ans, truncation=True, max_length=4096))
  189. # 编码生成答案和标准答案
  190. inputs_resp = self.tokenizer(resp, return_tensors='pt', padding=True, truncation=True, max_length=4096)
  191. inputs_ans = self.tokenizer(ans, return_tensors='pt', padding=True, truncation=True, max_length=4096)
  192. with torch.no_grad():
  193. outputs_resp = self.longformer_model(**inputs_resp).last_hidden_state.mean(dim=1) # 形状为 (1, 768)
  194. outputs_ans = self.longformer_model(**inputs_ans).last_hidden_state.mean(dim=1) # 形状为 (1, 768)
  195. # 计算余弦相似度
  196. similarity = self.cosine_similarity(outputs_resp.numpy(), outputs_ans.numpy())
  197. scores.append(similarity)
  198. return scores
  199. def combined_reward_func(self, prompts, completions, answer, **kwargs):
  200. """
  201. 综合多个奖励函数,动态调整权重。
  202. :param prompts: 输入提示
  203. :param completions: 模型生成的补全内容
  204. :param answer: 标准答案
  205. :return: 综合得分列表
  206. """
  207. # 计算各奖励函数的得分
  208. format_score = self.strict_format_reward_func(completions)
  209. semantic_score = self.semantic_correctness_reward_func(prompts, completions, answer)
  210. correctness_score = self.correctness_reward_func(prompts, completions, answer)
  211. # 动态调整权重
  212. combined_scores = []
  213. for fs, ss, cs in zip(format_score, semantic_score, correctness_score):
  214. if cs == 2.0: # 答案完全正确
  215. combined_scores.append(fs * 0.2 + ss * 0.3 + cs * 0.5)
  216. else: # 答案不完全正确
  217. combined_scores.append(fs * 0.4 + ss * 0.4 + cs * 0.2)
  218. return combined_scores
  219. @staticmethod
  220. def reasoning_quality_reward_func(completions, **kwargs):
  221. """
  222. 检查推理过程的质量。
  223. :param completions: 模型生成的补全内容
  224. :return: 推理过程质量得分列表
  225. """
  226. responses = [completion[0]["content"] for completion in completions]
  227. scores = []
  228. for response in responses:
  229. reasoning_match = re.search(r"<reasoning>\n(.+?)\n</reasoning>", response, re.DOTALL)
  230. if reasoning_match:
  231. reasoning_content = reasoning_match.group(1).strip()
  232. # 简单检查推理内容是否包含关键词
  233. if "诊断" in reasoning_content and "原因" in reasoning_content:
  234. scores.append(1.0)
  235. else:
  236. scores.append(0.5)
  237. else:
  238. scores.append(0.0)
  239. return scores
  240. @staticmethod
  241. def extract_xml_answer(text: str) -> str:
  242. """
  243. 从文本中提取 XML 格式的答案。
  244. :param text: 包含 XML 格式的文本
  245. :return: 提取的答案
  246. """
  247. answer = text.split("<answer>")[-1]
  248. answer = answer.split("</answer>")[0]
  249. return answer.strip()
  250. @staticmethod
  251. def count_xml(text) -> float:
  252. """
  253. 计算 XML 标签的数量和完整性。
  254. :param text: 包含 XML 格式的文本
  255. :return: XML 标签的完整性得分
  256. """
  257. count = 0.0
  258. if text.count("<reasoning>\n") == 1:
  259. count += 0.125
  260. if text.count("\n</reasoning>\n") == 1:
  261. count += 0.125
  262. if text.count("\n<answer>\n") == 1:
  263. count += 0.125
  264. count -= len(text.split("\n</answer>\n")[-1]) * 0.001
  265. if text.count("\n</answer>") == 1:
  266. count += 0.125
  267. count -= (len(text.split("\n</answer>")[-1]) - 1) * 0.001
  268. return count
  269. @staticmethod
  270. def xmlcount_reward_func(completions, **kwargs):
  271. """
  272. 计算 XML 标签的完整性得分。
  273. :param completions: 模型生成的补全内容
  274. :return: XML 标签的完整性得分列表
  275. """
  276. contents = [completion[0]["content"] for completion in completions]
  277. return [ModelTrainer.count_xml(c) for c in contents]
  278. @staticmethod
  279. def soft_format_reward_func(completions, **kwargs):
  280. """
  281. 检查补全内容是否符合软格式要求。
  282. :param completions: 模型生成的补全内容
  283. :return: 符合软格式要求的得分列表
  284. """
  285. pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
  286. responses = [completion[0]["content"] for completion in completions]
  287. matches = [re.match(pattern, r) for r in responses]
  288. return [0.5 if match else 0.0 for match in matches]
  289. @staticmethod
  290. def strict_format_reward_func(completions, **kwargs):
  291. """
  292. 检查响应是否符合严格的 XML 格式要求,并确保标签内容非空。
  293. :param completions: 模型生成的补全内容
  294. :return: 符合严格格式要求的得分列表
  295. """
  296. pattern = r"^<reasoning>\n(.+?)\n</reasoning>\n<answer>\n(.+?)\n</answer>\n$"
  297. responses = [completion[0]["content"] for completion in completions]
  298. scores = []
  299. for response in responses:
  300. match = re.match(pattern, response, re.DOTALL)
  301. if match:
  302. reasoning_content = match.group(1).strip()
  303. answer_content = match.group(2).strip()
  304. # 检查内容是否非空
  305. if reasoning_content and answer_content:
  306. scores.append(1.0) # 格式和内容均符合要求
  307. else:
  308. scores.append(0.5) # 格式符合但内容为空
  309. else:
  310. scores.append(0.0) # 格式不符合
  311. return scores
  312. @staticmethod
  313. def int_reward_func(completions, **kwargs):
  314. """
  315. 检查补全内容是否包含整数。
  316. :param completions: 模型生成的补全内容
  317. :return: 包含整数的得分列表
  318. """
  319. responses = [completion[0]['content'] for completion in completions]
  320. extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
  321. return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]
  322. @staticmethod
  323. def correctness_reward_func(prompts, completions, answer, **kwargs):
  324. """
  325. 检查补全内容是否正确。
  326. :param prompts: 输入提示
  327. :param completions: 模型生成的补全内容
  328. :param answer: 正确答案
  329. :return: 补全内容正确的得分列表
  330. """
  331. print("completions : \n ",completions)
  332. responses = [completion[0]['content'] for completion in completions]
  333. q = prompts[0][-1]['content']
  334. extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
  335. # print(f"\n Response:\n {responses}",f"\n Extracted:\n {responses}")
  336. print('-' * 20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
  337. return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]
  338. if __name__ == "__main__":
  339. try:
  340. # 加载配置文件
  341. config = load_config(f"../conf/conf_train.yaml")
  342. print("train config:\n",config)
  343. # 设置环境变量
  344. """
  345. # 多机多卡
  346. # export RANK=0 # 第一台机器的 rank
  347. # export WORLD_SIZE=4 # 总共有 4 台机器
  348. # export MASTER_ADDR=<主节点 IP>
  349. # export MASTER_PORT=12345
  350. # 单机多卡
  351. os.environ['RANK'] = '0'
  352. os.environ['WORLD_SIZE'] = '1'
  353. os.environ['MASTER_ADDR'] = 'localhost'
  354. os.environ['MASTER_PORT'] = '12345'
  355. # 根据操作系统选择后端
  356. backend = 'gloo' if os.name == 'nt' else 'nccl'
  357. # 使用文件初始化方法 2025-3-11 成功验证支持windows
  358. init_method = f'env://' # env:// # 文件路径需要所有进程都能访问
  359. dist.init_process_group(backend=backend, init_method=init_method)
  360. print(f"Initialized distributed training with backend: {backend}")
  361. """
  362. # 初始化 ModelTrainer
  363. trainer = ModelTrainer(config)
  364. # 加载模型和分词器
  365. model, tokenizer = trainer.load_model()
  366. # 加载数据集
  367. train_dataset = trainer.load_data(config.train_data_path)
  368. # 训练模型
  369. model = trainer.train(model, tokenizer, train_dataset)
  370. # 保存模型
  371. trainer.save_model(model, tokenizer, config.save_path)
  372. print("Training completed.")
  373. except Exception as e:
  374. print("exception \n ",e)
  375. finally:
  376. # # 确保进程组被销毁
  377. # if dist.is_initialized():
  378. # dist.destroy_process_group()
  379. print("end")