微调GPT-2时评估阶段触发CUDA内存不足错误求助

开发者阿泽

2026-6-24

问题：WSL环境下GPT-2微调评估阶段触发CUDA内存溢出

在WSL环境中使用自定义数据集（训练集10000行、验证集1000行）微调GPT-2模型时，训练过程正常，但执行每500步一次的评估阶段时，触发以下CUDA内存不足错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 7.72 GiB. GPU 0 has a total capacity of 24.00 GiB of which 13.38 GiB is free. Process 13650 has 17179869184.00 GiB memory in use. Process 48439 has 17179869184.00 GiB memory in use. Process 21062 has 17179869184.00 GiB memory in use. Process 40988 has 17179869184.00 GiB memory in use. Including non-PyTorch memory, this process has 17179869184.00 GiB memory in use. Of the allocated memory 9.14 GiB is allocated by PyTorch, and 187.59 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

复现代码

import torch.utils
import torch.utils.data
from transformers import (
    GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments, EarlyStoppingCallback, GPT2Config, get_linear_schedule_with_warmup, TrainerCallback)
import torch, np, GPUtil, os
from sklearn.metrics import accuracy_score

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

BATCH_SIZE = 1
EPOCHS = 100
OUTPUT_DIR = "Models/"
SAVE_STEPS = 1000
EVAL_STEPS = 500
ACCUMULATION_STEPS = 2
EARLY_STOPPING_PATIENCE = 3
DROPOUT = 0.25
MODEL_NAME = f"gpt2-bs{BATCH_SIZE}-ep{EPOCHS}-ss{SAVE_STEPS}-es{EVAL_STEPS}-esp{EARLY_STOPPING_PATIENCE}-drpt-{DROPOUT}"
DATASET_PATH = "trainingData.txt"
VAL_DATASET_PATH = "validationData.txt"
BLOCK_SIZE = 128

# Check if a GPU is available and set the device
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Training on device: {}".format(device))

# Load the pre-trained GPT model and tokenizer
model_name = 'gpt2'  # Change this to a larger model
model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
config = GPT2Config.from_pretrained(model_name, dropout=DROPOUT)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
print("Loaded model and tokenizer from {}".format(model_name))
print("Transfering model to device: {}".format(device))
model = GPT2LMHeadModel.from_pretrained(model_name, config=config).to(device)
print("Transfer Complete")

# Load the text messages dataset
dataset_path = DATASET_PATH  # Replace with the path to your file
dataset = TextDataset(tokenizer=tokenizer, file_path=dataset_path, block_size=BLOCK_SIZE)
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
print("Created dataset from {}".format(dataset_path))

# Load the validation dataset
val_dataset_path = "validationData.txt"  # Replace with the path to your file
eval_dataset = TextDataset(tokenizer=tokenizer, file_path=val_dataset_path, block_size=BLOCK_SIZE)
print("Created evaluation dataset from {}".format(val_dataset_path))

# Prepare the data collator for language modeling
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# Define the optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=0.1)
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)

# Calculate the total number of training steps
total_steps = len(train_dataloader) * EPOCHS

# Define the training arguments
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=1,
    save_steps=SAVE_STEPS,
    save_total_limit=2,
    evaluation_strategy="steps",
    load_best_model_at_end=True,
    metric_for_best_model="loss",
    weight_decay=0.01,
    learning_rate=1e-6,  # Tune the learning rate
    # gradient_accumulation_steps=ACCUMULATION_STEPS,
    eval_steps=EVAL_STEPS,
    fp16=True,
    # no_cuda=True,
)

# Create the Trainer instance and start training
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps=0, num_training_steps=total_steps
)

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    loss = torch.nn.functional.cross_entropy(predictions, labels)
    return {"accuracy": accuracy_score(labels, predictions), "loss": loss.item()}

class GPUMonitorCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        gpu = GPUtil.getGPUs()[0]  # Assuming you have 1 GPU
        print(f"GPU Utilization: {gpu.load*100}%, Memory Usage: {gpu.memoryUsed}, Free Memory: {gpu.memoryFree}, GPU Temperature: {gpu.temperature}°C")

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
    eval_dataset=eval_dataset,
    # callbacks=[GPUMonitorCallback()],
    callbacks=[EarlyStoppingCallback(early_stopping_patience=EARLY_STOPPING_PATIENCE), GPUMonitorCallback()],
    optimizers=(optimizer, scheduler),
    compute_metrics=compute_metrics,
)

print("Created data collator and trainer. Starting training...")

trainer.train().half()

print("Training completed. Saving the model...")

# Save the trained model and tokenizer
model.save_pretrained(MODEL_NAME)
tokenizer.save_pretrained(MODEL_NAME)
print("Model saved to {}".format(MODEL_NAME))

解决思路及代码修改建议

1. 修复重复加载模型的显存浪费

代码中两次加载GPT2LMHeadModel，第一次加载的模型未被释放，会占用额外显存。删除第一次无配置的模型加载，只保留带自定义dropout的加载：

# 删除以下无效代码行
# model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
config = GPT2Config.from_pretrained(model_name, dropout=DROPOUT)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
print("Loaded model and tokenizer from {}".format(model_name))
print("Transfering model to device: {}".format(device))
model = GPT2LMHeadModel.from_pretrained(model_name, config=config).to(device)

2. 修正优化器与训练参数的冲突

手动定义的optimizer设置了lr=0.1，但TrainingArguments中已经指定learning_rate=1e-6，两者冲突且0.1的学习率过高。建议删除手动定义的优化器和调度器，让Trainer自动生成适配的优化器：

# 删除以下代码块
# optimizer = torch.optim.AdamW(model.parameters(), lr=0.1)
# total_steps = len(train_dataloader) * EPOCHS
# scheduler = get_linear_schedule_with_warmup(
#     optimizer, num_warmup_steps=0, num_training_steps=total_steps
# )

# 创建Trainer时去掉optimizers参数
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
    eval_dataset=eval_dataset,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=EARLY_STOPPING_PATIENCE), GPUMonitorCallback()],
    compute_metrics=compute_metrics,
)

3. 修复`compute_metrics`函数的逻辑错误

语言模型的eval_pred中，predictions是形状为[batch_size, seq_len, vocab_size]的logits，原函数的维度处理和损失计算逻辑错误，会生成无效大张量占用显存。修改为：

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    # 转换为torch张量处理
    predictions = torch.tensor(predictions)
    labels = torch.tensor(labels)
    # 过滤掉pad token（label=-100）
    mask = labels != -100
    predictions = predictions[mask].view(-1, predictions.size(-1))
    labels = labels[mask].view(-1)
    # 计算交叉熵损失
    loss = torch.nn.functional.cross_entropy(predictions, labels)
    # 计算准确率
    preds = torch.argmax(predictions, dim=-1)
    accuracy = (preds == labels).float().mean().item()
    return {"accuracy": accuracy, "loss": loss.item()}

4. 修正训练启动代码的错误

trainer.train()返回的是TrainerState对象，调用.half()无意义，直接删除：

trainer.train()

5. 优化显存使用策略

开启梯度累积：取消TrainingArguments中gradient_accumulation_steps的注释，用小batch模拟大batch效果，减少显存占用：
```
gradient_accumulation_steps=ACCUMULATION_STEPS,
```

添加评估前显存清理：自定义回调函数在评估前清理显存碎片：

class EvalCleanupCallback(TrainerCallback):
    def on_evaluate(self, args, state, control, **kwargs):
        torch.cuda.empty_cache()

# 添加到callbacks列表中
callbacks=[EarlyStoppingCallback(early_stopping_patience=EARLY_STOPPING_PATIENCE), GPUMonitorCallback(), EvalCleanupCallback()],

尝试关闭混合精度：若WSL环境中fp16评估存在兼容性问题，可设置fp16=False

6. 解决WSL显存统计异常问题

错误中显示的进程显存占用数值异常是WSL2的CUDA统计bug，可尝试：

在启动脚本中添加环境变量：export PYTORCH_NO_CUDA_MEMORY_CACHING=1
关闭其他占用GPU的进程，重启WSL后再运行训练

内容的提问来源于stack exchange，提问作者Morris van den Bergh