硬件与张量 入门
这一轨前面的 primer 描述了网络在做什么。这一篇描述它跑在什么上面。 4 个短主题:GPU vs CPU,为什么每个现代模型都住在 GPU 上;VRAM(显存),现在「你能跑多大模型」最硬的那个约束;张量(tensor),深度学习的基本数据单位 —— 多维数组; 以及批量、填充、掩码,把现实世界里参差不齐的文本 打包成 GPU 友好的矩形的实战技术。这一篇实战里跳不掉 —— 你以后看到的每一条 「CUDA out of memory」错误信息,都是用这套词汇写出来的。
GPU vs CPU
CPU 聪明但少。GPU 笨但多。深度学习属于「第二种活」。
现代 CPU 就是「计算机」的柏拉图理想形态。少量(几个到几十个)又快又灵活的核 —— 笔记本电脑 8 个、工作站 64 个 —— 每一个都能高速做分支判断、复杂指令、 各种琐碎记账。优化的是「一个线程跑一长串多变的任务」。 操作系统、数据库、Web 服务器 —— 这是 CPU 的地盘。
现代 GPU 正好反过来。几千个又小又笨的核 —— NVIDIA H100 有 16,896 个 CUDA 核。每个核都不擅长分支、不擅长一般逻辑,只有在「跟所有邻居做同样的 简单操作」时才快。执行模型叫 SIMD(单指令多数据),或者更细粒度的表亲 SIMT。 整个架构是为「一个操作并行作用到几千份数据上」而生的。
而这第二种模式,恰好就是神经网络的全部组成:
- 矩阵乘法。Transformer 最底层那个操作。两个 4096×4096 矩阵相乘 是 670 亿次乘法 —— 每一次都可以并行跑。
- 逐元素的激活函数。ReLU、GELU、softmax —— 同一个函数应用到 几百万个值上,彼此独立。
- 归一化层。LayerNorm 和 RMSNorm 沿一个轴 reduce 一下、 再给每个元素 rescale —— 也是「尴尬级别的并行」。
想要个数字感的话:4096 × 4096 矩阵乘,快的 CPU 大约 50 ms;一张 H100 GPU 不到 1 ms。 这个差距随规模超线性增长,因为 4k × 4k 上 GPU 几乎还没热起来,而 CPU 已经顶满了。 训一个前沿 LLM,通常需要几千张 GPU 并行跑几个月。换成 CPU 来训, 同样的工作量得跑几百年。所以每个现代模型都住在 GPU 上。
GPU不擅长什么,简短版:
- 分支多的代码。如果不同「核」要走不同分支,它们就轮流来 —— 快的等慢的。深度学习的对策是写「没分支」的代码:用 padding 和 mask, 所有地方都做同一个操作。
- 小工作量。启动 GPU 本身有开销。一个 32 × 32 的小矩阵乘, CPU 反而比 GPU 快,因为 GPU 启动成本占主导。真正的深度学习总是把工作打成大包、 摊薄这个开销。
- CPU ↔ GPU 之间的内存传输。数据从系统 RAM 搬到 GPU 显存(§2) 要走 PCIe 总线,比 GPU 内部带宽慢得多。数据要尽量留在 GPU 上。
顺便:TPU 是 Google 的专用张量加速器 —— 比 GPU 还更专门、 专为「matmul + 激活」模式优化。Google 内部大部分模型(Gemini 等)都跑在上面。 功能上是同一个思路的更激进版本:大量笨的、并行的算术,加上受控的内存访问模式。
在 Transformer 里:每一层都是矩阵乘(Q·K、attention·V、FFN 权重) 加上少量逐元素操作(softmax、GELU、残差求和、LayerNorm)。整个前向传播就是 「给 GPU 喂一矩形一矩形的活」。设计架构的人是在为这个硬件现实而设计; 训练它的人,实际上是在为「能拿到多少 GPU 时间」竞争。
VRAM —— 最硬的瓶颈
你看到的每一条「CUDA out of memory」,故事都写在这里。
GPU 有自己专属的内存,跟 CPU 用的系统 RAM 是分开的。NVIDIA 叫它 VRAM, 新一点的芯片叫 HBM。H100 是 80 GB;A100 是 40 或 80;4090 是 24;4070 是 12。 这个数字就是「你能在这张卡上训或者跑多大模型」的硬约束。 模型和它的临时空间不够装,你就训不了。没得商量。
对一个 70 亿参数(7B)、fp16(每个数 2 字节)的模型,算一下:
推理(只装模型):
权重 7 × 10⁹ × 2 字节 = 14 GB
─────────────────
24 GB 的消费级显卡装得下
训练(模型 + 训练的一切配套):
权重 14 GB
梯度 14 GB ← 每个参数 1 份
Adam momentum 14 GB ← 每个参数 1 份
Adam variance 14 GB ← 每个参数 1 份
激活 ~30 GB ← 取决于 batch 和 seq_len
─────────
≈ 86 GB → 80 GB H100 装不下这里有两个惊喜。第一,对同一个模型,训练大约要用 5–10 倍于推理的 VRAM —— 因为优化器给每个参数都留了好几份,而且前向的激活值要存下来给反向用。 第二,推理时能在一张 24 GB 消费卡上轻松跑的 7B 模型,训练时连一张 80 GB H100 都装不下。 70B 模型光权重就是 140 GB —— 保证是「多卡领域」。
常见的减负手段,大致按人们伸手抓的顺序:
- 量化(Quantization)。用更低精度存权重 —— int8(字节数减半)、int4(四分之一),或者更花哨的 3 比特 / 2 比特方案。 fp16 下 14 GB 的 7B 模型,int4 只剩 3.5 GB。推理时很常见; 训练时比较麻烦,因为梯度想要更高精度。
- 混合精度训练。激活和梯度用 fp16 或 bf16, 优化器的「主拷贝」保留为 fp32 以保证稳定。激活内存差不多减半。
- 梯度 / 激活检查点(Gradient checkpointing)。不把前向的所有激活都存下来 —— 每几层存一次,反向时重新算中间那些。 用算力换内存。
- 模型分片 / 张量并行。把模型本身切开、放到多张 GPU 上。 ZeRO、FSDP、DeepSpeed、Megatron —— 每个分布式训练框架都是这套思路的一种变体。
- Flash attention。重新实现 attention,让 n × n 的注意力矩阵 从来不实化在 VRAM 里。在长上下文训练里能省 20–80 GB;现在每个现代 Transformer 实现 都内置了。
带宽是「内存」故事里的另一半,直到撞上之前没人提。H100 的 VRAM 带宽是 3 TB/s —— 巨大,但仍有上限。很多神经网络操作是内存瓶颈型的:GPU 等数据从 VRAM 过来的时间, 比真正算的时间还多。高效 kernel(FlashAttention、融合 MLP 等) 的一部分工作就是「尽量少往返 VRAM」。GPU 内部还有内存层级 —— 寄存器、shared memory / L1 缓存(每个 SM 256 KB)、L2 缓存(整卡共 50 MB)—— 比 VRAM 快得多,但小得多。性能优化大多是「让数据多待在更快的内存里」。
在 Transformer 里:每一个架构选择都带有 VRAM 决策的成分。KV 缓存(自回归解码时保留下来的过去 token 的 keys 和 values) 的规模是 batch × seq_len × n_layers × d_kv, 长上下文模型里它经常比权重本身还大。压 KV 缓存(Grouped-Query Attention、 Multi-Query Attention、滑动窗口 attention)是现代 LLM 能做到 128k 或 1M token 上下文的一半原因。
张量 —— 多维数组
深度学习的基本数据单位。网络里流的每个值,都是一个张量。
剥掉框架、剥掉模型、剥掉所有抽象 —— 剩下的就是张量。 张量就是数字的多维数组。整个现代 ML 栈(PyTorch、JAX、TensorFlow、NumPy) 本质上就是「高效创建、reshape、组合张量」的一堆函数。 理解了张量,你就理解了一个 LLM 代码里 90% 的事 —— 一行一行看下去都是张量在做事。
维度阶梯:
- 0-D(标量) —— 一个数。
shape = ()。 一次前向最后吐出来的 loss;一个学习率;一个概率。 - 1-D(向量) —— 一排数。
shape = (n,)。 一个词的 embedding(比如 768 维向量);一层的偏置。 - 2-D(矩阵) —— 行 × 列。
shape = (m, n)。 一个权重矩阵;一句话的 embeddings(序列 × d_model);一个协方差矩阵。 - 3-D(张量) —— 一摞矩阵。
shape = (a, b, c)。 Transformer 的家常形状:(batch, seq, d_model)。 - 更高维(n-D) —— 再加轴。多头 attention 在 4-D 里跑 (batch, heads, seq, head_dim);视频数据是 5-D(batch, time, channel, height, width)。
对每一个张量,你做事之前必须知道的三件事:
- Shape(形状)。一个整数元组 ——
(32, 512, 768)意思是 32 个样本,每个 512 个 token,每个 token 一个 768 维 hidden state。 神经网络里几乎所有 bug 都是形状对不上。 - Dtype(数据类型)。精度:
float32(4 字节)、float16/bfloat16(2 字节)、int8(1 字节)、bool。dtype 的选择一半是 VRAM(§2)决策、一半是数值稳定性决策。 - Device(设备)。字节住在哪 ——
cpu、cuda:0、cuda:1等。两个不同 device 上的张量做运算会报错; 先得.to(device)。
95% 时间会用到的操作:
- 逐元素:
+、*、tanh、relu。每个格子单独算;输出形状 = 输入形状。 - 归约:
sum、mean、max沿某个轴。(32, 512, 768).mean(dim=-1)→(32, 512)—— hidden 维被压扁了。 - Reshape:
view、reshape、permute、transpose。重新排列轴,不改数据。 Transformer 里 batch / seq / head 轴换来换去靠的全是这些。 - 矩阵乘法:PyTorch / NumPy 里写作
@。(32, 512, 768) @ (768, 768) → (32, 512, 768)。最大的那个操作; GPU 就是为这个生的。 - 索引 / 切片:
x[:, 0, :]取每个样本的第一个 token。x[mask]在mask为真的位置上取出来。
广播(broadcasting)是那种「初学者完全蒙圈、熟练后看不见」的操作。 你写 (32, 512, 768) + (768,) 时,小张量自动沿缺失的维度 「拉伸」一下,于是同一个 (768,) 的 bias 向量被加到每个样本的每个位置上。 没有真的复制内存;广播是虚拟的。神经网络代码里几乎每一行都靠广播写得很紧凑。
在 Transformer 里:输入 token 被 embed 成 shape 为(batch, seq, d_model) 的张量。这个形状在网络里一直保持 —— 每一层的输出又是同样形状的张量,只是数变了。 attention 内部,张量短暂 reshape 成 (batch, heads, seq, d_head) 来 并行算多头点积、然后再 reshape 回来。最后一层把它投影到(batch, seq, vocab_size) —— 每个位置在词表上一个 logit 向量。 每一步都是张量操作。
批量、填充与掩码
变长文本到底是怎么打包成 GPU 友好的矩形的。
真实的文本是一堆参差不齐的序列。真实的 GPU 想要矩形张量。 三件实战工具把两者桥起来 —— 它们出现在每一个 Transformer 训练循环的每一行里。 文本 primer §1 里点了个头;本节把机制钉死。
批量(Batch)。
你不会一次只喂一个句子。你一次喂 N 个句子,堆成 shape 为(N, max_len) 的张量,让 GPU 把这 N 个并行处理。N —— 就是 batch size —— 是深度学习里最重要的超参之一。 batch 越大,显存吃得越多,但 GPU 更忙、梯度平均得更狠、(在合理范围内)训练更快。 典型值:几张卡上微调 70B 模型时是 32;预训 GPT-4 量级模型时,batch 里有几百万 token。 具体数字取决于你的 VRAM 能装下多少。
填充(Padding)。
同一个 batch 里的句子长度不一样。要堆成 shape 为 (N, max_len)的矩形张量,得让它们一样长。你填充短的那些,在末尾追加特殊的[PAD] token —— 词表里约定俗成的 id 0。batch 里最长的那句话决定max_len;其它句子都追加 PAD 到这个长度。
填充上有两个实战旋钮。第一,max_len 的选择: 填到整个数据集的最长长度浪费算力;每个 mini-batch 内填到本批的最长(叫「动态填充」) 浪费就少得多。第二,填哪一边:大多数现代 Transformer 训练时右填充 (right-padding),自回归生成时左填充(left-padding)。原因细微但很要紧 —— 搞错了模型会悄悄学到一堆没意义的东西。
掩码(Mask)。
填充修好了形状,但带来一个新问题:attention 默认会愉快地计算「真 token 和 PAD」 之间的 attention 权重。模型会学到通过 padding 传信息 —— 而这毫无意义。 修法是掩码 —— 一个平行的 0/1 张量(经常是 bool), 告诉模型哪些位置是真的:
A 句: ["hi", ".", PAD, PAD, PAD, PAD, PAD] A mask: [ 1, 1, 0, 0, 0, 0, 0] B 句: ["the", "dog", "runs", "fast", ".", PAD, PAD] B mask: [ 1, 1, 1, 1, 1, 0, 0]
在 attention 块内部,掩码在 softmax之前应用:PAD 位置的 attention 分数被 设成 −∞,这样 softmax 后这些位置的权重精确为 0。 PAD token 对输出的贡献是 0;模型表现得跟「只在没填充的句子上跑」一模一样 (除了浪费的那点算力)。
Transformer 里常见的两种掩码:
- Padding mask —— 刚刚说的那种。标记哪些位置是真、哪些是 PAD。 基本每个 Transformer 都有。
- 因果掩码 / Causal(autoregressive)mask —— 用在 decoder-only 模型 比如 GPT 里。位置 t 只能关注位置 ≤ t。实现成一个三角形的
−∞矩阵,叠到 attention 分数上。和 padding mask 用逻辑「与」合成。
在 Transformer 里:每一次 attention 调用都带一个 mask。 训练时一个典型 batch 两者都有:为了自回归结构的因果掩码,加上为了变长的 padding mask。 输出是对的、GPU 看到的是漂亮的矩形张量、loss 只对真 token 计数 (loss 函数本身也要掩 —— 别因为模型在 PAD 位置预测的东西去惩罚它)。 这些不显眼的「水电管线」,正是「批量并行训练」能成立的全部基础。