如何避免微调Huggingface Pegasus模型时训练中途VRAM显存不足

之前设置固定截断长度仍出现显存波动，本质是遗漏了三个高频踩坑点，且因为数据集规模大，触发崩溃的极端样本占比极低，往往要连续运行数天才会碰到：

• 仅对输入文本做了截断，未对生成目标（标签）做截断设置：seq2seq模型的decoder注意力矩阵显存占用与标签序列长度呈平方关系，只要某个样本的摘要长度远高于阈值，整个批次的decoder计算图长度就会被显著拉长，显存会瞬间翻倍。
• Huggingface的DataCollatorForSeq2Seq默认按批次内最长样本做动态padding，即便给tokenizer传入了固定max_length，只要样本本身没被提前截断，collator就会把批次序列长度撑到样本的真实长度，不会主动做上限限制。
• 长期连续训练产生的显存碎片：PyTorch连续运行数天后，频繁申请、释放显存会产生大量不连续的内存块，哪怕单批次所需显存和首轮运行时完全一致，也会因为找不到连续的足够大的显存块触发OOM，表现为显存无理由突然飙升。

• 从源头消除极端样本触发的崩溃风险
  • 编写tokenize逻辑时，同时对输入文本、目标标签做明确截断，不能只处理输入侧。参考标准写法：

    def batch_tokenize(examples):
        model_inputs = tokenizer(
            examples["source_text"], 
            max_length=1024, 
            truncation=True
        )
        # 必须单独对目标文本设置截断，是多数人容易遗漏的关键步骤
        labels = tokenizer(
            text_target=examples["target_summary"], 
            max_length=128, 
            truncation=True
        )
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs
    

  • 正式启动全量训练前，随机抽取1%的全量数据集做长度遍历，统计所有样本的真实输入、标签长度，把长度超过设定max_length的异常样本直接从数据集中过滤，不要等训练到对应批次才临时处理触发崩溃。
• 从框架层硬约束显存使用上限
  • 训练启动前，给当前PyTorch进程设置最高可用显存占比：添加代码torch.cuda.set_per_process_memory_fraction(0.85)，为系统、临时计算开销预留15%的显存余量，避免把总显存占满触发硬崩溃。
  • 设置环境变量开启可扩展显存分配模式：添加代码os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"，大幅降低长期训练的显存碎片问题，避免连续运行数天后因为碎片化显存积累触发OOM。
• 进一步压缩显存波动幅度
  • 开启Trainer的group_by_length=True参数：自动将长度接近的样本分配到同一个批次，避免“绝大多数批次为短样本，突然混入一个超长样本拉高整个批次显存”的情况，让每个批次的序列长度维持在稳定区间。
  • 开启Pegasus原生支持的梯度检查点：模型初始化后添加model.gradient_checkpointing_enable()，可将激活值的显存占用降低40%-60%，为整个训练流程留足安全余量，不会因为小幅显存波动直接崩溃。

完成以上配置后，全训练流程的单批次显存占用会和首轮运行的观测值基本一致，不会出现无规律飙升的情况，无需运行数天之后才暴露显存不足的问题。

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