从单卡到多卡：用CUDA_VISIBLE_DEVICES和PyTorch DataParallel轻松管理你的炼丹炉
当你第一次尝试在PyTorch中使用多GPU训练时，可能会遇到各种令人困惑的问题：为什么代码明明指定了使用GPU 1，却仍然跑在GPU 0上？为什么DataParallel会报错说找不到某些设备？这些问题大多与CUDA的设备可见性管理有关。本文将带你从单卡调试开始，逐步掌握多GPU训练的核心技巧。
1. 单卡调试：精确控制你的实验环境
在深度学习开发中，我们常常需要在单卡上调试代码，即使服务器有多块GPU。这时候CUDA_VISIBLE_DEVICES环境变量就派上用场了。
# 在Linux/Mac上设置只使用第一块GPU export CUDA_VISIBLE_DEVICES=0 # 在Windows上设置 set CUDA_VISIBLE_DEVICES=0设置后，你的PyTorch程序将只能"看到"这一块GPU，无论代码中如何指定设备。这在多人共用服务器时特别有用，可以避免资源冲突。
验证当前可见设备的最简单方法：
import torch print(f"可用GPU数量: {torch.cuda.device_count()}") print(f"当前设备: {torch.cuda.current_device()}")注意：CUDA_VISIBLE_DEVICES的设置在Python进程启动后就无法更改。如果需要在运行时切换设备，需要在启动新进程前修改环境变量。
提示：在Jupyter notebook中，环境变量需要在启动notebook前设置，或者在代码中使用os.environ修改
2. 多卡数据并行：扩展你的训练能力
当你准备好将训练扩展到多GPU时，PyTorch的DataParallel是最简单的选择。它自动将数据分割到不同GPU上并行计算。
2.1 基础多卡设置
首先确保目标GPU可见：
# 使用GPU 0和1 export CUDA_VISIBLE_DEVICES=0,1然后修改你的训练代码：
model = MyModel() if torch.cuda.device_count() > 1: print(f"使用 {torch.cuda.device_count()} 块GPU") model = nn.DataParallel(model) model = model.cuda() # 这会自动使用所有可见GPU2.2 理解设备映射
DataParallel的device_ids参数决定了使用哪些GPU。关键点：
device_ids中的编号是相对于CUDA_VISIBLE_DEVICES的- 如果不指定
device_ids，默认使用所有可见GPU - 主GPU（接收梯度的设备）默认是
device_ids[0]
# 明确指定使用前两块可见GPU model = nn.DataParallel(model, device_ids=[0,1]) # 等价于 model = nn.DataParallel(model)2.3 常见问题排查
当遇到多GPU训练问题时，检查以下方面：
- 设备可见性：确认
torch.cuda.device_count()返回预期值 - 张量位置：确保所有输入张量都在GPU上
- batch size：总batch size是每GPU batch size乘以GPU数量
# 检查各GPU内存使用情况 print(torch.cuda.memory_allocated(0)) # GPU 0已分配内存 print(torch.cuda.memory_reserved(1)) # GPU 1保留内存3. 高级场景：灵活分配非连续GPU
在实际生产环境中，你可能需要更灵活的GPU分配策略。
3.1 使用非连续GPU
服务器可能有多个用户，需要共享不同GPU：
# 只使用GPU 0和2 export CUDA_VISIBLE_DEVICES=0,2在代码中，这些GPU将被重新编号为0和1：
# 使用重新编号后的设备 model = nn.DataParallel(model, device_ids=[0,1])3.2 多任务协同工作
当同时运行多个训练任务时，可以为每个任务分配不同的GPU：
# 终端1：运行任务A使用GPU 0,1 export CUDA_VISIBLE_DEVICES=0,1 python train_a.py # 终端2：运行任务B使用GPU 2,3 export CUDA_VISIBLE_DEVICES=2,3 python train_b.py3.3 动态GPU分配
在某些情况下，你可能需要根据GPU负载动态分配设备：
import os import torch def select_least_used_gpu(num_gpus=1): """选择使用率最低的GPU""" device_usage = [] for i in range(torch.cuda.device_count()): mem = torch.cuda.memory_allocated(i) device_usage.append((i, mem)) # 按内存使用量排序 device_usage.sort(key=lambda x: x[1]) selected = [str(x[0]) for x in device_usage[:num_gpus]] return ",".join(selected) # 动态设置可见设备 os.environ['CUDA_VISIBLE_DEVICES'] = select_least_used_gpu(2)4. 性能优化与最佳实践
多GPU训练不仅仅是让代码跑起来，还需要考虑效率问题。
4.1 数据加载优化
使用DataParallel时，数据加载器需要适当配置：
from torch.utils.data import DataLoader # 重要：设置num_workers > 0以避免数据加载成为瓶颈 loader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)注意：pin_memory=True可以加速CPU到GPU的数据传输，但会稍微增加CPU内存使用。
4.2 梯度累积技巧
当GPU数量增加但显存不足时，可以结合梯度累积：
optimizer.zero_grad() for i, (inputs, targets) in enumerate(loader): outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() if (i+1) % 2 == 0: # 每2个batch更新一次 optimizer.step() optimizer.zero_grad()4.3 混合精度训练
现代GPU支持混合精度训练，可以显著减少显存使用并加速训练：
from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for inputs, targets in loader: optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()在实际项目中，我发现合理组合这些技巧可以在保持batch size不变的情况下，将训练速度提升2-3倍。特别是在使用大型Transformer模型时，混合精度训练几乎成为标配。
