【Pytorch使用】CUDA 显存管理与 OOM 排查实战：以 PyTorch 联邦学习训练为例

Rabbit_QL2026-01-07 13:19

一、问题背景

​ 在深度学习训练过程中，特别是使用 GPU 进行联邦学习训练时，经常会遇到 torch.OutOfMemoryError: CUDA out of memory 错误。

1.1 现象发生的可能原因

  • 模型较大（如 VGG16、ResNet 等）
  • 批次大小（batch_size）过大
  • 多个客户端共享模型对象
  • 长时间训练导致内存碎片化
  • 未及时释放中间变量和缓存

1.2 典型错误信息

bash 复制代码 
torch.OutOfMemoryError: CUDA out of memory. 
Tried to allocate 6.12 GiB. 
GPU 2 has a total capacity of 44.42 GiB of which 1.67 GiB is free. 
Including non-PyTorch memory, this process has 42.74 GiB memory in use. 
Of the allocated memory 21.18 GiB is allocated by PyTorch, 
and 21.23 GiB is reserved by PyTorch but unallocated.

​ GPU 显存不足：当前进程已经几乎占满 GPU 显存，PyTorch 尝试再分配一块 6.12 GiB 的显存时失败。

项目 显存
GPU 总显存 44.42 GiB
当前可用 1.67 GiB
进程占用总量 42.74 GiB
PyTorch 实际使用 21.18 GiB
PyTorch 缓存预留 21.23 GiB
本次尝试申请 6.12 GiB ❌
  • 该 OOM 错误表明，尽管 GPU 总显存仍有少量空闲，但由于 PyTorch 显存缓存机制与当前显存碎片化，无法再为新的计算分配一块连续的 6.12 GiB 显存，从而触发 CUDA out of memory 错误。

二、CUDA 内存不足的原因分析

2.1 内存分配问题

  • 问题描述：

    • PyTorch 使用内存池（memory pool）管理 GPU 内存
    • 当内存碎片化严重时，即使有足够的空闲内存，也可能无法分配连续的大块内存

  • 原因：

    • 频繁的小块内存分配和释放
    • 未及时清理缓存
    • 多个模型实例同时占用内存

2.2 批次大小过大

  • 问题描述：

    • batch_size 过大导致单次前向传播需要大量显存
    • VGG16 模型在 224 × 224 224 \times 224 224×224 输入下，每个样本约需 12 − 15 M B 12-15 MB 12−15MB 显存
    • batch_size=512 需要约 6 − 7.5 G B 6-7.5 GB 6−7.5GB 显存（仅前向传播）

  • 计算公式：
    KaTeX parse error: Expected 'EOF', got '_' at position 30: ...} = \text{batch_̲size} \times \t...

2.3 模型对象共享

问题描述：

  • 多个客户端共享同一个模型对象引用
  • 在联邦学习中，所有客户端同时训练会导致内存累积
  • 模型参数的梯度信息会占用额外内存

2.4 计算图保留

问题描述：

  • 默认情况下，PyTorch 会保留计算图用于反向传播
  • 在计算参数差异时，如果不需要梯度，应该使用 torch.no_grad()

三、解决方案概述

​ 我们的解决方案采用多层次的内存优化策略：

  • 环境变量优化 - 减少内存碎片

  • 批次大小调整 - 降低单次内存需求

  • 模型副本管理 - 避免共享对象导致的内存累积

  • 主动内存清理 - 及时释放不需要的内存

  • 计算图优化 - 避免不必要的梯度计算

3.1 环境变量优化

  • 方法：

    在脚本开头设置环境变量，优化 CUDA 内存分配策略：

    python 复制代码 
    import os
# 优化 CUDA 内存分配，避免内存碎片
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

  • 作用：

    • expandable_segments:True 允许内存段动态扩展，减少内存碎片
    • 提高内存利用率，避免"有内存但无法分配"的情况

  • 参考文档： PyTorch Memory Management

3.2 批次大小调整

​ 根据模型大小和 GPU 显存容量选择合适的 batch_size：

python 复制代码 
conf = {
    "batch_size": 32,  
    # VGG16 + 224x224: 推荐 16-64
    # ResNet50 + 224x224: 推荐 32-128
    # 小型模型: 可以更大
}
(1) 选择原则
  • 大型模型（VGG16, ResNet50）: 16-64
  • 中型模型（ResNet18）: 32-128
  • 小型模型（MobileNet）: 64-256
  • 输入尺寸: 224x224 使用较小值，32x32 可以使用较大值
(2)内存估算
python 复制代码 
# 粗略估算公式
estimated_memory = batch_size * input_size * channels * 4 bytes * 3
# 3 倍包括：输入、梯度、优化器状态

3.3 模型副本管理

​ 为每个客户端创建独立的模型实例，避免共享对象：

python 复制代码 
from torchvision import models
import torch.nn as nn

# 为每个客户端创建独立的模型副本
client_model = models.vgg16(weights=None)
# 修改分类层
in_features = client_model.classifier[6].in_features
client_model.classifier[6] = nn.Linear(in_features, num_classes)
# 复制参数（不共享对象）
client_model.load_state_dict(server.global_model.state_dict())
client_model = client_model.to(device)
(1)关键点
  • ✅ 正确做法: 创建新模型实例，只复制参数
  • ❌ 错误做法 : 使用 copy.deepcopy() 或直接共享模型对象
(2)内存对比
  • 共享模型对象：所有客户端共享同一份模型内存 + 梯度内存累积
  • 独立模型副本：每个客户端独立内存，训练完成后可释放

3.4 主动内存清理

​ 在关键位置主动清理 GPU 缓存，但是需要注意频繁调用 empty_cache() 会引入同步开销，不建议在每个 batch 调用。

(1) 训练循环中的清理
python 复制代码 
for batch_id, (imgs, labels) in enumerate(self.train_loader):
    data, target = imgs.to(self.device), labels.to(self.device)
    optimizer.zero_grad()
    output = self.local_model(data)
    loss = loss_func(output, target)
    loss.backward()
    optimizer.step()
    
    # 清理中间变量释放内存
    del data, target, output, loss
    
    # 定期清理缓存
    if (batch_id + 1) % 10 == 0:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
  • 清理时机：
    • 每个批次后删除中间变量
    • 每 N 个批次清理一次缓存（N=10 或根据内存情况调整）
    • 客户端训练完成后清理
(2) 参数差异计算优化
python 复制代码 
# 使用 no_grad 避免保存计算图
with torch.no_grad():
    diff = dict()
    for name, data in self.local_model.state_dict().items():
        global_data = global_model.state_dict()[name]
        # 移到 CPU 避免 GPU 内存累积
        diff[name] = (data - global_data).cpu()

# 清理 GPU 缓存
if torch.cuda.is_available():
    torch.cuda.empty_cache()
  • 优化点：
    • torch.no_grad(): 不保存计算图，节省内存
    • .cpu(): 将结果移到 CPU，释放 GPU 内存
    • 计算完成后立即清理缓存
(2) 模型聚合优化
python 复制代码 
def model_aggregate(self, weight_accumulator):
    with torch.no_grad():
        for name, data in self.global_model.state_dict().items():
            # 确保张量在正确的设备上
            if isinstance(weight_accumulator[name], torch.Tensor):
                update_per_layer = weight_accumulator[name].to(data.device) * self.conf["lambda"]
            else:
                update_per_layer = weight_accumulator[name] * self.conf["lambda"]
            data.add_(update_per_layer)
    
    # 清理 GPU 缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

3.5 内存监控

​ 添加内存监控代码，便于调试和优化：

python 复制代码 
if torch.cuda.is_available():
    device_id = 2
    torch.cuda.set_device(device_id)
    device = torch.device(f'cuda:{device_id}')
    
    # 清理 GPU 缓存
    torch.cuda.empty_cache()
    
    # 显示内存使用情况
    allocated = torch.cuda.memory_allocated(device_id) / 1024**3
    total = torch.cuda.get_device_properties(device_id).total_memory / 1024**3
    print(f"GPU 内存使用情况: {allocated:.2f} GB / {total:.2f} GB")

四、最佳实践

4.1 批次大小选择

python 复制代码 
# 根据 GPU 显存选择 batch_size
gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3

if gpu_memory_gb >= 24:  # RTX 3090, A100 等
    batch_size = 64
elif gpu_memory_gb >= 12:  # RTX 3060, V100 等
    batch_size = 32
else:  # GTX 1080 等
    batch_size = 16

4.2 梯度累积（Gradient Accumulation）

​ 当显存不足但需要大批次训练时，使用梯度累积：

python 复制代码 
accumulation_steps = 4  # 累积 4 个批次
effective_batch_size = batch_size * accumulation_steps

for batch_id, (data, target) in enumerate(train_loader):
    output = model(data)
    loss = criterion(output, target) / accumulation_steps
    loss.backward()
    
    if (batch_id + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
        torch.cuda.empty_cache()

4.3 混合精度训练（Mixed Precision）

使用 FP16 可以显著减少显存占用：

python 复制代码 
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in train_loader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

4.4 检查点机制（Checkpointing）

对于非常大的模型，可以使用梯度检查点：

python 复制代码 
from torch.utils.checkpoint import checkpoint

# 在前向传播中使用检查点
output = checkpoint(model, input)

4.5 内存清理检查清单

在以下位置添加内存清理：

✅ 每个训练批次后删除中间变量

✅ 每个客户端训练完成后清理缓存

✅ 模型聚合后清理缓存

✅ 评估前清理缓存

✅ 使用 torch.no_grad() 进行不需要梯度的计算

五、故障排查

问题 1: 仍然出现内存不足

可能原因：

  • batch_size 仍然太大
  • 模型太大
  • 多个进程共享 GPU

解决方案：

python 复制代码 
# 1. 进一步减小 batch_size
conf["batch_size"] = 16  # 或更小

# 2. 使用梯度累积
# 3. 使用混合精度训练
# 4. 检查是否有其他进程占用 GPU
nvidia-smi  # 查看 GPU 使用情况

问题 2: 内存清理后仍然占用大量内存

可能原因：

  • PyTorch 的内存池机制会保留已分配的内存
  • 这是正常现象，不影响使用

解决方案：

python 复制代码 
# 如果需要强制释放，可以重启 Python 进程
# 或者在代码中定期重置
torch.cuda.empty_cache()
torch.cuda.ipc_collect()  # 收集 IPC 资源

问题 3: 训练速度变慢

可能原因：

  • 频繁的内存清理会影响性能
  • batch_size 太小

解决方案：

python 复制代码 
# 1. 减少清理频率
if (batch_id + 1) % 50 == 0:  # 改为每 50 个批次清理一次
    torch.cuda.empty_cache()

# 2. 在非关键路径清理
# 3. 使用更大的 batch_size（如果内存允许）

问题 4: 多 GPU 训练时的内存问题

解决方案：

python 复制代码 
# 使用 DataParallel 或 DistributedDataParallel
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
    # 注意：batch_size 会按 GPU 数量自动分配

六、代码示例总结

6.1 完整的客户端训练代码示例

python 复制代码 
def local_train(self, global_model):
    # 1. 复制模型参数（使用 no_grad）
    with torch.no_grad():
        self.local_model.load_state_dict(global_model.state_dict())
    
    # 2. 初始化优化器
    optimizer = torch.optim.SGD(self.local_model.parameters(), lr=0.001)
    loss_func = torch.nn.CrossEntropyLoss()
    
    # 3. 训练循环
    self.local_model.train()
    for epoch in range(self.conf["local_epochs"]):
        for batch_id, (imgs, labels) in enumerate(self.train_loader):
            # 数据移动到设备
            data, target = imgs.to(self.device), labels.to(self.device)
            
            # 前向传播和反向传播
            optimizer.zero_grad()
            output = self.local_model(data)
            loss = loss_func(output, target)
            loss.backward()
            optimizer.step()
            
            # 清理中间变量
            del data, target, output, loss
            
            # 定期清理缓存
            if (batch_id + 1) % 10 == 0:
                torch.cuda.empty_cache()
    
    # 4. 计算参数差异（移到 CPU）
    with torch.no_grad():
        diff = dict()
        for name, data in self.local_model.state_dict().items():
            diff[name] = (data - global_model.state_dict()[name]).cpu()
    
    # 5. 最终清理
    torch.cuda.empty_cache()
    return diff

参考资料

  1. PyTorch CUDA Memory Management
  2. PyTorch Performance Tuning Guide
  3. Understanding CUDA Memory Allocation
相关推荐
deephub1 天前
Agentic RAG：用LangGraph打造会自动修正检索错误的 RAG 系统
人工智能·大语言模型·rag·langgraph
2501_941878741 天前
在奥克兰云原生实践中构建动态配置中心以支撑系统稳定演进的工程经验总结
开发语言·python
weixin_443297881 天前
Python打卡训练营第31天
开发语言·python
坠金1 天前
方差、偏差
人工智能·机器学习
宏基骑士1 天前
Python之类中函数间的参数传递（有继承和无继承）
python
强盛小灵通专卖员1 天前
airsim无人机仿真深度强化学习自动避障辅导
人工智能·无人机·sci·深度强化学习·airsim·自动避障·小论文
540_5401 天前
ADVANCE Day41
人工智能·python·深度学习
yunni81 天前
智能办公软件的强大AI功能
人工智能·ai工具
云雾J视界1 天前
AI+IoT双轮驱动：构建风电设备预测性维护数字孪生体的全栈技术实践
人工智能·物联网·ai·lstm·iot·scada·金风科技
热门推荐
01GitHub 镜像站点02安娜的档案(Anna’s Archive) 镜像网站/国内最新可访问入口（持续更新）03Labelme从安装到标注：零基础完整指南04Linux下V2Ray安装配置指南05手把手教你通过Gemini3 pro 学生认证，白用一年，手慢无！06GitLab 零基础入门指南：从安装到项目管理全流程07jdk21下载、安装（Windows、Linux、macOS）08【踩坑笔记】50系显卡适配的 PyTorch 安装09UV安装并设置国内源102025-04-03 Latex学习1——本地配置Latex + VScode环境