训练时CUDA内存不足求助：双80GB GPU仍出现OOM错误

核心问题分析 #

从报错信息和代码来看，问题集中在单GPU独占负载、数据处理逻辑错误、训练配置未匹配模型类型这几个关键点：

1. 未启用多GPU训练 #

你的服务器配备2块80GB GPU，但代码默认仅将模型加载到GPU0，完全未利用第二块GPU的显存，导致单GPU承载全部训练负载，直接占满显存。

2. 标签设置不符合因果语言模型要求 #

代码中把labels设为[1] * len(example["input_ids"])，这是错误的：

• 因果语言模型的labels应该是input_ids的副本，仅对pad位置设置为-100（避免计算pad部分的损失）
• 全1的标签会让模型计算所有位置的损失，包括无效的pad段，既浪费显存，还会导致训练逻辑混乱

3. 数据收集器选型错误 #

你使用了DataCollatorForSeq2Seq，但DictaLM是因果语言模型（Causal LM），Seq2Seq收集器是为编码器-解码器结构设计的，会生成不必要的张量，额外占用显存。

4. 未开启半精度训练 #

默认全精度（FP32）训练会占用更多显存，开启半精度（FP16）可大幅降低内存占用，且对训练效果影响极小。

5. 未处理内存碎片 #

报错明确提示内存碎片问题，但你未设置对应的环境变量缓解。

具体修复步骤 #

1. 启用多GPU训练 #

在Trainer初始化时添加device_map="auto"，让Transformers自动分配模型到多GPU：

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    device_map="auto"  # 新增该行启用多GPU分配
)

2. 修正标签生成逻辑 #

将labels设为input_ids的副本，并把pad位置标记为-100：

def tokenize_function(examples):
    outputs = tokenizer(examples["Prompt"], max_length=512, truncation=True)
    # 用input_ids作为labels基础
    outputs["labels"] = outputs["input_ids"].copy()
    # 把pad位置设为-100，避免计算损失
    pad_token_id = tokenizer.pad_token_id
    for i in range(len(outputs["labels"])):
        outputs["labels"][i] = [x if x != pad_token_id else -100 for x in outputs["labels"][i]]
    return outputs

# 替换原tokenize_function调用
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=['Prompt']
)
# 删除原错误的labels生成代码

3. 更换适配因果LM的数据收集器 #

使用DataCollatorForLanguageModeling替代Seq2Seq收集器：

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False  # 因果LM用自回归训练，关闭掩码语言模型模式
)

4. 开启半精度训练 #

在TrainingArguments中添加fp16=True：

args = TrainingArguments(
    'results',
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    weight_decay=0.01,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    fp16=True  # 新增该行开启半精度训练
)

5. 缓解内存碎片 #

在启动训练脚本前设置环境变量：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

或者在Python代码开头添加：

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

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