GPU 如何训练模型：从算子到 Kernel

GPU 不会收到一句“把模型训练好”

Python 训练代码决定做哪些运算，框架把运算派发给设备实现，GPU 执行一批批 kernel。loss.backward() 也不是特殊硬件指令：autograd 沿计算依赖找到反向算子，继续启动矩阵乘、归约、逐元素运算和必要的通信。不熟悉这条梯度路径，可以先走梯度与反向传播旁路。

课程分析：理解性能的第一张图是数据什么时候在哪里，而非显卡型号排行榜。一个 step 的墙钟时间可能被数据加载、CPU 派发、HBM 访问、计算、跨卡等待或 checkpoint 拖慢；仅看到 GPU 利用率接近 100%，不能判断瓶颈。

CPU：准备 batch、组织计算图、启动设备任务
GPU 显存：参数、输入、激活、梯度与优化器状态
Forward kernels：矩阵乘、attention、归一化、损失
Autograd：按依赖启动 backward kernels
Gradient communication：多卡时同步或分片归约
Optimizer kernels：读取梯度与状态，写入新参数
下一个 batch：复用显存与执行计划

这是一条逻辑依赖链。真实系统允许数据传输、部分计算与通信重叠，不能将它误解为每一阶段必然串行。

算子、Kernel、SM 和 Tensor Core 各是哪一层

上述硬件层级与执行模型由 NVIDIA 官方性能文档说明；具体容量、支持精度和峰值需按目标硬件核验，本站不把某一代设备数字当成普遍规律。官方硬件背景

HBM 不是所有显存的统一名称：数据中心卡常见 HBM，消费卡也常用 GDDR。下文的 HBM 带宽模型描述设备外部显存的数据流，应用到具体卡时必须填入实际内存类型和有效带宽。

一个线性层，前向和反向到底在算什么

将 batch 与序列维合并，设 $X\in\mathbb R^{M\times K}$、$W\in\mathbb R^{K\times N}$，前向 $Y=XW$。来自上游的梯度 $G=\partial L/\partial Y$ 给出：

这三次主要计算都是 GEMM。按一次乘和一次加计两个 FLOP，单个前向约为 $2MKN$ FLOP。反向还需输入/参数梯度的矩阵乘，以及其他算子的反向；是否需要 $\partial L/\partial X$ 取决于上游是否训练。张量转置常能由 stride/库接口表达，并不意味着一定物理复制；真正的 contiguous 转换需要 profiler 确认。GEMM 计数依据

课程分析：BF16 参数只占每元素两字节，不代表训练只用两字节。梯度、FP32 累积或主副本、Adam 一二阶状态、保存激活及临时 buffer 共同占显存。MoE 的 active 参数适合估算部分计算量，驻留总参数才影响权重容量。先把这两本账分开。

Roofline：先判断“算不动”还是“搬不动”

设一次操作需要 $F$ FLOP、从 HBM 移动 $Q$ byte，GPU 对应精度峰值为 $P$ FLOP/s、有效内存带宽为 $B$ byte/s：

这只是理想下界，不含启动、同步、地址计算、低占用和访存不合并等开销。GEMM 的理想最少张量读写估算是 $Q=b(MK+KN+MN)$，其中 $b$ 是元素字节数；实际可能读写更多。强度与瓶颈方法

课程算例：BF16、$K=N=4096$ 时，$M=32$ 的理想算术强度约 31.5 FLOP/byte；把 $M$ 增至 2048，理想强度变为 1024 FLOP/byte。增加 token batch 可复用同一权重，提高矩阵乘效率，但也增加激活显存与端到端延迟。因此“小 batch 用不满卡”不必是 CUDA 配置出错。

对逐元素 y=x+b，每个元素少量计算却需要读取输入并写输出，通常更容易受访存影响。把一个 memory-bound 算子的数学表达减少一次加法，可能远不如避免一次中间张量写回有效。

融合与 FlashAttention：减少什么，没减少什么

课程分析：bias、activation、dropout、residual 若分别物化中间张量，会重复写 HBM 再读回。融合让中间结果尽可能留在片上，减少启动与流量；但过度融合可能增加寄存器压力、降低并发或阻碍其他并行。因此“kernel 数更少”不是最终验收指标，step 时间与数值一致性才是。

FlashAttention 的核心是分块计算精确 attention，通过在线 softmax 统计等手段避免把完整 $N\times N$ 注意力矩阵反复物化到 HBM；反向可以利用保存统计重算。它改善 IO 与存储，并没有把 dense attention 的数学计算复杂度自动从二次变成线性。论文依据

源码导读：固定 Dao-AILab/flash-attention 的 edb5c76ee329b18ed95d1f7ea9aa522a1331ab7d，读取 flash_attn/flash_attn_interface.py：flash_attn_func → FlashAttnFunc.apply → forward → _wrapped_flash_attn_forward，底层转入设备实现。前向上下文保存 Q、K、V、output、softmax LSE 与 RNG 状态，反向入口调用 _wrapped_flash_attn_backward。源码与 BSD-3-Clause 许可证

Q 的接口形状为 [batch,seqlen,nheads,headdim]，K/V 可以有更少 heads。由此检查 GQA 头数整除、causal mask 对齐、dropout 与梯度是否可用。该文件是一个具体接口路径，不能据此认定 PyTorch SDPA 在所有设备、dtype 与 mask 下都会选择同一 kernel。真实分派必须看 trace。

实机实验：先看测到了什么

项目不强制安装 PyTorch；已有 PyTorch 时运行下面的 CPU 观测：

python3 -m labs.torch_profile --device cpu --operation mlp --output /tmp/profile.json --trace /tmp/profile.trace.json

确认设备支持后可明确选择 CUDA：

python3 -m labs.torch_profile --device cuda --operation mlp --output /tmp/profile-cuda.json --trace /tmp/profile-cuda.trace.json

这是本站编写的微型算子实验。命令输出记录实际设备和 PyTorch 版本；没有 CUDA 时不能把 CPU 时间改名为 GPU 时间。matmul、softmax、mlp 用来对照计算密集、归约和多个算子串联。配置与环境问题见 labs/README.md；成本模型另运行 python3 -m labs.run --experiment systems。

本工作区实际验证：在 PyTorch 2.11.0+cu128、CPU 和 NVIDIA GeForce RTX 4060 Laptop GPU 上分别执行了 FP32、输入 [128,256]、权重 [256,1024] 与 [1024,256] 的微型 MLP。实际日志包含 aten::mm、GELU/softmax 及反向运算，梯度有限。这里只测前向和反向，没有 optimizer、NCCL 或完整语言模型；短时单机结果不能代替数据中心硬件 benchmark。

PyTorch profiler 支持 CPU/CUDA activities、shape、memory 和用户标记。接口来源 看表时先区分 self 时间与包含子调用的 total 时间：把父算子和子 kernel 全部求和会重复计算。CPU 启动时间也不等于 GPU 完成时间。CUDA 异步执行时用未同步的 Python 计时只测到派发；实验将计时与 profiler 开销分开，实际吞吐应看未开启 profiler 的稳定测量。

课程分析：先看时间线里的空洞，再看最长 kernel 和对应 shape。首次编译、autotune 与缓存预热应单列；内存测量区分 allocated 与 reserved。用 profiler 找假设，再关闭 profiler 验证优化，避免观测本身改变了被测对象。

MFU 与“GPU 利用率”为什么不一样

nvidia-smi 的 GPU utilization 表示采样时间内至少有 kernel 执行的时间比例，不等于 Tensor Core 达到峰值的比例。指标定义

课程分析：必须声明模型 FLOP 公式、精度峰值是否稀疏、是否计重计算、padding 和 attention。粗略 $6NT$ 适合某些 dense Transformer 估算，不能照搬到所有 MoE、超长上下文或混合注意力。持续执行访存 kernel 可以让利用率很高而 MFU 很低；这未必意味着程序有错，可能是硬件与工作负载的匹配问题。

调试任务：每次只改变一个瓶颈

工程任务：保留相同输入与损失，对三个 shape 测 matmul/softmax/MLP，提交未开启 profiler 的时间、实际设备、dtype、峰值显存和时间线解释。然后从 Triton 官方向量加法教程定位 program_id、load、store 的边界 mask；解释尾块为何必须 mask。可选源码练习 不需要为了这门课先重写高性能 GEMM；只有确认瓶颈、正确性与目标 shape，定制 kernel 才有明确目标。