多GPU训练时CUDA显存不足，为何仅识别单GPU显存？

问题描述 #

我正在基于RoBERTa相关教程从头训练模型，机器上有4块GTX 1080 Ti GPU，其中GPU 0、1被其他研究员占用，因此执行命令：

CUDA_VISIBLE_DEVICES=2,3 python script.py

但训练时出现显存不足错误，错误信息显示尝试在GPU 0（可见设备中的第一块，对应物理机GPU2）上分配显存，且只识别了单块GPU的10.91GiB容量，无法利用两块GPU的总显存来增大batch size。

nvidia-smi输出：

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 530.30.02              Driver Version: 530.30.02    CUDA Version: 12.1     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                  Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf            Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce GTX 1080 Ti      Off| 00000000:02:00.0 Off |                  N/A |
| 20%   54C    P2               83W / 250W|   4692MiB / 11264MiB |     45%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce GTX 1080 Ti      Off| 00000000:03:00.0 Off |                  N/A |
| 26%   60C    P2               73W / 250W|   4650MiB / 11264MiB |     44%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   2  NVIDIA GeForce GTX 1080 Ti      Off| 00000000:81:00.0 Off |                  N/A |
| 50%   71C    P0               84W / 250W|      0MiB / 11264MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   3  NVIDIA GeForce GTX 1080 Ti      Off| 00000000:82:00.0 Off |                  N/A |
| 30%   53C    P0               75W / 250W|      0MiB / 11264MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A   3494144      C   python                                     4690MiB |
|    1   N/A  N/A   3494896      C   python                                     4648MiB |
+---------------------------------------------------------------------------------------+

错误信息：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.07 GiB (GPU 0; 10.91 GiB total capacity; 8.36 GiB already allocated; 1.93 GiB free; 8.40 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

解答 #

1. 为什么只识别单GPU显存？ #

CUDA_VISIBLE_DEVICES=2,3只是把物理机的GPU2、3映射为当前进程可见的GPU0、1，但训练脚本默认只会把模型和数据加载到可见设备中的第一块GPU（即映射后的GPU0，对应物理机GPU2），没有启用多GPU并行训练机制，因此只会使用这一块GPU的显存，不会自动合并多块GPU的显存空间。

2. 如何利用多GPU增大batch size？ #

要利用多块GPU的显存，需要在代码中启用PyTorch的多GPU并行机制，以下是两种常用方案：

方案一：使用DataParallel（快速实现，适合单机多卡）

在模型定义完成后，添加以下代码将模型包装为多GPU并行：

import torch

# 假设model是你的RoBERTa模型实例
model = torch.nn.DataParallel(model, device_ids=[0, 1])
model = model.cuda()

• device_ids=[0,1]对应你可见的两块GPU（映射后的GPU0、1，即物理机的GPU2、3）
• 总batch size = 设置的batch_size参数 × GPU数量，因此可以直接设置单GPU能承载的最大batch size，总batch size自动翻倍。

方案二：使用DistributedDataParallel（推荐，效率更高）

这是PyTorch官方推荐的多GPU训练方案，效率比DataParallel更高，步骤如下：

1. 修改启动命令：

CUDA_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 script.py

• --nproc_per_node=2表示使用2块GPU

2. 在脚本开头添加分布式初始化代码：

import torch
import torch.distributed as dist
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

# 初始化分布式进程组
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
torch.cuda.set_device(rank)

# 数据加载时使用DistributedSampler，确保每个GPU加载不同的数据分片
dataset = 你的数据集实例
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler, batch_size=单GPU的batch size)

# 初始化模型并包装为DistributedDataParallel
model = 你的RoBERTa模型实例
model = model.to(rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])

• 每个GPU处理batch_size数量的数据，总batch size = 单GPU batch size × GPU数量，同样可以直接按单GPU最大容量设置batch size。

内容的提问来源于stack exchange，提问作者Higo Felipe Silva Pires