第3章 CUDA与PyTorch

CUDA与PyTorch：GPU加速背后的秘密，CUDA Stream与异步执行

说实话，很多人用PyTorch训练模型，觉得GPU加速就是“把数据扔到cuda上就完事了”。嗯，这话对了一半。真正让GPU火力全开的，是理解它背后的执行模型——CUDA Stream和异步执行。今天我就带你看看，这层“黑盒”里到底发生了什么。

GPU不是“瞬间完成”的魔法

先纠正一个常见误解。很多人以为调用cuda()之后，计算就立刻在GPU上跑完了。其实不是。CPU和GPU是异步工作的。你调用一个CUDA kernel，CPU只是把这个任务“丢”到GPU的命令队列里，然后立刻返回，继续干别的事。

我刚开始做模型部署时，就踩过这个坑。当时我写了一段代码，在GPU上做完推理后，立刻去读结果，结果读到的全是垃圾数据。为什么？因为GPU还没算完呢！

核心概念：CUDA Stream是一个按顺序执行的操作队列。同一个Stream里的操作，严格按照你提交的顺序执行。不同Stream之间，则可以并行执行。

默认Stream vs. 自定义Stream

PyTorch默认把所有操作都放在一个叫“默认Stream”的队列里。这就像只有一个收银台的超市——所有顾客（操作）都得排队。虽然GPU内部有成千上万个核心，但命令的提交和同步，还是得按顺序来。

那自定义Stream有什么用？说白了，就是多开几个“收银台”。

import torch

# 创建两个自定义Stream
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

# 在Stream s1中执行操作A
with torch.cuda.stream(s1):
    # 假设这是一个耗时的大矩阵乘法
    A = torch.randn(5000, 5000).cuda()
    B = torch.randn(5000, 5000).cuda()
    C = A @ B

# 在Stream s2中执行操作B
with torch.cuda.stream(s2):
    # 同时做另一个独立计算
    D = torch.randn(5000, 5000).cuda()
    E = torch.randn(5000, 5000).cuda()
    F = D @ E

# 等待所有Stream完成
torch.cuda.synchronize()

你看，操作A和操作B如果放在默认Stream里，就得一个算完再算另一个。但放在不同Stream里，只要GPU资源够，它们就能同时跑。我曾在一次推理优化项目中，用这个技巧把两路独立的视频流处理时间，从40ms降到了22ms——几乎翻倍。

异步执行：别让CPU闲着

CUDA的异步执行，说白了就是“CPU和GPU各干各的”。CPU提交完任务后，不用傻等GPU算完，可以先去准备下一批数据、做预处理、或者更新模型参数。

PyTorch里很多操作默认就是异步的。比如：

# 异步拷贝：CPU到GPU
x = torch.randn(10000, 10000)
x_gpu = x.cuda(non_blocking=True)  # 关键参数

# 这里CPU不会等拷贝完成，可以立刻做别的事
y = torch.randn(10000, 10000)
y_gpu = y.cuda(non_blocking=True)

# 真正需要结果时，再同步
torch.cuda.synchronize()

这个non_blocking=True，我建议你养成习惯。尤其是在数据加载环节。我曾经优化过一个数据加载器，把所有的cuda()调用都改成异步，配合多进程预处理，训练吞吐量直接提升了30%。

我的习惯：在DataLoader的collate_fn里，就把tensor异步拷贝到GPU。这样CPU在准备下一批数据时，GPU已经在算当前批次了。流水线作业，效率最高。

Stream同步：小心“野指针”

多Stream虽然好，但有个大坑——同步问题。不同Stream之间的操作，默认是互相不知道对方存在的。如果你在Stream A里写了一个tensor，然后在Stream B里去读它，那结果是不确定的。

为什么会这样？因为GPU执行Stream A和Stream B的顺序，是由硬件调度器决定的。可能A还没写完，B就开始读了。

避坑指南：我曾经在写一个多Stream的流水线时，忘了加同步事件，结果模型训练出来的loss忽高忽低，查了两天才发现是Stream之间数据竞争。从那以后，我只要跨Stream访问数据，必加torch.cuda.Event做同步。

# 正确的跨Stream同步方式
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

# 创建一个事件对象
event = torch.cuda.Event()

with torch.cuda.stream(s1):
    # 在Stream s1中计算结果
    result = some_computation()
    # 记录事件
    event.record(s1)

with torch.cuda.stream(s2):
    # 等待s1中的事件完成
    s2.wait_event(event)
    # 现在可以安全地使用result了
    use_result(result)

CUDA Graph：把“排队”变成“一键执行”

说到Stream，就不得不提CUDA Graph。这是NVIDIA在CUDA 10里引入的一个特性。它的思路很巧妙：把一系列CUDA操作“录制”成一个图，然后可以反复“回放”。

你想想看，每次训练迭代，我们做的操作其实差不多：前向、反向、优化器更新。每次都要CPU一个个提交命令，这中间有大量的CPU-GPU通信开销。CUDA Graph把这些操作“固化”成一张图，一次提交，GPU自己按图执行。

# PyTorch中使用CUDA Graph的简化示例
# 注意：需要PyTorch 1.10+ 和 CUDA 10+

# 先录制
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    # 这里放你希望“固化”的操作
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

# 之后每次迭代，只需要回放
# 但要注意：输入输出的内存地址不能变
for _ in range(1000):
    # 更新输入数据（原地修改）
    input.copy_(new_data)
    # 回放整个计算图
    graph.replay()
    # 读取结果
    result = output.clone()

我测试过，对于小batch size的模型，CUDA Graph可以减少30%-50%的CPU开销。尤其是那种“模型不大，但迭代次数很多”的场景，效果特别明显。

一张图总结核心逻辑

下面这张图，是我自己梳理的CUDA Stream与异步执行的核心流程。你看一眼就能明白：

从这张图你能看到：CPU提交任务到不同Stream，Stream里的操作被调度到GPU的不同计算单元上并行执行。最后通过同步事件，CPU再取回结果。整个过程，CPU和GPU都在高效运转，谁也不等谁。

实际项目中的Stream使用策略

说了这么多，到底什么时候该用多Stream？我总结了几条经验：

数据加载与计算重叠：一个Stream负责把数据从CPU拷贝到GPU，另一个Stream负责计算。这是最常见的用法。
多模型并行推理：如果同时跑多个不同的模型，每个模型放在自己的Stream里，互不干扰。
流水线并行：模型的不同层放在不同Stream里，前一层算完立刻通知后一层开始。

重要提醒：不是Stream越多越好。GPU的硬件资源是有限的。Stream太多，调度开销反而会增大。我一般控制在2-4个Stream，效果最好。

嗯，关于CUDA Stream和异步执行，今天就聊这么多。这些概念看起来有点抽象，但只要你动手写几行代码，跑一下profiling，就能直观感受到它们的威力。记住一句话：让GPU永远有活干，让CPU永远不等GPU——这就是性能优化的终极心法。

📚 课程目录

CUDA与PyTorch：GPU加速背后的秘密，CUDA Stream与异步执行

GPU不是“瞬间完成”的魔法

默认Stream vs. 自定义Stream

异步执行：别让CPU闲着

Stream同步：小心“野指针”

CUDA Graph：把“排队”变成“一键执行”

一张图总结核心逻辑

实际项目中的Stream使用策略