train_model_grpo_v1.py 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257
  1. import os
  2. import torch
  3. import torch.distributed as dist
  4. from unsloth import FastLanguageModel
  5. from unsloth import is_bfloat16_supported
  6. from trl import GRPOConfig, GRPOTrainer
  7. from datasets import load_dataset
  8. from conf_train import Config ,load_config # 导入配置文件
  9. import re
  10. class ModelTrainer:
  11. def __init__(self, config:Config):
  12. """
  13. 初始化 ModelTrainer 类,加载配置参数。
  14. :param config: 配置对象,包含模型训练所需的参数
  15. """
  16. self.config = config
  17. self.model_name = config.model_name
  18. self.max_seq_length = config.max_seq_length
  19. self.dtype = torch.float16 if config.dtype == "float16" else torch.bfloat16
  20. self.load_in_4bit = config.load_in_4bit
  21. self.lora_rank = config.lora_rank
  22. self.gpu_memory_utilization=config.gpu_memory_utilization
  23. def load_model(self):
  24. """
  25. 加载预训练模型和分词器。
  26. :return: 返回加载的模型和分词器
  27. """
  28. model, tokenizer = FastLanguageModel.from_pretrained(
  29. model_name=self.model_name,
  30. max_seq_length=self.max_seq_length,
  31. load_in_4bit=self.load_in_4bit,
  32. dtype=self.dtype,
  33. fast_inference=False,
  34. max_lora_rank=self.lora_rank,
  35. gpu_memory_utilization=0.6,
  36. )
  37. model = model.to_empty(device='cuda')
  38. # 初始化模型的权重
  39. for param in model.parameters():
  40. if param.is_meta:
  41. param.data = torch.randn_like(param)
  42. # 添加 LoRA 适配器
  43. model = FastLanguageModel.get_peft_model(
  44. model,
  45. max_seq_length=self.max_seq_length,
  46. r=self.lora_rank,
  47. target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
  48. "gate_proj", "up_proj", "down_proj"],
  49. lora_alpha=16,
  50. lora_dropout=0,
  51. bias="none",
  52. use_gradient_checkpointing="unsloth",
  53. random_state=3407,
  54. use_rslora=False,
  55. loftq_config=None,
  56. )
  57. return model, tokenizer
  58. def load_data(self, train_data_path):
  59. """
  60. 加载训练数据集。
  61. :param train_data_path: 训练数据路径
  62. :return: 返回加载的训练数据集
  63. """
  64. with open(train_data_path, 'r') as f:
  65. train_dataset = load_dataset("json", data_files={"train": train_data_path}, split="train")
  66. return train_dataset
  67. def train(self, model, tokenizer, train_dataset):
  68. """
  69. 训练模型。
  70. :param model: 预训练模型
  71. :param tokenizer: 分词器
  72. :param train_dataset: 训练数据集
  73. :return: 返回训练后的模型
  74. """
  75. print("is_bfloat16_supported()=", is_bfloat16_supported())
  76. print(f"Reserved memory: {torch.cuda.memory_reserved()}")
  77. print(f"Allocated memory: {torch.cuda.memory_allocated()}")
  78. train_loader = torch.utils.data.DataLoader(
  79. train_dataset, batch_size=1, shuffle=True, pin_memory=True
  80. )
  81. torch.cuda.empty_cache()
  82. training_args = GRPOConfig(
  83. use_vllm=False,
  84. learning_rate=self.config.learning_rate,
  85. adam_beta1=self.config.adam_beta1,
  86. adam_beta2=self.config.adam_beta2,
  87. weight_decay=self.config.weight_decay,
  88. warmup_ratio=self.config.warmup_ratio,
  89. lr_scheduler_type=self.config.lr_scheduler_type,
  90. optim=self.config.optim,
  91. logging_steps=self.config.logging_steps,
  92. bf16=is_bfloat16_supported(),
  93. fp16=not is_bfloat16_supported(),
  94. per_device_train_batch_size=self.config.per_device_train_batch_size,
  95. gradient_accumulation_steps=self.config.gradient_accumulation_steps,
  96. num_generations=self.config.num_generations,
  97. max_prompt_length=self.config.max_prompt_length,
  98. max_completion_length=self.config.max_completion_length,
  99. num_train_epochs=self.config.num_train_epochs,
  100. max_steps=self.config.max_steps,
  101. save_steps=self.config.save_steps,
  102. max_grad_norm=self.config.max_grad_norm,
  103. report_to=self.config.report_to,
  104. output_dir=self.config.output_dir,
  105. )
  106. trainer = GRPOTrainer(
  107. model=model,
  108. processing_class=tokenizer,
  109. reward_funcs=[
  110. self.xmlcount_reward_func,
  111. self.soft_format_reward_func,
  112. self.strict_format_reward_func,
  113. self.int_reward_func,
  114. self.correctness_reward_func,
  115. ],
  116. args=training_args,
  117. train_dataset=train_dataset,
  118. )
  119. trainer.train()
  120. return model
  121. def save_model(self, model, tokenizer, save_path):
  122. """
  123. 保存训练后的模型和分词器。
  124. :param model: 训练后的模型
  125. :param tokenizer: 分词器
  126. :param save_path: 保存路径
  127. """
  128. model.save_pretrained(save_path)
  129. tokenizer.save_pretrained(save_path)
  130. print(f"Model saved to {save_path}")
  131. @staticmethod
  132. def extract_xml_answer(text: str) -> str:
  133. answer = text.split("<answer>")[-1]
  134. answer = answer.split("</answer>")[0]
  135. return answer.strip()
  136. @staticmethod
  137. def count_xml(text) -> float:
  138. count = 0.0
  139. if text.count("<reasoning>\n") == 1:
  140. count += 0.125
  141. if text.count("\n</reasoning>\n") == 1:
  142. count += 0.125
  143. if text.count("\n<answer>\n") == 1:
  144. count += 0.125
  145. count -= len(text.split("\n</answer>\n")[-1])*0.001
  146. if text.count("\n</answer>") == 1:
  147. count += 0.125
  148. count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001
  149. return count
  150. @staticmethod
  151. def xmlcount_reward_func(completions, **kwargs):
  152. """
  153. Reward function that counts XML tags in the completion.
  154. """
  155. contents = [completion[0]["content"] for completion in completions]
  156. return [ModelTrainer.count_xml(c) for c in contents]
  157. @staticmethod
  158. def soft_format_reward_func(completions, **kwargs):
  159. """
  160. Reward function that checks if the completion has a specific format.
  161. """
  162. pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
  163. responses = [completion[0]["content"] for completion in completions]
  164. matches = [re.match(pattern, r) for r in responses]
  165. return [0.5 if match else 0.0 for match in matches]
  166. @staticmethod
  167. def strict_format_reward_func(completions, **kwargs):
  168. """
  169. Reward function that checks if the completion has a specific format.
  170. """
  171. pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"
  172. responses = [completion[0]["content"] for completion in completions]
  173. matches = [re.match(pattern, r) for r in responses]
  174. return [0.5 if match else 0.0 for match in matches]
  175. @staticmethod
  176. def int_reward_func(completions, **kwargs):
  177. """
  178. Reward function that checks if the completion contains an integer.
  179. """
  180. responses = [completion[0]['content'] for completion in completions]
  181. extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
  182. return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]
  183. @staticmethod
  184. def correctness_reward_func(prompts, completions, answer, **kwargs):
  185. """
  186. Reward function that checks if the completion matches the correct answer.
  187. """
  188. responses = [completion[0]['content'] for completion in completions]
  189. q = prompts[0][-1]['content']
  190. extracted_responses = [ModelTrainer.extract_xml_answer(r) for r in responses]
  191. print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
  192. return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]
  193. if __name__ == "__main__":
  194. # 加载配置文件
  195. config = load_config()
  196. # 设置环境变量
  197. """
  198. # 多机多卡
  199. # export RANK=0 # 第一台机器的 rank
  200. # export WORLD_SIZE=4 # 总共有 4 台机器
  201. # export MASTER_ADDR=<主节点 IP>
  202. # export MASTER_PORT=12345
  203. """
  204. # 单机多卡
  205. os.environ['RANK'] = '0'
  206. os.environ['WORLD_SIZE'] = '1'
  207. os.environ['MASTER_ADDR'] = 'localhost'
  208. os.environ['MASTER_PORT'] = '12345'
  209. # 初始化进程组
  210. dist.init_process_group(backend='nccl', init_method='env://')
  211. # 初始化 ModelTrainer
  212. trainer = ModelTrainer(config)
  213. # 加载模型和分词器
  214. model, tokenizer = trainer.load_model()
  215. # 加载数据集
  216. train_dataset = trainer.load_data(config.train_data_path)
  217. # 训练模型
  218. model = trainer.train(model, tokenizer, train_dataset)
  219. # 保存模型
  220. trainer.save_model(model, tokenizer, config.save_path)
  221. # 确保进程组被销毁
  222. if dist.is_initialized():
  223. dist.destroy_process_group()
  224. print("Training completed.")