3. GPU 异步计算

AI AI Infra
学习笔记- AI Infra
Publish Date:   2026-05-08

写 PyTorch 训练代码时,你会发现一些奇怪的现象:loss.backward() 后面紧跟一行 print(loss.item()),这一行突然变得很慢;明明用了 torch.cuda.synchronize() 测时间,profiler 里看到的耗时却完全不一样;DDP 训练里通信看似在反向之后才发生,实际却和反向”同时”在跑。

这些都是同一个东西的影子——异步计算 (asynchronous computation)。它不是某个高级 API,而是 GPU 编程模型的底层默认行为。理解它能解开很多”为什么训练比想象中快/慢”的谜团,也是看懂 DDP 计算-通信重叠、ZeRO-3 prefetch、Pipeline 调度这些大模型工程优化的前置知识。

本文按这个顺序展开:先讲清楚同步 / 异步的本质区别,再讲 CUDA 的 stream 模型,然后用 PyTorch 代码看实际怎么用,最后看三个典型场景(DDP / ZeRO-3 / Pipeline)是怎么把异步玩到极致的。

一、为什么需要异步:硬件本来就是并发的

1.1 一台机器里有几个”工人”

跑深度学习的服务器,至少有这些彼此独立的执行单元:

工人 干的事 谁指挥它
CPU 核心 Python 代码、数据预处理、kernel launch 你写的 Python
GPU SM(流处理器) 矩阵乘、conv、激活函数等计算 kernel CUDA driver
GPU DMA / Copy 引擎 Host↔Device、Device↔Device 数据搬运 CUDA stream
NVLink / PCIe / IB 网卡 跨 GPU、跨机通信 NCCL

关键事实:这些工人物理上是并行的——CPU 在跑 Python 时 GPU 不一定闲,GPU 在算 kernel 时网卡完全可以同时传梯度,DMA 在搬数据时 SM 也能继续算别的。

如果代码写得”全同步”(每一步都等上一步彻底干完才发下一条指令),就等于让其他工人都干等。异步就是为了让所有工人尽可能同时忙起来。

1.2 同步 vs 异步:一段代码看本质

考虑这样一个流程:算梯度 → 把梯度传给其他卡 → 用平均梯度更新参数。

同步版本(每一步都阻塞):

1
2
3
4
5
时间 →
CPU/GPU :  [算梯度          ][        ][用平均梯度更新参数]
网卡    :                    [传输/平均]
                                 ↑ GPU 此时干等
总耗时  :  T_compute + T_comm + T_update

异步版本(算下一批的同时把上一批的梯度传出去):

1
2
3
4
5
时间 →
CPU/GPU :  [算梯度       ][算下一批梯度       ][用平均梯度更新]
网卡    :                 [传输/平均上一批的梯度]
                          ↑ 同时进行,互不阻塞
总耗时  :  max(T_compute, T_comm) + T_update

如果 T_compute ≈ T_comm,异步版本几乎省掉一整段通信时间。这就是 DDP 跑得近似线性加速的根本原因——不是通信变快了,而是通信被计算掩盖了。

1.3 但异步不是免费的

异步带来三个新问题,后面每一节都会反复碰到:

  1. 依赖管理:某个 kernel 需要前一个 kernel 的结果,异步发射时怎么保证顺序?
  2. 资源竞争:两个 stream 都在用同一份显存怎么办?
  3. 可观测性:CPU 早就跑过那行代码了,GPU 实际还没开始,你怎么测时间、怎么处理报错?

CUDA 的 stream 模型就是为了解决这三件事而设计的。

二、CUDA 编程模型:Stream 与异步执行

2.1 两个时间轴:Host 与 Device

在 CUDA 里,CPU(host)和 GPU(device)有各自独立的指令队列。你在 Python 里调用一个 GPU 操作,实际上分两步:

  1. Host 提交 (launch):把”做什么”写一条指令丢进 GPU 的命令队列,立即返回
  2. Device 执行:GPU 在某个时刻按队列顺序实际做这件事
1
2
3
4
5
                    时间 →
Host (CPU)  :  [launch K1][launch K2][launch K3][launch K4][... 继续 Python ...]
Device (GPU):              ↓ 排队
                          [    K1 执行    ][    K2 执行    ][    K3 执行    ][K4 ...]
                            ↑ 这里 host 早跑过 launch K1 那行代码了

Host 的时间和 Device 的时间是错开的——这是 CUDA 异步模型的最核心事实。

代码层面看:

1
2
3
4
5
6
7
8
9
import torch

x = torch.randn(4096, 4096, device="cuda")
y = torch.randn(4096, 4096, device="cuda")

t0 = time.time()
z = x @ y                      # ← 这一行立刻返回!GPU 还没算完
t1 = time.time()
print(t1 - t0)                 # 通常 < 1ms,这是 launch 时间

x @ y 看起来像在算矩阵乘,实际只是把矩阵乘任务提交给 GPU。如果你想测真实计算时间,必须强制等 GPU 干完:

1
2
3
4
5
6
torch.cuda.synchronize()       # 阻塞 host 直到 GPU 队列清空
t0 = time.time()
z = x @ y
torch.cuda.synchronize()       # 再阻塞一次,确保 z 算好
t1 = time.time()
print(t1 - t0)                 # 现在才是真实计算时间

这是 99% 的”GPU 测时间错”的根源——很多人以为代码慢的地方,其实 host 早就跑过了,真实瓶颈在另一个地方。

2.2 Stream:GPU 的指令队列

每张 GPU 上不只有一条命令队列,而是有多个独立的队列,每个叫一个 stream。

  • 同一个 stream 内:任务严格按提交顺序串行执行(有依赖关系)
  • 不同 stream 之间:任务可以并发执行(相互独立)

PyTorch 默认所有操作都跑在一个”默认 stream”上,所以你写的代码看起来是顺序的。但只要你把不同任务派到不同 stream,GPU 的硬件就能让它们并行起来。

1
2
3
4
默认 stream     :  [matmul][add ][relu][matmul][...]    ← 所有计算排成一队
自定义 stream 1 :        [memcpy h2d][memcpy h2d]       ← 数据搬运并发进行
自定义 stream 2 :              [NCCL AllReduce]         ← 通信并发进行
                Time →

GPU 内部:SM 跑计算 kernel,Copy 引擎跑 memcpy,NVLink 引擎跑通信——三组硬件资源,三个 stream 各占一个,真的同时在跑。

2.3 不同 stream 之间的并发示意

举一个具体例子:数据搬运 + 计算重叠。

1
2
3
4
5
时间轴      0    1    2    3    4    5    6    7
默认 stream :           [matmul A    ][matmul B    ]
copy stream :  [H2D x→A][H2D y→B   ]
                   ↑ 这两件事在 GPU 上同时进行
                   一个用 SM,一个用 DMA,不抢资源

如果你不开 copy stream,数据搬运和计算就会串行:

1
2
3
时间轴      0    1    2    3    4    5    6    7    8    9
默认 stream :  [H2D x→A][matmul A    ][H2D y→B][matmul B    ]
                                        ↑ 此时 SM 闲着等 DMA

直观就能看到差距——多 stream 是在用已经存在但没被利用的硬件并行性。

2.4 同步原语:让乱序的世界保持正确

有了多 stream,就要回答”怎么保证依赖顺序”。CUDA 提供几把锁:

torch.cuda.synchronize():阻塞 host,直到 device 上所有 stream 都干完。最暴力,主要用来测时间和 debug。

stream.synchronize():阻塞 host,直到这一个 stream 干完。

stream.wait_stream(other):让 stream 在 GPU 内部等 other 干完才继续——不阻塞 host,host 该干嘛干嘛。这是异步重叠的关键 API。

Event:更细粒度的标记。在某个 stream 上 event.record() 打个标记,另一个 stream event.wait() 等这个标记被达到。常用于建跨 stream 的依赖图。

1
2
3
4
5
6
7
8
9
10
11
12
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

with torch.cuda.stream(s1):
    a = compute_A()           # 在 s1 上算

ev = torch.cuda.Event()
ev.record(s1)                  # 在 s1 上打标记

with torch.cuda.stream(s2):
    ev.wait()                  # s2 等 s1 的标记
    b = use(a)                 # 现在 s2 才用 a,保证 a 算完

整个过程host 不阻塞,只是在 GPU 那边建了”s2 的某个 kernel 必须等 s1 那个 event 才能跑”的依赖。这就是异步重叠的标准写法。

三、PyTorch 里怎么用异步

3.1 你早就在用异步了

只要你在 PyTorch 里跑 .cuda() 张量做运算,默认就是异步的。下面这段代码看起来”按顺序”运行:

1
2
3
4
y = model(x)        # forward
loss = criterion(y, target)
loss.backward()     # backward
optimizer.step()    # update

实际上,这四行代码执行完时,GPU 上的工作可能一个都还没真正做完——它们都被排进了默认 stream 的队列里,host 已经跑到下一行了。直到你做”必须看到结果”的事(打印、保存、转 CPU),host 才会被迫等 GPU。

3.2 哪些操作会”强制同步”

这些是 host 与 device 的同步点,会让 host 阻塞等 GPU:

操作 为什么阻塞
tensor.item() 必须把数取回 CPU 内存
tensor.cpu() / .numpy() 数据搬到 CPU,要等 GPU 写完
print(tensor) 打印要看到值
tensor.tolist() 同上
if tensor > 0: 条件判断需要值
torch.cuda.synchronize() 显式阻塞

最常见的 footgun:训练循环里手贱写一个 print(loss.item()),每个 iteration 都会逼 host 等 GPU 队列清空,本来异步重叠的计算-通信瞬间变同步。生产代码里这种打印要么定期(每 100 步)做,要么累积成 tensor 最后再 .item()。

3.3 数据搬运:non_blocking=True

最容易拿到的”免费”异步收益是数据加载。默认情况下 tensor.to('cuda') 是同步的(host 等 H2D 拷贝完),但加一个参数就异步了:

1
2
3
4
5
6
7
8
9
10
11
# DataLoader 用 pin_memory=True 才能开 non_blocking
loader = DataLoader(dataset, ..., pin_memory=True)

for x, y in loader:
    x = x.to(device, non_blocking=True)   # 异步 H2D
    y = y.to(device, non_blocking=True)
    
    # 此时 H2D 还在 copy stream 跑,但下面的 forward 在默认 stream
    # 默认 stream 看到 x,y 就会自动等 copy 完成(PyTorch 帮你管依赖)
    pred = model(x)
    loss = criterion(pred, y)

这样下个 batch 的搬运可以和当前 batch 的计算重叠,DataLoader 不再是瓶颈。

3.4 测时间的正确姿势

既然 time.time() 测的是 host 时间,不是 GPU 时间,那怎么准确测?三种方法:

1
2
3
4
5
6
7
# 方法 1:最简单,粗粒度
torch.cuda.synchronize()
t0 = time.time()
heavy_gpu_work()
torch.cuda.synchronize()
t1 = time.time()
print(f"{t1 - t0:.3f}s")
1
2
3
4
5
6
7
8
9
10
# 方法 2:CUDA Event,精度高,不阻塞 host
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
heavy_gpu_work()
end.record()

end.synchronize()              # 等 end 这个事件被达到
print(f"{start.elapsed_time(end):.3f} ms")
1
2
3
4
# 方法 3:torch.profiler,看清楚每个 kernel 在哪个 stream 跑多久
with torch.profiler.profile() as prof:
    heavy_gpu_work()
print(prof.key_averages().table())

第二种方法是 production 测速首选——它在 GPU 自己的时钟上打时间戳,不受 host 异步行为干扰。

四、典型场景一:DDP 的计算-通信重叠

把前面的概念套到一个实际场景。DDP 反向时要做的事:算各层的梯度 → AllReduce 跨卡同步 → 优化器更新。

4.1 同步版的反向(基线)

如果直接写”反向算完→AllReduce→更新”:

1
2
3
4
5
6
时间 →
反向计算  :  [L_N][L_{N-1}][L_{N-2}]...[L_2][L_1]
NCCL 流   :                                    [AR L1][AR L2]...[AR LN]
更新      :                                                            [step]

总时间    =  T_backward + T_allreduce + T_step

通信和计算完全串行——通信时间一秒不少地加在 wall-clock 上。

4.2 DDP 实际怎么做(异步重叠)

DDP 用了我们前面讲的所有招数:

  1. Hook 触发:每个参数 register_hook,grad 一就绪就通知 DDP
  2. Bucket 凑齐就发:dist.all_reduce(..., async_op=True)——任务进 NCCL 的 stream,host 立即返回继续反向
  3. NCCL 走自己的 stream:NCCL 通信 kernel 跑在专门的 NCCL stream,不抢默认 stream 的 SM
  4. Bucket 反向逆序:最后一层在第一个 bucket,反向才刚开始就能发出第一波 AllReduce

时序图变成:

1
2
3
4
5
6
7
时间 →
默认 stream (反向) :  [L_N ][L_{N-1} ][L_{N-2}]...[L_2][L_1]
NCCL stream (通信) :       [AR_bucket_0]   [AR_bucket_1]...[AR_bucket_K]
                                                                     ↑ 尾部还有少量等待
host 在做啥        :  [bk_N: launch L_N + check bucket][bk_{N-1}: launch + check]...[wait_all]

总时间    ≈  max(T_backward, T_allreduce) + 很短的 tail

通信几乎完全被反向计算掩盖。如果反向比通信慢,通信就是”白送”的;如果通信比反向慢(模型小、跨节点带宽差),通信会从尾部露出来,这时候才是 DDP 的瓶颈。

4.3 一行关键代码

DDP 内部的核心异步逻辑就这两行(简化):

1
2
3
4
5
6
7
8
9
# bucket 凑齐 → 异步发射,不阻塞反向
bucket.handle = dist.all_reduce(bucket.flat_grad,
                                op=ReduceOp.SUM,
                                async_op=True)

# backward 末尾,一次性等所有 bucket 收尾
for bucket in buckets:
    bucket.handle.wait()
    bucket.flat_grad.div_(world_size)

整个 backward 过程中只有最后那几个 wait() 是真正的同步点,前面全是异步堆积。

五、典型场景二:ZeRO-3 / FSDP 的参数预取

ZeRO-3 把参数本身分片,每张卡只持有 1/N。前向时每用到一层都要先 AllGather 一下把这层参数凑齐。如果同步做:

1
2
3
4
时间 →
计算 :  [---等---][L1 fwd][---等---][L2 fwd][---等---][L3 fwd]
通信 :  [AG L1  ]         [AG L2  ]         [AG L3  ]
                ↑ GPU SM 干等

每次 AllGather 都让计算干等,相当于把通信完全暴露在 wall-clock 里,慢得离谱。

FSDP 的解法是 prefetch:计算第 i 层的同时,在另一个 stream 上异步发起第 i+1 层的 AllGather:

1
2
3
4
5
6
7
8
时间 →
默认 stream (计算)  :  [L1 fwd][L2 fwd][L3 fwd][L4 fwd]
NCCL stream (通信)  :  [AG L1]
                            [AG L2]    ← 在 L1 计算时已发起
                                  [AG L3]
                                        [AG L4]

整体时间 ≈ max(计算, 通信),通信被掩盖

实现关键:用 stream + event 把”L_{i+1} 的 AllGather 必须在 L_i 计算开始之后”这条依赖建好,然后让 NCCL stream 自己往前跑。FSDP 的 forward_prefetch=True 和 backward_prefetch=BACKWARD_PRE 就是开/关这个机制的开关。

这是异步重叠在 ZeRO-3 里的应用——和 DDP 的思路一样(把通信丢到独立 stream + 提前发射),只是触发时机从”反向梯度就绪”换成”前向到达某层之前”。

六、典型场景三:Pipeline 并行的 1F1B

异步的另一种形态——不是计算和通信重叠,而是多个 micro-batch 在不同 stage 上同时流动。

6.1 朴素调度的”气泡”

设 4 个 stage、4 个 micro-batch。如果一个 micro-batch 跑完前向再跑反向才换下一个:

1
2
3
4
5
6
7
时间 →
Stage 0:  [F1][F2][F3][F4]                              [B4][B3][B2][B1]
Stage 1:      [F1][F2][F3][F4]                      [B4][B3][B2][B1]
Stage 2:          [F1][F2][F3][F4]              [B4][B3][B2][B1]
Stage 3:              [F1][F2][F3][F4][B4][B3][B2][B1]
          ←warmup→                              ←cooldown→
            ↑ 大量 stage 干等              ↑ 大量 stage 干等

气泡时间正比于 stage 数,stage 越多浪费越大。

6.2 1F1B(One Forward One Backward)

Megatron-LM 等用的调度:每个 stage 在 warmup 后,轮流交替跑一次前向和一次反向,把流水线塞满:

1
2
3
4
5
6
时间 →
Stage 0:  [F1][F2][F3][F4][B1][F5][B2][F6][B3][F7][B4][F8][B5][B6][B7][B8]
Stage 1:      [F1][F2][F3][B1][F4][B2][F5][B3][F6][B4][F7][B5][F8][B6][B7][B8]
Stage 2:          [F1][F2][B1][F3][B2][F4][B3][F5][B4][F6][B5][F7][B6][F8][B7][B8]
Stage 3:              [F1][B1][F2][B2][F3][B3][F4][B4][F5][B5][F6][B6][F7][B7][F8][B8]
                       ↑ warmup 后立刻交替 F/B,各 stage 几乎不空闲

气泡只剩下 warmup 和 cooldown 这两段,占比 (stage_count - 1) / num_microbatches,micro-batch 越多气泡越被稀释。

底层用的还是同样的异步思想:stage 之间用通信传 activation/gradient,通信走独立 stream,计算和通信重叠。1F1B 调度只是规定了”每个 stage 该按什么顺序往队列塞 F 和 B kernel”,硬件层面的并发还是 stream 模型在管。

七、易踩的坑

异步带来效率,也带来一堆”看起来不科学”的现象。常见几条:

7.1 用 time.time() 测出来的 GPU 耗时是错的

time.time() 测的是 host 走过那行代码的时间,kernel launch 本身只要几 μs。除非 host 撞到强制同步点,否则你测出来的是”提交时间”,不是”GPU 干完时间”。永远用 torch.cuda.synchronize() 包裹或 CUDA Event,前面 §3.4 讲过。

7.2 错误暴露得很晚

GPU kernel 执行错(比如越界、NaN 触发某些 assert),host 不会立刻收到——它已经跑过那行代码了。错误要等到下一次同步点(下一个 .item()、synchronize()、甚至下一个 iteration 才暴露。看到报错栈指向某行 PyTorch 代码,真凶可能在很多行之前。

调试这种问题的招:加环境变量 CUDA_LAUNCH_BLOCKING=1,强制每次 kernel launch 都等 GPU 跑完,把异步关掉。错误就会指向真正出错的那行,代价是训练慢几倍——只用于调试。

7.3 OOM 也会延迟

显存不够爆 OOM 同样不一定在你期待的那行报。如果某次显存分配在另一个 stream 上 lazy 触发,栈可能完全是误导的。同样可以用 CUDA_LAUNCH_BLOCKING=1 看真实位置,或者用 torch.cuda.memory._record_memory_history() 抓显存事件回放。

7.4 多 stream 的依赖必须显式建

PyTorch 默认 stream 上的操作彼此自动建依赖(因为是同一个队列)。但你自己开 stream 后,跨 stream 的依赖必须自己用 event 建。忘了建依赖就等于在程序里悄悄写了个 race condition,出来的数据可能是上一个 batch 的、可能是部分写好的。

1
2
3
4
5
6
7
8
9
10
11
12
13
# 错的:s2 用 a,但 s1 还没算完
with torch.cuda.stream(s1):
    a = compute()
with torch.cuda.stream(s2):
    use(a)            # ⚠️ undefined behavior

# 对的:用 event 建依赖
with torch.cuda.stream(s1):
    a = compute()
ev = torch.cuda.Event(); ev.record(s1)
with torch.cuda.stream(s2):
    ev.wait()
    use(a)

PyTorch 的 tensor.record_stream(stream) 还能告诉显存分配器”这块 tensor 还在被 stream 用着,别回收”,防止显存被提前释放。这些细节是手写多 stream 代码时绕不开的。

八、一句话总结

GPU 是个异步加速器——你写的每行 .cuda() 操作几乎都不阻塞 host,只是把任务提交到某个 stream 的队列。Stream 让多种硬件资源(SM / DMA / NVLink)同时忙起来,event 负责跨 stream 建依赖。DDP 的计算-通信重叠、ZeRO-3 的参数预取、Pipeline 的 1F1B,都是同一套”独立 stream + 提前发射 + 显式同步”思想的具体应用。会用异步,大模型训练的吞吐才能逼近硬件理论上限;不理解异步,测出来的时间和报错的位置都会让人摸不着头脑。

Author: YANG
Link: https://SGDBS.github.io/AI_Infra/1.DataParallel/AsyncCompute/
Reprint policy: All articles in this blog are used except for special statements CC BY 4.0 reprint polocy. If reproduced, please indicate source YANG !
AI AI Infra
赏
  TOC