Amazon SageMaker上Longformer模型CUDA内存不足及多GPU未利用问题

核心问题在于ml.p3.8xlarge实例的4块V100 GPU仅GPU0被占用，导致单GPU内存过载触发CUDA OOM；同时CPU实例虽可运行但训练效率极低。已尝试的并行训练、混合精度等优化未解决多GPU利用率问题，是当前OOM的主要诱因。

1. 修正PyTorch多GPU训练配置 #

你提到尝试了并行训练，但大概率是未正确适配SageMaker环境的多GPU逻辑，以下是两种可靠实现方式：

方式1：使用DataParallel（快速适配） #

import torch
from transformers import LongformerForSequenceClassification

# 初始化模型
model = LongformerForSequenceClassification.from_pretrained(
    "allenai/longformer-base-4096",
    num_labels=2,
    low_cpu_mem_usage=True
)

# 自动适配多GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model.to(device)

注意：DataParallel会将模型副本放到所有GPU，数据拆分后分配到各GPU计算，最后汇总梯度。

方式2：使用DistributedDataParallel（高效推荐） #

DDP是PyTorch官方推荐的多GPU训练方案，内存效率更高，适配SageMaker环境需额外配置：

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from transformers import LongformerForSequenceClassification, DataCollatorWithPadding

# 初始化分布式环境
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29500"
dist.init_process_group(backend="nccl", rank=0, world_size=4)

# 加载模型并移动到指定GPU
local_rank = 0
model = LongformerForSequenceClassification.from_pretrained(
    "allenai/longformer-base-4096",
    num_labels=2,
    low_cpu_mem_usage=True
).to(local_rank)
model = DDP(model, device_ids=[local_rank])

# 配置数据加载器（必须用DistributedSampler拆分数据）
train_dataset = ... # 你的训练数据集
sampler = DistributedSampler(train_dataset)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=8, # 单GPU批次大小，总批次=8*4=32
    sampler=sampler,
    collate_fn=data_collator
)

2. SageMaker Notebook环境适配 #

• 确认使用的Kernel为PyTorch 1.10+版本，旧版本对分布式训练支持不完善；
• 在Notebook开头执行以下命令设置环境变量，确保多GPU通信正常：

%%bash
export NCCL_DEBUG=INFO
export OMP_NUM_THREADS=1

3. 进一步降低单GPU内存负载 #

结合多GPU配置，叠加以下优化可彻底解决OOM：

• 梯度累积：若单GPU批次仍过大，通过累积梯度实现等效大批次，降低单步内存占用：

accumulation_steps = 4 # 每4步更新一次梯度，等效批次=8*4*4=128
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

for step, (inputs, labels) in enumerate(train_loader):
    inputs = {k: v.to(local_rank) for k, v in inputs.items()}
    labels = labels.to(local_rank)
    
    with torch.cuda.amp.autocast(): # 配合混合精度
        outputs = model(**inputs)
        loss = outputs.loss / accumulation_steps
    
    loss.backward()
    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
    
    # 及时清理无用张量
    del outputs, loss
    torch.cuda.empty_cache()

• 调整Longformer参数：减小attention_window（如从512改为256），降低注意力机制的内存开销；
• 启用模型量化：使用torch.ao.quantization工具对模型进行INT8量化，可减少约75%的内存占用：

model.eval()
model.qconfig = torch.ao.quantization.get_default_qconfig('fbgemm')
torch.ao.quantization.prepare(model, inplace=True)
# 用少量校准数据跑一次前向
for batch in train_loader:
    model(**{k: v.to(local_rank) for k, v in batch.items()})
    break
torch.ao.quantization.convert(model, inplace=True)

4. 验证多GPU利用率 #

执行以下代码确认模型是否正确分配到所有GPU：

print(f"可用GPU数量: {torch.cuda.device_count()}")
print(f"模型参数所在设备: {next(model.parameters()).device}")

若输出显示多个GPU，再运行!nvidia-smi查看内存占用，正常情况下4块GPU的内存使用率应接近。

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