☰
PyTorch GPU加速比差异从哪来:训练循环、显存与多卡拆解
2026/10/26 5:45:29 网站建设 项目流程

显卡装进机箱、驱动跑通、torch.cuda.is_available()返回 True——很多人以为到这一步,PyTorch 中的 GPU 加速就自动生效了,模型训练速度自然会翻十几倍。我在几个项目里见过同一个现象:同一份代码,有人跑出 8 倍加速,有人跑出 0.7 倍加速,还有人加了一张卡之后训练时间反而涨了。差距不在运气,而在于你有没有搞清楚 GPU 到底在替谁干活、哪一段耗时真的被并行掉了。这篇就把 PyTorch 里 GPU 加速这件事从硬件特性、框架执行路径到训练循环逐段拆开,讲清楚加速比是怎么来的、怎么测才算数、显存这个硬天花板怎么绕,以及那些"上了卡反而更慢"的现场是怎么形成的。内容偏向已经能跑通训练脚本、想把训练时间真正压下去的读者,刚接触深度学习的人也能从前面几节拿到判断自己该不该用卡的依据。

1. 一次对照实验:快的到底是什么部分

1.1 同一份代码,两种截然不同的加速比

我印象最深的一次对比,是同一个小组里的两份作业。模型 A 是一个三层的全连接网络,输入维度 784,中间两层 512,输出 10,用一个小规模的手写数字数据集训练,batch size 设到 256。这台机器的 CPU 上跑一个 epoch 大概要五分钟出头,把模型和数据.to(device)之后,一个 epoch 掉到 25 秒左右。同组另一个同学做的是模型 B,两层小卷积核的卷积网络,输入尺寸 32×32,batch size 只设了 8,在 CPU 上跑一个 epoch 五十多秒,换到 GPU 上变成 30 秒上下,加速不到两倍。

硬件是同一张卡,框架版本一样,代码风格也接近,加速比差了四倍多。原因不神秘:模型 A 的每一层都是一个大矩阵乘法,计算量随维度呈立方增长,而搬运的数据量只随平方增长,这种活儿天生适合切成几万份同时干;模型 B 的 batch 只有 8,每层卷积的并行任务数被卡得很死,同时在跑的线程填不满显卡的计算单元,剩下的时间都花在启动和等待上。

我把这个现象总结成一句话,后面所有内容都围绕它展开:GPU 加速的收益等于"能被并行的计算量"减去"必须串行付出的开销"。你看到的加速比,是这两者的比值,不是一个固定常数。

1.2 用计算强度这个尺子去预判加速空间

判断一段计算值不值得搬上显卡,业内常用的指标叫计算强度,也就是每搬运一个字节的数据,能顺带做多少次浮点运算。显卡的显存带宽通常在几百 GB/s 到 TB/s 量级,而 CPU 和显卡之间的数据通道带宽只有几十 GB/s,两者差了一个数量级还多。这意味着什么?如果你的算子每读一个数只做一两次运算,那它就是在给数据通道打工,算力再强也发挥不出来;如果每读一个数能做几十上百次运算,通道的瓶颈就被掩盖了,算力才是主角。

矩阵乘法是典型的后者。一个 M×K 的矩阵乘 K×N 的矩阵,需要搬运的数据量是 MK+KN+MN,计算量是 2MKN。当三个维度都是 1024 时,计算强度大约在几百这个量级,远超通道能承受的阈值。卷积在通道数足够时也属于这一类。而逐元素操作,比如 ReLU、加法、LayerNorm 里的缩放,计算强度基本在个位数,这类操作在显卡上跑,性能曲线几乎是贴着带宽走的。

这张表可以拿来快速判断你模型里的瓶颈在哪:

操作类型计算强度大致量级通常受限于加速表现
大矩阵乘法、大通道卷积数百算力明显,数倍到数十倍
中等卷积、注意力中的批量矩阵乘数十算力与带宽混合明显,但受形状影响大
逐元素运算、归一化、激活函数个位数显存带宽有限,常被内核启动开销吃掉
索引、条件判断、循环控制接近零指令调度几乎没有,甚至变慢

实际网络里这几类操作是混在一起的,所以真实加速比总是介于最好和最坏之间。看到别人说"我的模型快了十倍",先去问他的 batch size 和模型结构,否则这个数字对你毫无参考价值。

1.3 加速比的分子分母要拆开看

把一次训练迭代的耗时写成加和形式,问题会清楚很多。在 CPU 上,迭代耗时大致是数据准备、前向计算、反向计算、参数更新的总和,这几段都在同一块芯片上串行执行。到了 GPU 上,式子变成数据准备、主机到设备拷贝、前向、反向、参数更新、以及若干次同步等待,其中后四项在设备上执行,但由 CPU 发起。

关键在于:只有"前向、反向、参数更新"这三项被真正加速了,数据准备和拷贝没有,同步等待反而是新增的。假设原始迭代耗时 100 个单位,其中计算占 70、数据占 30,计算部分被加速 20 倍,那么新耗时是 30 + 3.5 = 33.5,加速比约 3 倍,而不是 20 倍。这就是阿姆达尔定律在训练场景里的具体形态。

很多人第一次做 GPU 迁移时只盯着模型那几行代码,把.to(device)加上就以为完事了,结果数据管道还是单线程在 CPU 上做重增强,同步点一个没少。这时候模型本身快了二十倍,整体只快了三倍,然后得出"显卡没那么神"的结论。其实不是显卡不行,是你只优化了七成里的七成。

2. 张量搬上显卡之后,PyTorch 到底替你做了什么

2.1 从 host 到 device:一次显式拷贝和它的同步语义

tensor.to(device)或者早些年常用的.cuda(),做的是把数据从主机内存复制到显存。这个动作是同步的,函数返回时拷贝已经完成,主机端可以立刻读到结果。拷贝走的通道带宽有限,一个 100MB 的批次数据搬过去,按几十 GB/s 的有效带宽算,光拷贝就要几毫秒,如果这个批次的计算只需要一两毫秒,那拷贝就是纯亏。

默认情况下这个拷贝不带锁页内存优化,主机端的内存可能被操作系统换出,驱动需要先做一次中转,实际带宽会打折扣。给 DataLoader 打开pin_memory=True,数据会被放进锁页内存,驱动可以直接做异步传输,配合non_blocking=True使用,拷贝就能和前一次迭代的计算重叠起来。这是我在所有涉及大批量图像或长序列的项目里必开的一个开关,改动量只有两行,收益在数据重的任务上非常明显。

需要提醒的是,non_blocking=True只对锁页内存生效,如果源头还在普通内存里,这个参数就是个装饰品,甚至可能因为语义误判导致数据还没到就被读取。判断方法很简单:开这个参数后如果训练结果出现莫名其妙的偏差,第一件事就是检查pin_memory有没有一起打开。

2.2 异步执行:CPU 扔完指令就走,不等结果

PyTorch 在显卡上的大部分操作是异步的。你写下y = model(x),CPU 把内核启动指令塞进队列就返回了,真正的计算在设备上排队执行。这个设计让 CPU 可以提前准备下一批数据、发起下一次拷贝,形成流水线。

异步带来一个副作用,也是我见过最多的测量错误:用time.time()包住一段显卡代码,测出来的往往是"CPU 把指令发完的时间",而不是"计算做完的时间"。一段真实需要 80 毫秒的前向,可能测出来只有 3 毫秒。等你拿着这个数字去估算总训练时间,误差会大得离谱。

正确的做法是在计时结束前显式同步一次,torch.cuda.synchronize()会阻塞到队列里所有任务完成。代价是破坏了流水线,所以只在测量时加,不要塞进训练循环里。这一点后面第 4 节会展开讲测量方法。

2.3 内核启动开销:为什么小算子堆多了会变慢

每次在显卡上启动一个计算内核,都有固定开销,量级在几微秒。单看很小,但一个 Transformer 层的反向传播可能涉及几十上百个内核,一个 epoch 几千次迭代,累积起来就是可观的数字。当单个内核的计算时间只有一两微秒时,启动开销的占比会超过一半,此时并行度再高也没用。

这个机制解释了一个反直觉的现象:把一个大矩阵乘法拆成十个小矩阵乘法,理论浮点运算量不变,实际耗时可能翻倍。也是为什么逐元素操作多的模型,在显卡上的表现经常不如预期——它不是一个活儿,是几百个小活儿,每个都要单独排队。

框架层的应对手段是把相邻的逐元素算子融合成一个内核,PyTorch 2.x 的torch.compile就在做这件事,把 Python 层面的多次调用编译成少数几个融合内核。我在一个以 LayerNorm 和激活函数为主的模型上试过,开启编译后单步耗时下降了约三成,代价是首次编译要等一两分钟,且遇到动态形状会触发重新编译。是否值得开,取决于你的模型是静态形状为主还是变长输入为主。

2.4 显存分配器:为什么第二次迭代比第一次快

第一次迭代往往比后续迭代慢一截,这不是错觉。显存分配器在首次遇到某个尺寸的张量时会向驱动申请新的一块,后续再遇到同样尺寸就直接从缓存池里取,省掉了申请和释放的系统调用。这个缓存机制叫缓存分配器,它只增不减,池子会一直占着不放。

理解这一点有两个实际用处。一是测速时必须先跑几轮热身,把池子预热到稳定状态,否则测出来的是分配器冷启动的性能。二是当你看到显存占用居高不下、想手动回收时,torch.cuda.empty_cache()只能把空闲块还给驱动,正在被张量引用的部分动不了,而且频繁调用会破坏缓存池、让后续分配变慢。我一般只在两个 epoch 之间、或者切换模型结构时调一次,训练循环里绝不放。

3. 把训练循环拆成零件,看加速比被谁吃掉

3.1 前向、反向、更新三段各自的性格

前向传播是计算密度最高的一段,卷积和矩阵乘集中在这里,显卡的算力主要靠它发挥。反向传播的计算量大约是前向的两倍,因为它要算两份梯度,一份对输入、一份对权重,同时还要把前向保存的中间激活读回来,所以它对带宽的压力也更大。参数更新这一步,如果是 SGD,就是一次逐元素乘加,计算强度极低;如果是 Adam,还要读一阶矩、二阶矩两份状态,读写量是参数量的好几倍,这段在显卡上跑纯属浪费算力,但因为要和参数放在同一设备上才能做原地更新,一般也就顺手在显卡上做了。

三段的比例决定了加速上限。卷积网络里前向反向占绝对大头,迁移收益高;而一些小模型或者以嵌入查表为主的推荐模型,更新段和查表段的占比反而更高,这些段落对显卡不友好,加速比自然上不去。

一个粗略的经验:如果你的模型参数量在百万级以下、且以逐元素运算为主,先别急着上卡,把 CPU 端的向量化做好可能收益更直接。参数量上到千万级、以矩阵乘和卷积为主,显卡的性价比才开始显现。

3.2 数据管道:最容易被忽视的那半条流水线

训练循环里有一段代码完全不碰显卡,就是取数据。for x, y in dataloader这一行背后,可能是读图、解码、裁剪、翻转、归一化一整套流程。如果num_workers=0,这些活儿全在主进程里串行做,显卡就在那儿干等着,利用率曲线会呈现明显的锯齿状——算一下、停一下、算一下、停一下。

把num_workers调到 4 到 8,让子进程并行准备数据,锯齿就会被填平。这个数字不是越大越好,每个 worker 都要占内存和 CPU 核,开太多反而会引起进程调度争抢,我一般按 CPU 物理核数的三分之一到二分之一来设,再根据实际利用率微调。

另外两个参数值得一起配:pin_memory=True让数据进锁页内存,为异步拷贝铺路;persistent_workers=True避免每个 epoch 结束时销毁重建 worker 进程,在数据量小、epoch 多的情况下能省下可观的启动时间。这三个参数组合起来,在图像类任务上我见过数据侧耗时下降六成的情况。

3.3 同步点:那些偷偷让流水线停摆的调用

异步执行最怕被打断。训练循环里每出现一次需要读回主机端数值的操作,前面排队的任务就必须全部做完,流水线清空一次。最常见的触发点是打印损失:print(loss)或者loss.item(),这会把一个标量从显存搬回主机,强制同步。放在每个 iteration 里,等于每步都踩一次刹车。

我在别人的代码里见过更隐蔽的写法:用一个列表把每步的 loss 收集起来,最后统一画图。如果收集的是张量本身,那些张量会一直挂在计算图上,显存越占越多,几十步之后就爆了。正确的做法是.item()之后存浮点数,而且要控制频率,比如每五十步记一次。日志、指标统计、学习率调度里的条件判断,凡是需要读数值的地方,都要有意识地减少次数。

还有一类同步来自显存与内存之间的来回搬运。tensor.cpu()、tensor.numpy()、在显卡张量上做 Python 的if判断,每一个都是潜在的停顿点。写代码时养成一个习惯:能在设备上算完的,就别中途搬回来。

3.4 用一张表把耗时归位

排查加速不达预期时,我习惯先把一轮迭代的各个环节按"在哪个设备上执行""是否引入同步"分类,做成下面这样的表,再逐项去测:

环节执行设备是否引入同步常见优化手段
数据读取与增强CPU 多进程否提高 worker 数、简化增强、预解码
主机到设备拷贝拷贝引擎视是否锁页而定pin_memory+non_blocking
前向计算显卡否提高 batch、融合算子、混合精度
反向计算显卡否同上,注意激活占用
参数更新显卡否融合优化器、梯度累积降低频率
指标记录与打印主机是降低频率、先累积后统一搬回
检查点保存主机是降低频率、异步保存

这张表填完,通常能一眼看出哪一格是红的。如果拷贝和同步占了三四成,说明问题在管道不在算力;如果前向反向本身就慢,那才轮到讨论模型结构和混合精度。

4. 测量方法:让跑出来的数字站得住脚

4.1 热身、同步、重复,三步缺一不可

我在第 2 节提过异步执行和分配器冷启动这两个坑,它们共同决定了测速脚本必须做三件事。第一,热身至少十轮,把显存池和计算图缓存都预热到稳定状态;第二,每轮计时前后加同步,确保测的是真实完成时间;第三,重复多轮取中位数,而不是平均值,因为偶发的系统调度抖动会把均值拉偏。

用事件计时的写法比time.time()更贴近设备真实耗时,因为它记录的是设备时间线上的时间戳:

import torch def benchmark_step(step_fn, warmup=10, iters=50): for _ in range(warmup): step_fn() torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) times = [] for _ in range(iters): start.record() step_fn() end.record() torch.cuda.synchronize() times.append(start.elapsed_time(end)) times.sort() return times[len(times) // 2]

注意start.record()要在同步之后调用,如果前面还有没做完的任务,第一次的计时会被污染。另外step_fn里不要包含打印和保存,那些是测量噪声。

4.2 显存要看峰值,不看瞬时值

nvidia-smi显示的显存占用是进程级的瞬时快照,采样间隔还大,很容易漏掉峰值。训练过程中真正决定你能不能跑起来的,是峰值显存,因为一旦某个瞬间超了,就是 OOM。PyTorch 提供了记录峰值的接口,需要在训练开始前重置一次:

torch.cuda.reset_peak_memory_stats() # ... 训练若干步 ... peak = torch.cuda.max_memory_allocated() / 1024**3 print(f"peak allocated: {peak:.2f} GB")

max_memory_allocated统计的是被张量实际占用的部分,而max_memory_reserved统计的是缓存池向驱动申请的总量,后者通常更大。排查 OOM 时两个都要看:如果 allocated 远小于 reserved,说明是碎片问题;如果两者接近且都顶到卡的上限,那就是真的不够用。

4.3 用 profiler 找到真正的热点

前面几节讲的是方法论,落到具体模型上还是得靠数据说话。torch.profiler能给出每个算子占用的设备时间、调用次数、以及主机端的调度耗时,是我定位性能问题的主力工具:

from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: for _ in range(5): step_fn() print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

看输出时我会关注三件事:排在前几位的算子是不是预期的计算主力、有没有大量耗时极短但调用次数惊人的小算子、以及 CPU 总时间是否显著高于设备总时间。最后一条如果成立,说明主机端成了瓶颈,通常意味着数据管道或者 Python 层的控制流太重。

一个真实的排查经历:某个模型的单步耗时比预期多了四成,profiler 里排第一的不是卷积,而是一个 reshape 之后紧跟 transpose 的组合,调用次数上千。原因是代码里在每次前向时都做了一次不必要的维度调整,改成在初始化时调好就再没出现过。这种问题靠读代码很难发现,靠 profiler 一眼就看出来了。

5. 显存:比算力更早到顶的那堵墙

5.1 显存被谁吃了:一份粗略的账

很多人以为模型有多少参数就占多少显存,实际占用是参数量的好几倍。以标准的全精度训练为例,每个参数要占:参数本身 4 字节、梯度 4 字节、如果用的是 Adam 还要存一阶矩和二阶矩各 4 字节,加起来 16 字节。一个一亿参数的模型,光这些就是 1.6GB。

真正的大头往往是中间激活。前向过程中为了反向传播要用到的中间结果,都必须留在显存里,它的规模正比于 batch size、序列长度、通道数和层数。一个视觉模型在 batch 64 时,激活占用可能比参数占用还大。此外还有计算过程中的临时工作区、通信缓冲区、以及框架自身的一些开销。

所以估算显存要按"参数 + 梯度 + 优化器状态 + 激活峰值 + 余量"来算,余量至少留两成。只按参数量估算,是新手最容易犯的错,也是最容易在训练跑到一半时炸掉的隐患。

5.2 混合精度:用一半的位宽换一半的搬运量

把前向和大部分运算降到半精度,是当前性价比最高的显存与速度优化手段之一。半精度张量占的字节数是全精度的一半,搬运量减半,而且在支持张量核心的硬件上,半精度矩阵乘的吞吐可以是全精度的数倍。

PyTorch 里的用法是torch.autocast配合GradScaler。自动混合精度会在合适的算子上自动降精度,同时把累加、归一化、损失计算这些对数值范围敏感的操作保持在单精度,避免溢出。GradScaler负责在反向传播前把损失放大,防止半精度下的梯度下溢成零,更新参数前再缩回来:

scaler = torch.amp.GradScaler("cuda") for x, y in loader: optimizer.zero_grad(set_to_none=True) with torch.autocast(device_type="cuda", dtype=torch.float16): out = model(x) loss = criterion(out, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

几个实操经验。optimizer.zero_grad(set_to_none=True)比默认的置零更省,因为它把梯度张量释放掉而不是写零,省了一次全量写显存。GradScaler只在反向和更新时用,前向在autocast上下文里就行。如果训练中出现连续的梯度缩放因子下降,说明有数值不稳定,通常是某个自定义算子在半精度下溢出了,可以给那个算子单独加autocast的禁用装饰器。

半精度带来的加速在不同模型上差异很大。计算密集、矩阵乘为主的模型收益明显;而本来就受限于带宽的模型,收益主要来自搬运量减半,幅度会小一些,但显存占用下降是实打实的,能让你把 batch 调大一倍,间接又把并行度提上去了。

5.3 梯度累积与激活重算:用时间换空间的两条路

显存不够时,最直接的想法是减小 batch,但 batch 太小会让显卡填不满,加速比掉下来。梯度累积提供了一个折中:用小 batch 跑多次前向反向,把梯度累加起来,再统一更新一次参数。这样等效的 batch 变大了,但峰值显存只由单次前向的 batch 决定。

accum_steps = 4 for i, (x, y) in enumerate(loader): with torch.autocast(device_type="cuda", dtype=torch.float16): loss = criterion(model(x), y) / accum_steps scaler.scale(loss).backward() if (i + 1) % accum_steps == 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True)

注意损失要除以累积步数,否则梯度会等比例放大,等效学习率就变了。另外 BatchNorm 这类依赖批次统计的层,在梯度累积下统计的是小批次的分布,和真正的大批次不完全等价,如果模型里有这类层,需要评估一下影响。

另一条路是激活重算,在前向时不保存中间激活,反向时按需重新算一遍。代价是多一次前向的计算,收益是激活占用大幅下降,在深层模型上往往能把显存砍掉一半以上。PyTorch 提供了torch.utils.checkpoint,包在需要的层外面即可。我一般只在显存实在吃紧时才用,因为它会让单步耗时上升两到三成,等于用时间买空间。

5.4 碎片问题:一个容易被误判的 OOM

有一种 OOM 很迷惑:报错说显存不足,但nvidia-smi显示还有空闲,而且把 batch 调小一点又能跑。这种情况多半是碎片导致的——缓存池里有很多零散的小块,凑不出一块连续的大块给新的分配请求。

处理办法是设置分配器的切分策略,通过环境变量控制最大可切分块的大小,让大块尽量保持完整。这个参数需要在进程启动前设置,PYTORCH_CUDA_ALLOC_CONF支持按最大切分块大小来限制,具体数值要结合你的张量尺寸来调,没有通用值。另一个更省事的办法是在训练开始前做一次足够大的预热分配再释放,把池子的块尺寸结构定下来。

6. 那些加了显卡反而变慢的现场

6.1 小 batch 撑不起并行度

这是最常见的一类。batch 只有 4 或 8,每一次前向的计算量太小,显卡的计算单元填不满,大部分时间花在内核启动和调度上。同一份代码在 CPU 上跑,虽然单次慢,但至少没有启动开销,所以测出来的比值可能小于一。

判断方法很直接:固定其他条件,把 batch 从 8 逐步调到 64、128,看单步耗时怎么变。如果单步耗时几乎不随 batch 增长而增长,说明你处在启动开销主导的区间,增加 batch 几乎免费;当单步耗时开始随 batch 线性上升,才说明计算单元被填满了。跨过那个拐点,加速比才有意义。

实在受限于显存或者任务本身(比如在线学习、序列极长的生成任务),可以考虑第 5.3 节的梯度累积,或者干脆接受小 batch 的现实,把优化重心放到数据管道和同步点上。

6.2 显卡利用率曲线上的锯齿说明了什么

用显卡状态查看命令盯着利用率看,如果曲线是接近满格的平稳线,说明计算单元基本没闲着;如果是规律性的锯齿,说明它在等数据;如果是忽高忽低的大起大落,可能有同步点或者周期性的大开销操作。

锯齿最常见的原因是数据管道。看到锯齿先做两件事:把num_workers加上去,把数据增强里那些用 Python 循环实现的、纯 CPU 单线程的操作换成向量化实现或者更高效的库。我在一个项目里见过用 PIL 逐张做随机裁剪的写法,改成批量的张量运算之后,数据准备时间从每批 40 毫秒降到 6 毫秒。

大起大落则要怀疑周期性的重活,比如每 N 步保存一次检查点、每轮结束做一次完整验证集评估。这些操作本身不慢,但它们会打断流水线,让显卡空转。把评估和保存的频率降下来,或者放到单独的后台流程里,曲线会平滑很多。

6.3 模型太小,迁移成本大于收益

有些任务的模型只有几十万参数,输入也很小,比如简单的表格数据分类、短文本分类。这类模型在显卡上的计算时间可能只有一两毫秒,而主机到设备的拷贝、内核启动、同步加起来的固定开销就要好几毫秒。搬上卡,总耗时反而增加。

我的判断标准是看单步耗时能否稳定超过十毫秒。低于这个量级,固定开销的占比太高,不值得折腾。高于这个量级,且随着 batch 增大耗时增长不明显,那就值得上卡。这条标准不绝对,但能帮你快速筛掉一批不适合的场景。

还有一种情况是推理而非训练。单条样本的推理,batch 为 1,显卡的并行度完全用不上,延迟未必比 CPU 低。这时候要么做批处理攒够样本一起推,要么就老老实实用 CPU,别为了"用上显卡"而用显卡。

6.4 设备之间来回搬数据

代码里如果出现x.cpu()之后再x.to(device)的往返,每一次都是一次完整的同步加两次拷贝。我见过在训练循环里把特征搬到 CPU 做拼接、再搬回显卡的写法,光这一步就让单步耗时涨了三倍。后来改成在显卡上直接用拼接操作,耗时几乎归零。

同类问题还有在显卡张量上做 Python 条件判断,比如if tensor.max() > 0,这会触发同步把数值读回来。改成用张量运算表达同样的逻辑,把分支消除掉,或者把判断频率降下来,都能解决。

7. 多卡之后,加速比为什么不是线性的

7.1 数据并行做了什么,通信开销从哪来

单卡不够用的时候会想到多卡。数据并行的基本做法是每张卡持有一份完整的模型副本,各自吃不同的数据切片,算完梯度后做一次全局归约,把各卡的梯度求平均,再各自更新参数。听起来很自然,但那次全局归约就是瓶颈。

归约的通信量正比于参数量,而计算量正比于 batch size 乘以模型规模。卡越多,每张卡分到的数据越少,计算时间越短,但通信时间不降反升。当计算时间短到和通信时间同一量级,加卡就没有收益了,甚至因为同步等待而变慢。

缓解手段主要有两个方向。一是提高单卡的计算量,比如每张卡上仍然保持较大的 batch,让计算时间远大于通信时间;二是用计算通信重叠的技术,在反向传播逐层算梯度的时候就异步发起对应层的归约,让通信藏在计算后面。后者在框架提供的分布式封装里已经默认做了,前提是你用的是推荐的接口而不是早期那个单进程多线程的实现。

7.2 接口选择:优先用分布式的那一套

PyTorch 里有两套数据并行接口。早期那套多线程实现只支持单机多卡,全局归约走主机内存,效率低,而且和梯度累积、混合精度混用时容易出现精度问题,官方已经明确不推荐。现在推荐的是分布式数据并行,每个进程管一张卡,进程间用高效的集合通信库通信,支持多机多卡,性能也更好。

从旧接口迁移到新接口的改动量其实不大,主要是进程启动方式变了,需要用启动器拉起多个进程,每张卡一个。代码层面主要改三处:初始化进程组、给模型包上分布式封装、把数据加载器换成带分片采样的版本。我建议新项目直接上分布式那套,别为了省事用旧接口,后面想扩展到多机的时候会推倒重来。

7.3 卡数上去之后,学习率要不要跟着调

这个问题在实践里经常被忽略。数据并行下,全局 batch 是单卡 batch 乘以卡数,等效的梯度噪声变小了,如果学习率还按单卡的设置来,收敛会变慢,甚至卡在一个平坦的区域出不来。常见的处理是按 batch 放大的比例线性放大学习率,同时加几轮热身让训练初期更稳。

线性放大不是万能的,batch 大到一定程度后收益会递减,此时可以改用平方根缩放。我自己的做法是先按线性放大试,观察前几百步的损失曲线,如果下降明显比单卡时慢,就往回收一点。这个调整没有标准答案,但一定要意识到"加卡之后超参需要重新看"这件事,直接沿用旧配置是最容易出问题的做法。

最后分享一个我在多个项目里反复验证的小习惯:每次改完训练脚本,先别急着跑完整训练,用固定的随机种子跑二十步,记录单步耗时的中位数和峰值显存,和改动前的基线对比。这个动作花不到两分钟,但能挡住绝大多数"改了半天其实更慢"和"跑到第三轮才 OOM"的情况。训练速度这件事,最怕的不是优化手段不够,而是凭感觉判断快慢。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询