使用PyTorch Lightning多GPU微调扩散模型时CUDA内存不足

问题背景 #

在4张RTX3090（24GB）显卡上基于PyTorch Lightning 2.1微调扩散模型，遇到以下内存异常：

• 单GPU配置（devices=1，strategy="ddp"），batch_size=4时模型占用约22542MB内存，运行正常
• 切换为4GPU（devices=4）且未修改其他配置时，首轮训练直接触发“CUDA out of memory”错误
• 单步调试显示training_step函数仅占用约17000MB内存，但退出该函数后才出现内存不足，怀疑反向传播阶段存在内存泄漏或超额占用
• 单GPU+DDP策略下也会出现相同的内存不足问题

复现代码 #

from share import *
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from torch.utils.data import DataLoader
from tutorial_dataset import MyDataset,BoothDataset,CatDataset,DreamDataset,read_pic,RemoveDataset,LaionDataset
#from cldm.logger import ImageLogger
from cldm.model import create_model, load_state_dict
from torchvision.utils import save_image
import copy
from PIL import Image
########seed#########
import random
import os
import numpy as np
import torch

def main():
    torch.set_float32_matmul_precision('medium')
    pl.seed_everything(42,workers=True)
    # Configs
    resume_path = './model/pure.ckpt' 
    batch_size = 4
    learning_rate = 1e-4  

    # First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.
    model = create_model('./model/cldm_v15.yaml').cpu()
    model.load_state_dict(load_state_dict(resume_path, location='cpu'))
    model.learning_rate = learning_rate

    # Misc
    dataset = LaionDataset()

    dataloader = DataLoader(dataset, num_workers=4, batch_size=batch_size, shuffle=True)

    checkpoint_callback = ModelCheckpoint(
        monitor='train_loss',
        dirpath='./model_save/',
        filename='{epoch:04d}-{train_loss:.4f}-{global_step:.0f}' ,
        every_n_epochs = 1,
    )
    trainer = pl.Trainer(precision=16,
                            accelerator="auto", devices=4, strategy="ddp",
                            callbacks=[checkpoint_callback] ,
                            max_epochs = 1,
                            deterministic=True,  #seed
                            )
    # Train!
    trainer.fit(model, dataloader)


if __name__=='__main__':
    main()

问题分析与解决方案 #

核心原因 #

1. DDP批次大小误解：DDP模式下batch_size是单GPU的批次量，而非全局总批次。设置batch_size=4后，4张GPU各自处理4个样本，全局总批次为16，反向传播时单GPU需存储对应梯度，内存占用远高于单GPU训练。
2. 单GPU+DDP的额外开销：即使单GPU运行DDP，也会引入进程组初始化、梯度同步缓冲区等额外内存开销，导致比普通单GPU训练占用更多内存。
3. 反向传播的内存峰值：training_step仅执行前向传播，反向传播需要存储前向过程的所有中间张量用于梯度计算，这部分内存不会在training_step内体现，退出函数后才会达到内存峰值。

具体解决措施 #

• 降低单GPU批次大小：将batch_size从4调整为2或1，匹配多GPU下的内存承载能力。
• 优化DDP策略：
  • 使用strategy="ddp_find_unused_parameters_false"（确认模型无未使用参数时），减少不必要的梯度内存占用。
  • 添加gradient_clip_val=1.0到Trainer配置，限制梯度范数，避免梯度爆炸导致的内存激增。
• 调整确定性训练设置：deterministic=True会强制使用确定性算法，增加内存开销，改为deterministic=False或deterministic="warn_only"可释放部分内存。
• 优化数据加载：根据CPU核心数调整num_workers（比如4卡设置为8），避免数据加载占用过多CPU内存间接影响GPU。
• 精简模型参数：检查模型是否有可冻结的冗余参数，减少需要计算梯度的参数数量，降低反向传播的内存消耗。

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