1. 动手提速之前先把瓶颈找准很多人一上来就折腾分布式训练、改网络结构结果折腾半天提速效果还不如人意。我见过不少同学跑训练时 GPU 利用率只有 20%还一个劲儿往模型里加注意力模块——方向完全反了。加速 PyTorch 模型训练这件事第一位不是调什么参数而是搞清楚你的训练时间到底花在了哪里。先做一个简单的耗时拆分。拿你现在的训练脚本在 DataLoader 取数据的部分、forward 和 backward 的部分、以及 optimizer.step 周围各打一个时间戳跑一个 epoch 看看占比。我习惯用time模块或者tqdm粗暴一点直接在循环里对每个阶段累加时间。如果数据加载占了 40% 以上说明瓶颈在数据管线不在 GPU 计算如果 GPU 利用率已经很高但整体时间还是很长那才是模型本身或者硬件算力的问题。判断 GPU 利用率有个更直接的办法。训练跑起来之后在另一个终端执行nvidia-smi -l 1观察 Volatile GPU-Util 这个字段。如果利用率经常在 90% 以上说明 GPU 被喂饱了这时候优化的重点应该转向混合精度、算子融合这些方向。如果利用率在 50% 以下乱跳说明 GPU 经常在等数据你把网络结构改成花都不会有本质提升先把数据管线搞定再说。还有一个很多人忽略的点IO 瓶颈不一定只在 DataLoader 里。如果你的训练数据是从网络盘或者机械硬盘读的每读一个 batch 都要卡一下这个延迟在nvidia-smi上看不出来但实际训练时间会拉得很长。我自己之前踩过这个坑数据放在 HDD 上step 耗时忽高忽低换了 SSD 之后训练时间直接少了四分之一。所以在动手之前我建议你先建立一个基准记录——把当前训练脚本跑一个固定 epoch记录总耗时、单 step 平均耗时、GPU 利用率和显存占用。后面每做一项优化就重新跑一次拿数据说话。下面的内容总结了一张自检表你可以对照自己的情况快速判断优先改哪里。现象可能的瓶颈优化优先级GPU 利用率低CPU 跑满数据加载跟不上最高先解决GPU 利用率高但单 step 慢计算本身或 IO 延迟高显存占用接近上限偶尔 OOM显存利用率不足或数据量太大中高GPU 利用率一般训练曲线波动大batch 太小或学习率不合适中以上都正常但还是慢模型结构过于复杂或硬件受限低2. DataLoader 优化常常被忽略的免费加速2.1 先解决 GPU 等数据的问题数据加载是 PyTorch 训练加速里性价比最高的优化点而且几乎没有风险只需要改几个参数。DataLoader 默认是单进程读数据CPU 读图、解码、归一化、转 Tensor 这一套流程全都串行GPU 算完一个 batch 之后要干等。这个等待时间在很多实际项目里占到总训练时间的 30% 到 50%——我在做目标检测和图像分类项目时都遇到过。最简单的改法是调num_workers。这个参数控制数据加载用的子进程数量多开几个进程并行读数据相当于把单线程的流水线变成多线程GPU 就不用老是等了。但这里有个常见的误区很多新手以为num_workers开得越大越好结果一开就是 16、32反而把 CPU 和内存占满了性能也不升反降。因为进程切换、锁竞争和内存带宽都是有开销的开到一定程度之后收益迅速衰减。num_workers的值具体怎么定我的经验是先看机器物理核心数从nproc或者任务管理器里确认然后取物理核心数的一半到四分之三作为起点。比如一台 16 核的机器设 8 到 12 个 worker 一般就够了如果你做的是简单的 NumPy 数组加载worker 开 4 个甚至 2 个就够了因为数据本身读取很快开太多反而浪费。之后你可以跑一个 epoch 比较时间逐步往上加两个找到拐点。pin_memoryTrue这个参数几乎可以无脑开。它会把数据放到锁页内存里让 GPU 可以直接通过 DMA 访问省掉一次 CPU 到 GPU 的隐式拷贝。特别在你用to(device)把数据搬到 GPU 上时这个参数带来的提升非常明显。代价是占用更多主机内存但只要内存没有爆建议一律开启。persistent_workersTrue则让 DataLoader 的子进程在每轮 epoch 之间保持存活省去反复创建进程的开销配合上面两个参数一起开实测数据加载部分能快 20% 到 30%。train_loader DataLoader( dataset, batch_size128, shuffleTrue, num_workers8, # 根据物理核心数调整 pin_memoryTrue, # 加速 CPU 到 GPU 的数据拷贝 persistent_workersTrue, # epoch 之间不销毁 worker 进程 prefetch_factor4, # 每个 worker 预取 4 个 batch )prefetch_factor是 PyTorch 1.11 之后开始稳定支持的参数它让每个 worker 提前预取多个 batch 的数据放到队列里进一步掩盖加载延迟。默认值通常是 2可以试着调到 4 或 6。不过这里有个前提如果样本预处理特别重比如超大分辨率图像加上复杂的增强worker 往往来不及预取调高prefetch_factor效果不明显你更需要的是增加num_workers。2.2 预处理太重怎么办很多实际项目的训练脚本里__getitem__方法中做了一堆事情——读图、随机裁剪、翻转、色彩抖动、转 Tensor、归一化。如果这些操作本身是纯 CPU 的即便开了多进程 worker每个 worker 也要花大量时间在图像解码和增强上。遇到这种情况有几个进阶操作可以试试。第一个思路是把图片提前解码成内存中的数组。如果数据集不是特别大比如几万张图完全可以在训练前把所有数据一次性读入内存。我做过一个几万张商品图的分类任务数据集总共不到 20GB直接全部加载进 RAM训练时__getitem__只做索引和增强速度比每次从磁盘读图快了好几倍。第二个思路是用 NVIDIA DALI。这个库能把图像解码和部分增强操作放到 GPU 上从数据读取那一步就开始用 GPU 并行。代价是 DALI 的数据管道和原生 PyTorch Dataset 接口不完全兼容需要额外写一些代码而且调试也不像原生那么直观。我的建议是只有当原生 DataLoader 优化到极限、数据加载仍然拖后腿时再考虑 DALI。还有一点如果数据集很大比如几十万张以上缓存整个数据集到内存不现实可以考虑把数据转成 WebDataset 或者 tar 包格式实现流式读取。这样每个 batch 从磁盘连续读取减少了随机 IO 寻址的开销。这个方案配合 DALI 或原生 DataLoader 都行但需要先做一轮数据格式转换。2.3 踩过的坑与心得关于 DataLoader我总结几条实战经验。第一num_workers不是越多越好。我有一次在 32 核的服务器上把num_workers设成 32结果每个 epoch 反而慢了很多后来发现是每个 worker 都在大量占用内存带宽加上系统频繁切换进程直接把性能拖垮了。后来降到 12速度立刻提了上来。第二如果你的数据读取逻辑里有随机的 NumPy 操作或者用了比较重的第三方库可以考虑把预处理的结果提前缓存下来用磁盘空间换训练时间。比如一次训练中数据集预处理后的 Tensor 是固定的就用torch.save存一份下次直接加载避免重复计算。第三验证阶段同样要优化。很多人只盯着训练 DataLoader验证集就随意设个num_workers0结果每个 epoch 结束验证要跑半天。验证 DataLoader 可以用同样的多进程参数只是不需要 shuffle设置shuffleFalse就行。这样整体的训练时间也能省下一大截。3. 混合精度训练快一倍还不掉精度3.1 为什么 FP16 能快这么多混合精度训练是 PyTorch 加速里效果最显著、改动最少的方案之一。它的核心思路是用 FP16半精度浮点来存中间激活值和做大部分矩阵乘法运算同时用 FP32 保存一份权重主副本只在梯度更新时用 FP32。为什么这样能快因为现代 GPU尤其是 NVIDIA Ampere 和 Hopper 架构的 FP16 算力通常是 FP32 的两倍甚至更多同时 FP16 占用的显存只有 FP32 的一半意味着更大的 batch size 也能塞进显存。AMPAutomatic Mixed Precision是 PyTorch 官方提供的自动混合精度工具它由两个核心组件构成torch.cuda.amp.autocast和GradScaler。autocast上下文管理器会自动判断哪些算子适合用 FP16哪些算子必须保持 FP32从而在不需要手改网络代码的前提下享受加速。GradScaler则负责处理梯度缩放为了防止 FP16 下梯度值过小导致下溢underflow为 0它会在反向传播前把 loss 放大若干倍计算完梯度后再缩小回来。用起来非常方便只需要修改训练循环里的一小段代码。下面是我在图像分类任务里最常用的一套写法scaler torch.cuda.amp.GradScaler() for images, labels in train_loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意这里和普通训练循环有三点不一样forward 必须包在autocast()上下文里反向传播用的是scaler.scale(loss).backward()优化器更新用的是scaler.step(optimizer)并且要调用scaler.update()来更新缩放系数。这 4 行代码的改动在支持 Tensor Core 的 GPU 上通常能带来 1.5 到 2 倍的训练速度提升而且精度基本不掉。3.2 BF16 是更好的选择吗如果你用的是 Ampere 架构之后的 GPU比如 A100、V100 比较老RTX 30 系、RTX 40 系、H100 这些那还有一个更值得尝试的选择BF16BFloat16。BF16 和 FP16 相比指数位和 FP32 一样多范围更大所以几乎不会出现溢出问题也不需要复杂的 loss scaling 逻辑。虽然在尾数精度上 BF16 比 FP16 差但得益于其超大的数值范围在很多模型上 BF16 的表现甚至比 FP16 更稳定训练曲线更平滑。PyTorch 里用 BF16 非常简单只需要在autocast里指定数据类型with torch.cuda.amp.autocast(dtypetorch.bfloat16): outputs model(images)如果 GPU 支持 BF16你会发现连GradScaler都可以省掉因为数值范围大了基本不会下溢。不过也要注意BF16 的加速效果依赖硬件支持在老的 GPU 上可能还不如 FP16。判断方式很简单执行torch.cuda.is_bf16_supported()返回True就直接用 BF16。3.3 什么情况不推荐开 AMP混合精度不是万能的。如果你的网络里有一些对数值精度极度敏感的算子比如某些自定义的损失函数、涉及到极小的浮点累加、或者使用了大量torch.exp、torch.log这类数值范围很大的函数开启 AMP 后可能会出现训练不稳定、loss 不收敛或者直接 NaN 的情况。检测方法很简单开启 AMP 后跑几十个 step观察 loss 曲线是不是比 FP32 时更抖。另外如果你的模型本身很小或者 GPU 利用率本来就不高比如小跑一个 MLPAMP 的收益可能很有限。因为 FP16 带来的加速主要在带宽密集和计算密集的大矩阵乘法上小模型体现不出来。我建议的做法是先开启 AMP 跑一个固定 epoch记录时间和精度再关掉对比一次用数据判断要不要开。在 AMP 训练中还有一个常见事故loss出现 NaN但GradScaler没报错。这种情况往往不是 AMP 本身的问题而是 loss 计算里有自定义函数把数值算爆了或者梯度裁剪和 loss scaling 的配合出了问题。处理梯度裁剪时注意要先把 loss 缩放回去或者使用scaler.unscale_(optimizer)否则裁剪的阈值是按未缩放梯度设计的数值会不对。正确写法是scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update()4. 训练策略与算子优化榨干 GPU 的每一分算力4.1 先调 batch size再谈其他如果 DataLoader 已经吃饱GPU 利用率也上去了下一步该看 batch size 的设定。这里有一个反直觉的规律很多人在显存允许的情况下batch size 明明可以开得很大却一直用 16、32 的保守值导致 GPU 利用率上不去。把 batch size 增大之后单位时间处理的总样本数会明显增加训练收敛速度也会变快因为算力被用得更满了。增大 batch size 的代价是可能会影响收敛的泛化性能。这个问题的标准处理方式是把学习率按比例往上调。经验法则是batch size 变成原来的 k 倍学习率大约也调成原来的 k 倍或者用 sqrt(k) 倍视任务而定。我有一个图像分类项目batch size 从 64 调到 256学习率从 0.01 调到 0.04训练到同样的测试精度总时间几乎减半。如果显存不够撑起更大的 batch另一个常用手段是梯度累积。逻辑是把一个大 batch 拆成几个小 batch分别计算梯度但不更新参数累加到一定步数后再一起更新。代码上也很简单核心是用optimizer.zero_grad()只在累积步时调用一次accumulation_steps 4 optimizer.zero_grad() for i, (images, labels) in enumerate(train_loader): with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) / accumulation_steps scaler.scale(loss).backward() if (i 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()需要注意这里把 loss 除以accumulation_steps是为了让梯度累积后的总幅度和直接用大 batch 时保持一致。梯度累积能在不增加显存的情况下模拟更大的 batch但这并不能直接提升 GPU 利用率只是解决了 batch size 受限的问题。如果你想要的是吞吐量提升焦点还是应该放在把真实 batch 尽量开大。4.2 torch.compilePyTorch 2.x 的免费午餐PyTorch 2.0 之后torch.compile可以说是最值得尝试的加速特性之一。它会把你的模型编译成优化后的计算图做算子融合Fusion和内核自动调优从而减少 kernel 启动开销、减少中间张量的读写。我以前在 A100 上测试过一个 ResNet 模型开启torch.compile后单 step 时间从 50 毫秒降到 40 毫秒左右提升差不多 20%。对于 Transformer 架构特别是注意力计算比较多的模型提升会更明显有时能达到 30% 以上。用起来很简单一行代码model torch.compile(model)但有些坑需要提前知道。第一第一次运行编译时会经历一个比较长的预热期因为系统在尝试不同的 kernel 组合。这个预热时间可能持续几十秒甚至几分钟很容易让人误以为卡死了其实耐心等一等就好。第二torch.compile对动态 shape 支持得不好如果你的输入尺寸在每个 batch 都不太一样比如 NLP 里不定长的序列编译缓存会被频繁重建速度反而变慢。这时候可以试试torch.compile(model, dynamicTrue)来声明支持动态形状但性能提升幅度会打折。第三如果你用了很多自定义的、purely Python 的控制流逻辑或者有一些不规范的 tensor 操作torch.compile可能直接报编译错误。遇到这种情况建议先把模型的 forward 尽可能简化成标准的 PyTorch 算子组合或者干脆放弃torch.compile用其他方式加速。在工程实践里我的判断标准是主流经典模型ResNet、ViT、BERT 类直接开遇到明显自定义的模型先试跑几个 step不报错再继续训练。4.3 一些小而有效的显存与算子优化除了上面这些大项还有一些细节优化加起来也能积少成多。第一个是torch.backends.cudnn.benchmark True。这个开关会让 cuDNN 自动选择最适合当前输入尺寸的卷积算法。如果你输入尺寸固定开启后能带来小幅但稳定的加速。代价是第一次迭代时会有额外的自动调优开销而且输入尺寸老是变的话反复调优反而拖慢速度。所以只建议在输入尺寸固定的场景下开。第二个是channels_last内存格式。对于卷积网络把 tensor 的内存布局从 NCHW 改成 NHWC可以在部分 GPU 上配合 FP16 获得额外的访存效率提升。PyTorch 里只需要model model.to(memory_formattorch.channels_last)数据也调用.to(memory_formattorch.channels_last)。我实测过在部分模型上配合 AMP 可以再多 3% 到 8% 的提升代价是兼容性上偶尔会踩坑。第三个是显存管理。很多人在 OOM 后习惯调用torch.cuda.empty_cache()这其实只是把显存还给 PyTorch 缓存池并不能真正提升训练速度。更实用的是设置PYTORCH_CUDA_ALLOC_CONF环境变量比如PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128可以减少显存碎片。如果你用的是 3090 或 4090 这种有 24GB 显存的卡显存碎片问题尤其明显设置了这个环境变量之后OOM 概率能降低不少。第四个是减少无谓的.item()和 GPU 到 CPU 的同步。训练循环里如果你频繁读取loss.item()或者每步都打印一堆指标会强制 GPU 和 CPU 同步严重拖慢训练。正确做法是把指标先累积在 GPU 上每 N 个 step 才输出一次。这个看似很小的改动在短 step 的训练里能省下 5% 到 10% 的时间。5. 分布式训练单卡瓶颈下的多卡破局5.1 从 DP 到 DDP为什么要换当单卡已经优化到极限还不够快时分布式训练就成了必经之路。PyTorch 提供了两种主要方案DataParallelDP和DistributedDataParallelDDP。很多入门教程喜欢教 DP因为代码改动少只要一行model nn.DataParallel(model)就行。但 DP 有个致命的问题它在一个进程里用多线程跑多卡每步前向和反向时都要把梯度或输出汇聚到主卡上主卡的内存和通信开销成了瓶颈多卡利用率很不均匀。DDP 则是多进程方案每张卡一个进程各自维护一份模型副本只在前向之前同步一次模型参数反向传播时不显式同步而是通过 all-reduce 在每张卡的梯度之间做一次通信然后各自独立更新参数。这样通信量小很多扩展效率远高于 DP。同样是 4 张卡DDP 在多多数情况下能跑到 3.5 倍以上的加速比而 DP 能到 2 倍已经很不错了。从代码角度看DDP 改动其实也不多。核心有三步初始化进程组、包装模型、用分布式采样器替换普通 DataLoader 的 shuffle。下面是一个最小可用的模板import torch.distributed as dist from torch.utils.data.distributed import DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(nccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) sampler DistributedSampler(dataset, shuffleTrue) train_loader DataLoader(dataset, batch_sizebatch_size, samplersampler) model model.to(local_rank) model DDP(model, device_ids[local_rank]) for epoch in range(num_epochs): sampler.set_epoch(epoch) # 保证每个 epoch 数据被打乱方式不同 for images, labels in train_loader: # 训练逻辑保持不变启动方式用torchrun最方便torchrun --nproc_per_node4 train.py这里有个细节我提醒一下大家DistributedSampler的set_epoch一定要在每次 epoch 前调用否则每个 epoch 的 shuffle 结果都一样模型的泛化能力会受影响。另外DDP 要求每个进程看到的数据是不同的所以 batch size 在 DDP 里是指每张卡上的 batch size总 batch 等于单卡 batch 乘以卡数这一点调整学习率时别漏掉。5.2 DDP 的通信开销和扩展性DDP 也不是万能的。当模型很小、单卡 step 时间很短比如几十毫秒时进程间通信相对成本就会比较高多卡加速比可能只有 1 倍、2 倍。这时候你需要先检查一下是不是batch_size太小导致每张卡上的计算时间太短。通常单卡 batch 至少要有几十上百DDP 的通信开销才能被计算时间掩盖。还有一个容易踩的坑DDP 打印日志时每个进程都会执行print如果你不加判断控制台会被刷屏。更麻烦的是有一些 checkpoint 保存逻辑如果每个进程都执行可能出现文件写入冲突。正确做法是只在主进程rank 0里执行打印和保存判断方式if dist.get_rank() 0: torch.save(model.module.state_dict(), best_model.pth) print(fEpoch {epoch} loss {loss})分布式训练这块我的经验是从 2 卡开始先验证代码能跑通、速度有提升再往 4 卡、8 卡扩展。不要一上来就搞 8 卡因为集群环境下的网络故障、显存不一致等问题排查起来非常耗时先把链路走通再说。6. 实战中的问题清单与避坑技巧6.1 排查实录加速优化过程中我踩过的坑不算少整理几个高频问题分享给各位。问题一开了多进程 DataLoader 后报File descriptor相关错误。这个通常是num_workers开得太大或者数据集内部有无法被 pickled 的资源比如打开的文件句柄、锁对象。解决办法是把num_workers调小或者在__getitem__里避免持有跨进程的句柄。如果你用了自定义的 Dataset 并且里面包含 lambda 函数或者局部类也会出现 pickling 错误改成模块级函数即可。问题二开启 AMP 后 loss 不下降或者出现grad scale is not finite警告。这是 GPA 缩放系数反复溢出导致的。处理方法包括降低初始init_scale比如GradScaler(init_scale2**10)检查学习率是否太大检查网络结构中是否有数值不稳定的自定义层。实在不行可以换 BF16或者退回 FP32先把模型训通再考虑优化。问题三DDP 训练时 GPU 利用率不均匀某张卡明显更慢。我遇到过 Graphcore 别的卡都在忙卡 3 利用率只有 50% 的情况。原因通常是数据不平衡比如某个进程处理的样本特别复杂或者机器 CPU 核数不够导致部分进程数据加载跟不上。先检查是不是num_workers设置得太小其次可以用DistributedSampler确保样本被均匀打散最后检查机器上是否有其他任务占用 CPU导致某些进程抢不到资源。问题四torch.compile之后显存反而变大了。编译模型有时会为了算子融合多缓存中间结果导致显存上升。如果显存本来就紧建议不要开torch.compile或者用torch.compile(model, modereduce-overhead)的变体它对显存更友好一些但加速幅度也会打折。6.2 我的个人经验总结做了这么多加速尝试之后我对优化顺序有了比较明确的把握。第一步永远是数据管线把num_workers、pin_memory、persistent_workers配上这步收益最快、风险最低第二步开混合精度改动小、收益明显第三步调整 batch size 和学习率把算力真正用满第四步用torch.compile锦上添花只有前面的都做完了还嫌慢才考虑上 DDP。我见过太多人跳过前两步直接上分布式结果数据加载跟不上多卡训练比单卡还慢——这真不是段子我就遇到过四卡训练因为数据管线没跟上反而比单卡慢了不少。加速的前提是单卡的瓶颈已经解决分布式才能发挥出真正的威力。另外每次改动最好只动一个变量测完再动下一个。这样如果出了性能回退或者精度问题你能立刻知道是哪个改动引起的回滚也方便。我自己习惯每次跑完记录一份日志包括 batch size、学习率、是否 AMP、是否 DDP、单 step 耗时、显存占用积累一段时间之后你对自己模型的调参边界会有非常清晰的感觉。