Update train.py

train.py

@@ -17,7 +17,7 @@ else:
 
 # 定义教师模型与学生模型
 teacher_model_name = "Qwen/Qwen1.5-7B-Chat"  # 教师模型（较大模型）
-student_model_name = "gpt2"                   # 学生模型（较小模型，可更换）
+student_model_name = "distilgpt2"            # ✅ 学生模型，建议用 distilgpt2 替代 gpt2
 
 # 加载教师模型（仅用于生成软标签，不参与梯度计算）
 teacher = AutoModelForCausalLM.from_pretrained(
@@ -74,7 +74,7 @@ class DistillationTrainer(Trainer):
         super().__init__(*args, **kwargs)
         self.teacher = teacher  # ✅ 传入教师模型
 
-    def compute_loss(self, model, inputs, return_outputs=False):
+    def compute_loss(self, model, inputs, return_outputs=False):  # ❌ 去掉 num_items_in_batch
         labels = inputs["input_ids"]
 
         # ✅ 计算学生模型的输出
@@ -115,6 +115,7 @@ training_args = TrainingArguments(
     logging_steps=100,
     save_strategy="epoch",
     remove_unused_columns=False,  # ✅ 关键设置，确保 Trainer 不删除未识别的列
+    gradient_checkpointing=True,  # ✅ 允许梯度检查点，节省显存
     fp16=True if torch.cuda.is_available() else False
 )