1. 显存瓶颈的根源

1.1 一个让我印象深刻的教训

几年前我接手过一个图像处理项目。模型结构不复杂,计算量也不大,按理说跑起来应该很轻松。结果呢?训练一个batch要等好几秒。

我当时第一反应是「模型太深了,得剪枝」。折腾了两天,效果微乎其微。后来一查性能分析工具——好家伙,GPU计算单元的空闲率高达70%。

说白了,GPU一直在等数据。等谁?等PCIe把数据从CPU内存搬过来。

这个坑,我至今记忆犹新。从那以后,我养成了一个习惯:优化显存之前,先看数据搬运的路径。

1.2 数据搬运的「高速公路」有多窄?

你想想看,GPU的计算能力每年翻倍,但PCIe的带宽呢?

拿PCIe 4.0 x16来说,理论带宽32 GB/s。听起来不小对吧?但对比一下:

硬件 带宽 备注
PCIe 4.0 x16 32 GB/s CPU ↔ GPU 主通道
GPU显存带宽(HBM2e) ~1.6 TB/s GPU内部访问
L2缓存带宽 ~4 TB/s 片上缓存
寄存器带宽 ~20 TB/s 最靠近计算单元

看到了吗?PCIe带宽比显存带宽慢了整整50倍。比寄存器带宽慢了600倍。

这就是问题的根源。计算单元跑得再快,数据过不来,一切都是白搭。

1.3 内存层级:为什么「远近亲疏」这么重要?

我经常跟团队说一句话:GPU里的数据,也是有「阶级」的。

离计算单元越近,速度越快,但容量越小。离得越远,容量越大,但速度越慢。这个层级结构,你必须要刻在脑子里:

  • 寄存器:最快,KB级别,每个线程私有
  • 共享内存:次快,几十KB,同一个block的线程共享
  • L1/L2缓存:中等速度,MB级别,自动管理
  • 显存(HBM/GDDR):较慢,GB级别,程序员手动管理
  • CPU内存:最慢,通过PCIe访问,延迟以毫秒计

为什么会这样?说白了,物理限制摆在那里。芯片面积有限,不可能把几百GB的显存都做成寄存器那么快。成本也扛不住。

核心原则:尽量让数据待在「近处」,减少跨层级的搬运。一次PCIe传输的成本,够你在显存内部做几百次计算了。

1.4 一张图看懂数据流动

下面这张图,是我每次做性能优化时都会先画一遍的。它帮你理清数据从磁盘到GPU计算单元的全路径:

GPU数据搬运层级结构 磁盘 / 网络 ~1 GB/s CPU 内存 ~50 GB/s PCIe 4.0 x16 32 GB/s (理论) 显存 (HBM) ~1.6 TB/s L2 缓存 ~4 TB/s L1 / 共享内存 ~8 TB/s 寄存器 ~20 TB/s CUDA 核心 计算单元 越往上,带宽越大,容量越小 PCIe 是整个链路中最窄的瓶颈

这张图我建议你保存下来。每次做优化时,问自己一个问题:我的数据现在在哪一层?它要去哪一层?

1.5 避坑指南:我曾经踩过的三个坑

坑一:频繁的CPU-GPU数据交换

我曾经写过一个训练循环,每个step都把loss拷回CPU打印。结果呢?每个step都在等PCIe。后来改成每100个step拷一次,速度直接快了3倍。

教训:尽量减少host↔device的数据传输次数。能批量传输就别一次传一个。

坑二:小数据量也走PCIe

有个同事写了个kernel,每次只处理4个float。但每次都要从CPU拷过去。我说你想想看,PCIe一次传输的固定开销就有几微秒,你传4个float才16字节,这性价比太低了。

后来改成攒够一批再传,问题解决。

坑三:忽略内存对齐

这个坑比较隐蔽。我记得有一次,数据传输总是比预期慢30%。查了半天,发现是内存没对齐,导致PCIe无法发挥最大带宽。对齐之后,速度立刻上来了。

小技巧:使用 cudaMallocHost 分配页锁定内存,可以显著提升PCIe传输效率。我个人习惯在需要频繁传输数据时,优先用这个。

1.6 总结:记住这三句话

  • 计算不贵,搬运才贵——PCIe是GPU性能的第一道坎
  • 数据要「就近」处理——尽量让数据待在显存里,别来回搬
  • 批量传输,减少次数——一次搬1GB比搬1000次1MB快得多

嗯,这一章就到这里。下一章我们会深入聊聊,如何用代码来「驯服」这些数据搬运的野兽。