1. 多GPU训练的基本概念与准备工作第一次接触多GPU训练时我完全被各种设备编号和并行策略搞晕了。直到踩过几次坑之后才明白其实核心就是两件事控制哪些GPU对程序可见以及如何把模型分布到这些GPU上。我们先从最基础的准备工作说起。在开始之前你需要确认几件事情服务器上确实安装了多块NVIDIA GPU可以用nvidia-smi命令查看正确安装了CUDA和cuDNNPyTorch版本支持多GPU训练一般1.0以上版本都没问题检查GPU状态的命令如下nvidia-smi这个命令会显示所有GPU的状态包括内存使用率、GPU利用率等。输出结果中最左边一列就是GPU的物理编号通常是0,1,2,3这样的数字。2. 控制GPU可见性的两种方式2.1 使用CUDA_VISIBLE_DEVICES环境变量这是最常用的方法可以在程序启动前就确定哪些GPU对程序可见。比如CUDA_VISIBLE_DEVICES0,1 python train.py这条命令会让程序只能看到0号和1号GPU。注意几个关键点编号之间用英文逗号隔开不能有空格这个设置必须在启动Python程序之前完成在程序内部这些GPU会被重新编号为从0开始我遇到过的一个典型错误是# 错误的写法有空格 CUDA_VISIBLE_DEVICES0, 1 python train.py这种写法会导致只有0号GPU被识别逗号后面的空格会让参数解析出错。2.2 在代码中使用os.environ设置如果你不想在命令行中指定也可以在Python代码的最开始设置import os os.environ[CUDA_VISIBLE_DEVICES] 0,1但要注意这行代码必须放在任何与GPU相关的操作之前包括import torch之前因为PyTorch导入时会初始化CUDA环境我曾经犯过一个错误把设置代码放在了模型定义之后结果完全没生效调试了半天才发现问题。3. DataParallel的基本使用与原理3.1 最简单的使用方式当你已经通过CUDA_VISIBLE_DEVICES设置了可见GPU后最简单的多GPU使用方法就是model torch.nn.DataParallel(model).cuda()这样PyTorch会自动使用所有可见的GPU进行训练。DataParallel的工作原理是在前向传播时将输入数据切分到不同GPU每个GPU计算自己的那部分在主GPU默认是可见列表中的第一个上汇总结果计算损失并反向传播3.2 自定义设备编号有时候我们想更精细地控制使用哪些GPU可以这样model torch.nn.DataParallel(model, device_ids[0,1]).cuda()这里的device_ids参数指的是经过CUDA_VISIBLE_DEVICES过滤后的逻辑编号不是物理编号。这是最容易混淆的地方。举个例子物理服务器有4块GPU0,1,2,3你设置了CUDA_VISIBLE_DEVICES2,3那么在代码中device_ids[0,1]对应的就是物理GPU 2和33.3 关于主卡的注意事项主卡默认是device_ids[0]有几个特殊职责负责梯度汇总和参数更新承担更多的显存开销因为要保存完整的模型参数一些操作只能在主卡上执行这就引出了一个常见错误如果你指定的device_ids不包含主卡就会报错# 错误示例 os.environ[CUDA_VISIBLE_DEVICES] 0,1,2,3 model torch.nn.DataParallel(model, device_ids[1,2,3]).cuda()这会引发RuntimeError因为device_ids[0]主卡必须是可见的。4. 常见错误与解决方案4.1 Invalid device id错误这个错误通常发生在device_ids指定的编号超出了可见GPU范围。比如os.environ[CUDA_VISIBLE_DEVICES] 0,1 model torch.nn.DataParallel(model, device_ids[1,2]).cuda() # 错误这里device_ids[1,2]但可见的GPU只有两个逻辑编号0和1所以2是无效的。解决方案检查CUDA_VISIBLE_DEVICES的设置确保device_ids中的所有编号都小于可见GPU数量4.2 显存分配不均问题使用DataParallel时经常发现主卡的显存占用比其他卡多不少。这是因为主卡需要保存完整的模型参数输出结果会在主卡上汇总一些中间变量也存储在主卡上缓解方法使用更大的batch size让计算更均衡考虑使用DistributedDataParallel更高级的并行方式调整模型结构减少主卡负担4.3 数据加载的瓶颈当使用多GPU时数据加载可能成为瓶颈。我常用的优化方法是增加DataLoader的num_workers数量使用pin_memory加速数据传输预加载部分数据优化后的DataLoader示例from torch.utils.data import DataLoader train_loader DataLoader( dataset, batch_size256, shuffleTrue, num_workers8, pin_memoryTrue )5. 高级技巧与最佳实践5.1 混合精度训练多GPU训练时使用混合精度可以显著减少显存占用并加速训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, targets in train_loader: inputs inputs.cuda() targets targets.cuda() with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5.2 梯度累积当单卡batch size受限时可以通过梯度累积模拟更大的batch sizeaccumulation_steps 4 for i, (inputs, targets) in enumerate(train_loader): inputs inputs.cuda() targets targets.cuda() outputs model(inputs) loss criterion(outputs, targets) loss loss / accumulation_steps loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()5.3 模型保存与加载多GPU训练保存模型时要注意去掉module.前缀# 保存 torch.save(model.module.state_dict(), model.pth) # 加载 model Model() model torch.nn.DataParallel(model) model.load_state_dict(torch.load(model.pth))如果直接保存DataParallel包装后的模型会导致加载时出现问题。6. 监控与调试6.1 实时监控GPU状态训练时可以另开一个终端窗口使用以下命令监控watch -n 1 nvidia-smi这会每秒刷新一次GPU状态方便观察显存使用和利用率。6.2 PyTorch内置工具PyTorch提供了一些有用的工具函数# 检查CUDA是否可用 torch.cuda.is_available() # 获取当前设备 torch.cuda.current_device() # 获取设备名称 torch.cuda.get_device_name(0) # 清空缓存 torch.cuda.empty_cache()6.3 日志记录建议在代码中添加详细的日志记录方便调试import logging logging.basicConfig(levellogging.INFO) logger logging.getLogger(__name__) logger.info(fUsing GPUs: {torch.cuda.device_count()}) logger.info(fCurrent device: {torch.cuda.current_device()})7. 从DataParallel到DistributedDataParallel当你在单机多卡上熟练使用DataParallel后可能会遇到它的局限性只能单机多卡不能多机主卡负载不均衡效率不如DistributedDataParallel高升级到DistributedDataParallel的基本步骤import torch.distributed as dist dist.init_process_group(backendnccl) torch.cuda.set_device(args.local_rank) model torch.nn.parallel.DistributedDataParallel( model, device_ids[args.local_rank] )不过这个话题比较大值得单独写一篇教程来讲解。在实际项目中我发现对于大多数单机多卡场景DataParallel已经足够好用特别是快速原型开发阶段。