重写Hugging Face TrainerCallback的on_step_end引发CUDA内存溢出求助

核心原因 #

1. 训练阶段模型处于训练模式（model.train()），BatchNorm、Dropout等层会保留梯度和中间激活张量，评估训练集时未切换到eval模式，导致训练+评估的内存占用叠加。
2. CTCTrainer的evaluate默认复用训练的batch size配置，但训练集规模远大于验证集，加上训练过程中未释放的计算图、梯度缓存，内存负载骤增。
3. 回调中直接调用trainer.evaluate会绑定当前训练的模型状态，训练累积的优化器状态、梯度张量未临时释放，进一步挤占GPU内存。

可行解决办法 #

1. 评估时强制切换模型模式并禁用梯度计算 #

修改自定义回调的on_step_end方法，手动控制模型状态，避免训练模式下的额外内存开销：

def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
    if control.should_evaluate:
        control_copy = deepcopy(control)
        # 切换到评估模式，关闭梯度计算
        self._trainer.model.eval()
        with torch.no_grad():
            self._trainer.evaluate(eval_dataset=self._trainer.train_dataset, metric_key_prefix='train')
        # 切回训练模式
        self._trainer.model.train()
        torch.cuda.empty_cache()
        return control_copy

2. 单独降低训练集评估的batch size #

在调用evaluate时指定更小的评估batch size，避免和训练batch size一致导致内存过载：

self._trainer.evaluate(
    eval_dataset=self._trainer.train_dataset,
    metric_key_prefix='train',
    per_device_eval_batch_size=2  # 根据GPU内存调整，建议比训练batch size小一半以上
)

3. 减少训练集评估的频率 #

当前eval_steps=5过于频繁，训练集数据量大，每5步评估一次会持续占用内存。建议改为按epoch评估，或大幅调大评估步数：

training_args = TrainingArguments(
    # ... 其他参数
    evaluation_strategy="epoch",  # 改为按epoch评估训练集
    # 或保留step模式但调大步数：eval_steps=100
)

4. 临时清空训练阶段的梯度缓存 #

在评估前手动清空优化器的梯度，释放部分内存：

def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
    if control.should_evaluate:
        control_copy = deepcopy(control)
        # 临时清空梯度，释放内存
        self._trainer.optimizer.zero_grad(set_to_none=True)
        self._trainer.model.eval()
        with torch.no_grad():
            self._trainer.evaluate(eval_dataset=self._trainer.train_dataset, metric_key_prefix='train')
        self._trainer.model.train()
        torch.cuda.empty_cache()
        return control_copy

5. 调大eval_accumulation_steps分散内存压力 #

通过分批计算评估指标，减少单次评估的内存占用：

training_args = TrainingArguments(
    # ... 其他参数
    eval_accumulation_steps=15  # 数值越大，每次评估占用内存越小，按需调整
)

验证建议 #

• 用torch.cuda.memory_allocated()和torch.cuda.memory_reserved()在评估前后打印内存使用，定位内存泄漏点。
• 优先尝试切换模型模式+降低评估batch size的组合方案，这是解决此类OOM最直接的手段。

内容的提问来源于stack exchange，提问作者Daniel Farrugia