瀏覽代碼

调整train_model_grpo.py 启用vLLM 减少gpu_memory_utilization参数在模型加载或推理时的内存占用比例

zhouyang.xie 3 月之前
父節點
當前提交
ef858b9d68
共有 1 個文件被更改,包括 6 次插入6 次删除
  1. 6 6
      src/train_model_grpo.py

+ 6 - 6
src/train_model_grpo.py

@@ -117,9 +117,9 @@ class ModelTrainer:
             max_seq_length=self.max_seq_length,
             load_in_4bit=self.load_in_4bit, # 值为True 以 4 bit量化进行微调,为False LoRA 16bit。这将内存使用量减少了 4 倍,使我们能够在免费的 16GB 内存 GPU 中实际进行微调。4 位量化本质上将权重转换为一组有限的数字以减少内存使用量。这样做的缺点是准确度会下降 1-2%。如果您想要这种微小的额外准确度,请在较大的 GPU(如 H100)上将其设置为 False。
             dtype=self.dtype,
-            fast_inference = False, # Enable vLLM fast inference
+            fast_inference = True, # Enable vLLM fast inference
             max_lora_rank = lora_rank,
-            gpu_memory_utilization=0.8,# Reduce if out of memory
+            gpu_memory_utilization=0.1,# Reduce if out of memory
         )
         
         # 添加 LoRA 适配器
@@ -149,7 +149,7 @@ class ModelTrainer:
     def train(self, model, tokenizer, train_dataset):
         print("is_bfloat16_supported()=",is_bfloat16_supported())
         training_args = GRPOConfig(
-            use_vllm = False, # use vLLM for fast inference!
+            use_vllm = True, # use vLLM for fast inference!
             learning_rate = 5e-6,
             adam_beta1 = 0.9,
             adam_beta2 = 0.99,
@@ -203,17 +203,17 @@ if __name__ == "__main__":
     # 配置参数
     model_name = os.path.join('..', 'models', 'pretrained', 'DeepSeek-R1-Distill-Qwen-1.5B')
     # model_name: 预训练模型的路径
-    max_seq_length = 1024 # 2048  # 最大序列长度
+    max_seq_length = 2048  # 最大序列长度
     dtype = torch.float16  # 数据类型
     load_in_4bit = True  # 是否以4位精度加载模型
-    lora_rank=32
+    lora_rank=16
 
     # 定义训练集和测试集路径
     train_data_path = os.path.join('..', 'data', 'processed', 'train.jsonl')
     # train_data_path: 训练数据路径
 
     # 初始化 ModelTrainer
-    trainer = ModelTrainer(model_name, max_seq_length, dtype, load_in_4bit)
+    trainer = ModelTrainer(model_name, max_seq_length, dtype, load_in_4bit,lora_rank)
     
     # 加载模型和分词器
     model, tokenizer = trainer.load_model()