原文链接0. 引言此章并非Pytorch入门这一章表面上在讲pytorch的张量、矩阵乘法、反向传播、参数初始化、数据加载、优化器、训练循环但真正的重点应该是从资源的角度重新理一遍这些本来就会写的代码。我对pytorch的训练原本也不算陌生至少会写个能跑的baseline但显然我们的时间金钱不是无限的学了这章我需要开始关注每一步到底在算什么、吃了多少显存、要多少FLOPs、为什么要特定的精度格式等等对模型训练需要的时间和内存有一个大概的概念。此篇笔记将不再赘述我原本熟悉的东西。1. 张量显然张量并不是一个数组或者一个单纯的容器在pytorch中输入、参数、激活值、梯度、优化器状态本质上都可以看成张量是整个训练过程最基础的表达形式。1.1 view、transpose和contiguous张量看上去是多维的但层本质上仍是一块线性内存。view、transpose这样的操作不会重新复制一遍数据而只是改变了张量对底层内存的解释方式数据在内存上物理的排列没有改变这叫做zero-copy十分高效。但transpose也会带来non-contiguous的问题即便数据在物理上的排列没有改变但数据的逻辑排列还是变了即元数据形状欸或步长变了逻辑索引顺序与内存存储顺序不再一致一个转置之后的张量不能再直接view必须先.contiguous()而.contiguous()会创建新的连续内存有额外的内存和时间开销。1.2 einops和jaxtyping用原生的.view()和.transpose()操作维度时需要时刻记住张量的维度顺序很容易搞晕一旦模型维度变复杂代码就很容易变得难读。课程推荐用jaxtyping给维度加标签用einops做einsum、reduce、rearrange这类操作尽管增加了少量语法开销但其清晰的维度命名显著降低了调试难度更多用法可见 Einops tutorial。2. 内存2.1 内存不只属于参数除了参数本身还有梯度、激活值和优化器状态在资源核算时也要按这几类来算。不能简单地用 总显存 / 参数字节数 去简单估算模型是否可跑尤其是用了 AdamW这类优化器优化器状态本身会占掉很大一部分显存。2.2 FP32、FP16和BF16FP32 训练稳定但显存和计算代价高。FP16 省显存、速度快但动态范围太小训练过程中容易溢出造成致命后果。BF16 大小不变的情况下牺牲一点精度保留了更大的动态范围训练里更稳现在的主流选择。FP8 极致压缩显存精度极低训练极不稳定主要用于推理的量化。深度学习训练里损失小数点后10位的精度并无所谓但一旦溢出就是致命的所以BF16截断了FP32的尾数保留了指数位它能表示的数值范围与FP32一样大极大地提升了训练稳定性优于FP16。3. 矩阵乘法、FLOPs和MFU3.1 矩阵乘法是主要计算开销来源深度学习模型里大部分的计算都来自矩阵乘法无论是线性层还是Transformer底层核心都离不开它。如果输入x的形状是(B, D)权重w的形状是(D, K)一次矩阵乘法可以写成C i k ∑ D A i j B j k C_{ik} \sum^{D} A_{ij}B_{jk}Cik​∑D​Aij​Bjk​乘法D DD次加法D − 1 D-1D−1次则单次元素的浮点运算数近似为2 × D 2 \times D2×D次总输出元素个数为B × K B\times KB×K那么y x w的FLOPs近似为2 × B × D × K 2 \times B \times D \times K2×B×D×K这个公式虽然简单但它是后面训练成本估算的基础。3.2 FLOPs、FLOP/s 和 MFU几个概念FLOPs总计算量FLOP/s每秒浮点运算次数MFU实际吞吐相对理论峰值的利用率MFU实测FLOPS硬件理论峰值 / FLOPSMFU 0.5被认为是相当不错的性能但这个公式忽略了通信和系统开销只关注纯粹的计算效率。4. 训练开销文档里用线性模型推出了一个很重要的经验结论前向传播约为2 × 参数量 × token数反向传播大约是前向的两倍所以前向 反向总共可近似写成6 × 参数量 × token数 6 \times \text{参数量} \times \text{token数}6×参数量×token数它背后其实就是矩阵乘法在forward和backward里的工作量拆分。5. 模型训练的数据加载numpy.memmap对于超大规模语料可以用numpy.memmap把磁盘文件映射成一个“按需访问”的映射对象就像一个磁盘上的“指针“一样当真正访问某一段数据时系统才会把那部分数据调入内存。pin_memory()默认情况下CPU张量放在可分页内存里而GPU在搬运数据前需要先复制到固定的非分页内存区域调用pin_memory()后张量会被放到固定内存中这样GPU能更直接地访问减少一次中间拷贝提高传输效率。non_blockingTrue如果数据已经在固定内存里那么.to(device, non_blockingTrue)时传输就可以尽量异步地进行不必让Python线程原地等待。这样一来GPU处理当前batch的同时CPU还能继续准备下一个batch不容易出现GPU等数据空转。6. 理解与反思本章核心是带来了资源核算的视角以往我基本只关注模型运行效果现在需要兼顾显存、精度、FLOPs、优化器状态、数据传输等硬件与系统层面问题将训练从单纯的优化问题拓展为资源约束下的系统工程问题。资源开销估算也建立起了对训练成本的直观认知能快速判断模型配置的开销与性能瓶颈。承接前文分词器的基础内容本章补齐了模型完整训练的核心逻辑是后续学习大模型训练细节的关键铺垫。