\documentclass[a4paper]{article}

% Basic packages
\usepackage[fontset=fandol]{ctex}
\usepackage{amsmath, amssymb}
\usepackage{graphicx}
\usepackage[margin=2.5cm]{geometry}
\usepackage[most]{tcolorbox}
\usepackage{etoolbox}
\usepackage{listings}
\usepackage{hyperref}
\usepackage{booktabs}
\usepackage{subcaption}
\usepackage{float}
\usepackage{tikz}
\usetikzlibrary{positioning,arrows.meta}
\IfFileExists{pgfplots.sty}{
    \usepackage{pgfplots}
    \pgfplotsset{compat=1.18}
}{}

% Highlight boxes
\newtcolorbox{knowledgebox}[1]{
    enhanced, colback=blue!5!white, colframe=blue!75!black, colbacktitle=blue!75!black,
    coltitle=white, fonttitle=\bfseries, title=#1, attach boxed title to top left={yshift=-2mm, xshift=2mm},
    boxrule=1pt, sharp corners
}

\newtcolorbox{importantbox}[1]{
    enhanced, colback=yellow!10!white, colframe=yellow!80!black, colbacktitle=yellow!80!black,
    coltitle=black, fonttitle=\bfseries, title=#1, sharp corners
}

\newtcolorbox{warningbox}[1]{
    enhanced, colback=red!5!white, colframe=red!75!black, colbacktitle=red!75!black,
    coltitle=white, fonttitle=\bfseries, title=#1, sharp corners
}

% Listings
\lstset{
    language=Python,
    basicstyle=\ttfamily\small,
    keywordstyle=\color{blue},
    stringstyle=\color{red!60!black},
    commentstyle=\color{green!60!black},
    breaklines=true,
    frame=single,
    numbers=left,
    numberstyle=\tiny\color{gray},
    captionpos=b,
    extendedchars=false
}

% Front-page metadata
\newcommand{\notetitle}{CS336 Lecture 2: Building a Model in PyTorch}
\newcommand{\noteauthors}{基于 Tatsu Hashimoto 授课内容整理}
\newcommand{\notedate}{2025年4月}
\newcommand{\videochannel}{Stanford Online}
\newcommand{\videopublishdate}{2025-04}
\newcommand{\videoduration}{1:19:18}
\newcommand{\videourl}{}
\newcommand{\videocoverpath}{cover.jpg}
\newcommand{\repourl}{https://github.com/hqhq1025/ai-course-notes}


% Footer with repo info
\usepackage{fancyhdr}
\pagestyle{fancy}
\fancyhf{}
\fancyfoot[C]{\thepage}
\fancyfoot[R]{\footnotesize\color{gray}\href{\repourl}{AI Course Notes}}
\renewcommand{\headrulewidth}{0pt}
\renewcommand{\footrulewidth}{0pt}

\begin{document}

\begin{titlepage}
\centering
{\Large 课程笔记\par}
\vspace{1.2cm}
{\huge\bfseries \notetitle\par}
\vspace{0.8cm}
{\large \noteauthors\par}
\vspace{0.3cm}
{\large \notedate\par}
\vspace{1.2cm}

\ifdefempty{\videocoverpath}{
{\small 请在生成笔记时填入视频封面图的本地路径。\par}
}{
\includegraphics[width=0.82\textwidth,height=0.45\textheight,keepaspectratio]{\videocoverpath}\par
}

\vfill
\begin{tcolorbox}[width=0.9\textwidth, colback=black!2!white, colframe=black!60, sharp corners]
\textbf{视频作者/频道}：\videochannel\par
\textbf{发布日期}：\videopublishdate\par
\textbf{视频时长}：\videoduration\par
\ifdefempty{\videourl}{}{
\textbf{视频链接}：\href{\videourl}{\nolinkurl{\videourl}}\par
}
\par
\textbf{项目主页}：\href{\repourl}{\nolinkurl{\repourl}}\par
\end{tcolorbox}
\end{titlepage}

\tableofcontents
\newpage

\section{引言：从概览走向实现}

上一讲先概览了语言模型、tokenization 以及从零构建模型的动机。本讲进入真正的实现层面，目标不是再讲 Transformer 的整体结构，而是把训练模型所需要的 \textbf{PyTorch 原语} 一层层搭起来。

\begin{importantbox}{本讲的三件事}
\begin{enumerate}
    \item 从张量开始，搞清楚模型参数、梯度、优化器状态和数据到底如何占用内存。
    \item 从矩阵乘法开始，学会对 FLOPs 和 FLOP/s 做资源核算。
    \item 从 `nn.Parameter`、`nn.Module`、优化器和训练循环出发，搭出一个最小但完整的训练系统。
\end{enumerate}
\end{importantbox}

讲义里的核心方法很朴素：\textbf{先算账，再写代码}。如果不知道一个张量有多少字节、一次矩阵乘法有多少 FLOPs、一次训练步会额外产生多少梯度和优化器状态，就很难判断模型为什么慢，也很难知道什么时候该换数据类型、换实现方式或者换训练策略。

\subsection{这一讲的阅读方式}

这一讲最值得反复看的不是某一段代码，而是背后的判断方式：
\begin{itemize}
    \item \textbf{定义}：这个对象是什么，PyTorch 里对应什么类型。
    \item \textbf{直觉}：为什么它会占内存，为什么它会花算力。
    \item \textbf{机制}：PyTorch / CUDA / autograd 实际怎么做。
    \item \textbf{应用}：这个原语在训练系统里被怎么使用。
    \item \textbf{局限}：这个近似或实现在哪里会失效。
\end{itemize}

阅读顺序建议：\textbf{Definition $\to$ Intuition $\to$ Mechanism $\to$ Application $\to$ Limitations}。

\begin{knowledgebox}{两类资源}
本讲只盯住两种资源：
\begin{itemize}
    \item \textbf{内存}：参数、梯度、激活值、优化器状态都要装得下。
    \item \textbf{计算}：训练和推理的成本通常都能归结成 FLOPs。
\end{itemize}
\end{knowledgebox}

\subsection{一个先算清楚的 napkin math}

先看两个典型估算。

\begin{enumerate}
    \item 训练一个 70B 参数模型，在 15T tokens 上，用 1024 张 H100，大概需要多久？
    \item 如果只看显存，8 张 80GB H100 上，使用朴素 AdamW，最多能训练多大的 dense 模型？
\end{enumerate}

\begin{importantbox}{6ND 规则——训练 FLOPs 的万能近似}
对于密集 Transformer，训练的总 FLOPs 可以用一个极其简洁的公式近似：
$$
C \approx 6 \times N \times D
$$
其中 $N$ 是模型参数量，$D$ 是训练 token 数。这里的 6 来自前向传播（$\approx 2ND$）加反向传播（$\approx 4ND$）。

这个公式忽略了 embedding 层、attention mask 计算、layernorm 等非矩阵乘法操作，但在规模估算时误差通常不超过 10\%--20\%。它是做训练预算时最重要的单个公式。
\end{importantbox}

\subsubsection{Worked Example: 70B 模型在 1024 张 H100 上训练多久}

我们把上面的估算展开为一步一步的计算：

\textbf{Step 1: 计算总 FLOPs}

模型参数量 $N = 70 \times 10^9$，训练数据量 $D = 15 \times 10^{12}$ tokens。由 6ND 公式：
$$
C = 6 \times 70 \times 10^9 \times 15 \times 10^{12} = 6.3 \times 10^{24} \text{ FLOPs}
$$

\textbf{Step 2: 计算集群有效算力}

H100 SXM 在 bfloat16 下的峰值算力约为 $989.5 \times 10^{12}$ FLOP/s $\approx 10^{15}$ FLOP/s。假设 MFU = 0.5（即有效利用率 50\%），则单卡有效算力为：
$$
\text{单卡有效} = 10^{15} \times 0.5 = 5 \times 10^{14} \text{ FLOP/s}
$$

1024 张卡的集群每天提供的 FLOPs：
$$
\text{日 FLOPs} = 1024 \times 5 \times 10^{14} \times 86400 = 4.42 \times 10^{22} \text{ FLOPs/day}
$$

\textbf{Step 3: 换算训练天数}
$$
\text{训练天数} = \frac{6.3 \times 10^{24}}{4.42 \times 10^{22}} \approx 143 \text{ 天}
$$

\begin{table}[H]
\centering
\begin{tabular}{p{0.28\textwidth}p{0.25\textwidth}p{0.30\textwidth}}
\toprule
\textbf{参数} & \textbf{数值} & \textbf{来源} \\
\midrule
模型参数量 $N$ & $70 \times 10^9$ & 模型架构定义 \\
训练 token 数 $D$ & $15 \times 10^{12}$ & 训练数据量 \\
总 FLOPs $C = 6ND$ & $6.3 \times 10^{24}$ & 6ND 公式 \\
H100 bf16 峰值 & $\sim 10^{15}$ FLOP/s & NVIDIA 规格表 \\
假设 MFU & 0.5 & 经验估计 \\
GPU 数量 & 1024 & 集群规模 \\
训练天数 & $\sim$143 天 & 总 FLOPs / 日 FLOPs \\
\bottomrule
\end{tabular}
\caption{70B 模型训练时间估算的完整参数表}
\end{table}

\begin{warningbox}{MFU 的敏感性}
上面的计算中 MFU 设为 0.5，但实际 MFU 可能在 0.3--0.6 之间波动。如果 MFU 从 0.5 降到 0.35，训练时间就会从 143 天变成 204 天——差出两个月。因此 MFU 的微小变化会导致训练成本的巨大波动，这也是工程优化的核心目标之一。
\end{warningbox}

\subsubsection{Worked Example: AdamW 训练 7B 模型的内存分解}

以一个 7B 参数模型为例，详细分解使用 AdamW 优化器时的 GPU 显存占用。

\textbf{假设}：使用混合精度训练，参数和梯度为 bf16（2 bytes），优化器状态中的主权重（master weights）和一阶/二阶矩为 fp32（4 bytes）。

$$
N = 7 \times 10^9 \text{ 个参数}
$$

\begin{table}[H]
\centering
\begin{tabular}{p{0.28\textwidth}p{0.14\textwidth}p{0.15\textwidth}p{0.28\textwidth}}
\toprule
\textbf{组件} & \textbf{dtype} & \textbf{每参数字节} & \textbf{7B 模型内存} \\
\midrule
模型参数（训练副本） & bf16 & 2 & 14 GB \\
梯度 & bf16 & 2 & 14 GB \\
主权重（master weights） & fp32 & 4 & 28 GB \\
一阶矩 $m_t$ & fp32 & 4 & 28 GB \\
二阶矩 $v_t$ & fp32 & 4 & 28 GB \\
\midrule
\textbf{静态总计} & & \textbf{16} & \textbf{112 GB} \\
\bottomrule
\end{tabular}
\caption{AdamW 混合精度训练 7B 模型的内存分解（不含激活值）}
\end{table}

这意味着仅模型参数和优化器状态就需要 112 GB 显存，还没有算激活值。以 80 GB 的 H100 来算，\textbf{至少需要 2 张卡}才能装下静态状态，实际上考虑激活值可能需要 4 张或更多。

如果用纯 fp32 训练（不做混合精度），每个参数需要：
$$
4\;(\text{参数}) + 4\;(\text{梯度}) + 4\;(m_t) + 4\;(v_t) = 16 \text{ bytes}
$$
7B 模型静态内存为 $7 \times 10^9 \times 16 = 112$ GB，和混合精度的总量恰好相同——但混合精度的优势在于前向/反向计算速度更快、激活值更小。

\begin{knowledgebox}{为什么混合精度还要存 fp32 主权重}
混合精度训练的核心思想是：用低精度做前向和反向计算以加速，但用 fp32 精度累积梯度更新。如果直接在 bf16 参数上做梯度更新，微小的更新（如 $\eta \cdot g \ll 1$）会因为 bf16 的有限精度被截断为零，导致训练停滞。fp32 主权重保证了累积更新的精度。
\end{knowledgebox}

\section{Memory Accounting}

\subsection{张量是深度学习的基本单位}

在 PyTorch 里，几乎所有东西最终都是张量：
\begin{itemize}
    \item 模型参数
    \item 梯度
    \item 优化器状态
    \item 数据
    \item 激活值
\end{itemize}

\begin{knowledgebox}{张量的两件事}
一个 tensor 的内存成本只看两件事：
\begin{itemize}
    \item 元素个数 `numel()`
    \item 每个元素占多少字节 \texttt{element\_size()}
\end{itemize}
因此最简单的内存估计就是
\[
\text{memory} = \text{numel} \times \text{element size}.
\]
\end{knowledgebox}

\subsection{Worked example: 一个小张量到底花多少内存}

先看一个最小例子。一个 \(4\times 8\) 的 float32 tensor 有 32 个元素，每个元素 4 字节，所以总内存是：
\[
4 \times 8 \times 4 = 128 \text{ bytes}.
\]
这类算式很小，但它的逻辑和大模型完全一样。训练系统里的很多 bug，本质上就是“把 128 bytes 的直觉误搬到 128 GB 的规模上”。

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.23\textwidth}p{0.18\textwidth}p{0.2\textwidth}}
\toprule
\textbf{张量} & \textbf{形状} & \textbf{dtype} & \textbf{内存} \\
\midrule
Toy matrix & \(4\times 8\) & float32 & 128 B \\
Hidden state batch & \(32\times 2048\times 4096\) & bfloat16 & \(\approx 512\) MB \\
FFN weight & \(4096\times 16384\) & float32 & \(\approx 256\) MB \\
Adam states & same as weight & float32 \(\times 2\) & \(\approx 512\) MB \\
\bottomrule
\end{tabular}
\caption{从玩具张量到大模型参数，内存公式本质上没有变}
\end{table}

\begin{importantbox}{实战里最常见的误解}
很多人只看参数量，不看激活值和 optimizer state。实际上，训练时最贵的往往不是模型前向本身，而是“参数 + 梯度 + 优化器状态 + 激活值”四者叠加后的总账。
\end{importantbox}

\subsection{Worked example: 为什么激活值会突然吃掉显存}

参数内存是静态的，激活值内存是随 batch size 和 sequence length 线性增长的。对于一个形状为 \(B\times L\times H\) 的激活张量，如果使用 bfloat16，它的内存大约是：
\[
2 \times B \times L \times H \text{ bytes}.
\]
这就是为什么你明明觉得“模型参数没那么大”，训练时还是会 OOM：真正爆掉的可能是中间激活，而不是权重本身。

\begin{table}[H]
\centering
\begin{tabular}{p{0.23\textwidth}p{0.18\textwidth}p{0.22\textwidth}p{0.23\textwidth}}
\toprule
\textbf{对象} & \textbf{典型形状} & \textbf{每元素字节数} & \textbf{增长方式} \\
\midrule
参数 & 固定形状 & 2 或 4 & 随模型规模增长 \\
激活值 & \(B\times L\times H\) & 2 或 4 & 随 batch 和序列长度增长 \\
梯度 & 与参数同形状 & 2 或 4 & 每次反向都要存 \\
优化器状态 & 与参数同形状 & 4 或更多 & 只和参数量相关 \\
\bottomrule
\end{tabular}
\caption{训练显存最重要的差别：参数是固定成本，激活值是动态成本}
\end{table}

\begin{importantbox}{为什么这件事和 Assignment 1 直接相关}
你在实现 Transformer 时，batch size、sequence length、hidden size、attention head 数都会联动影响激活值的总量。只盯参数量，往往会把最关键的显存瓶颈看漏。
\end{importantbox}

\subsection{常见浮点类型}

浮点数的表示由三部分组成：\textbf{符号位}（sign）、\textbf{指数位}（exponent）和\textbf{尾数位}（mantissa/significand）。不同数据类型在这三部分的分配上做了不同的权衡。

\begin{table}[H]
\centering
\begin{tabular}{lccccl}
\toprule
\textbf{类型} & \textbf{字节} & \textbf{指数位} & \textbf{尾数位} & \textbf{动态范围} & \textbf{训练常见用途} \\
\midrule
float32 (FP32) & 4 & 8 & 23 & $\sim 10^{\pm 38}$ & 主权重、优化器状态 \\
float16 (FP16) & 2 & 5 & 10 & $\sim 10^{\pm 5}$ & 部分前向/激活 \\
bfloat16 (BF16) & 2 & 8 & 7 & $\sim 10^{\pm 38}$ & 训练和推理的主流选择 \\
fp8 (E4M3) & 1 & 4 & 3 & $\sim 10^{\pm 2}$ & 前沿训练/推理实验 \\
fp8 (E5M2) & 1 & 5 & 2 & $\sim 10^{\pm 5}$ & 梯度计算 \\
\bottomrule
\end{tabular}
\caption{常见数值类型的位分配与动态范围对比}
\end{table}

这里最重要的结论是：\textbf{bfloat16 不是”更小的 float32”，而是”保留了 float32 的指数范围，但牺牲一些尾数精度”}。在深度学习里，这种取舍通常比 float16 更稳。

\begin{warningbox}{bf16 vs fp16：看起来都是 2 字节，但行为完全不同}
fp16 只有 5 位指数，动态范围仅约 $10^{\pm 5}$，这意味着绝对值小于 $\sim 6 \times 10^{-8}$ 的数会被截断为零（underflow），大于 65504 的数会溢出为 inf。在训练中，梯度值经常非常小，fp16 的 underflow 会导致梯度”消失”，必须配合 \textbf{loss scaling} 来人为放大 loss，使梯度落入 fp16 的可表示范围。

bf16 有 8 位指数，动态范围和 fp32 一样大（$\sim 10^{\pm 38}$），因此\textbf{不需要 loss scaling}。代价是尾数只有 7 位（fp16 有 10 位），精度较低。但实践表明，对于大多数深度学习任务，动态范围比尾数精度更重要。

\textbf{结论}：如果硬件支持 bf16（如 A100、H100），优先选 bf16 而非 fp16。
\end{warningbox}

\begin{importantbox}{为什么训练比推理更难低精度化}
训练时会同时涉及参数、梯度、优化器状态和反向传播链式法则。数值误差会在更新过程中累积，因此训练对数值稳定性更敏感；推理阶段没有反向传播，通常可以更激进地量化。
\end{importantbox}

\subsubsection{Worked Example: 混合精度节省了多少计算和内存}

考虑一个矩阵乘法 $Y = XW$，其中 $X \in \mathbb{R}^{4096 \times 4096}$，$W \in \mathbb{R}^{4096 \times 4096}$。

\textbf{内存对比}：
\begin{table}[H]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{dtype} & \textbf{$X$ 内存} & \textbf{$X + W$ 总内存} \\
\midrule
float32 & $4096^2 \times 4 = 64$ MB & 128 MB \\
bfloat16 & $4096^2 \times 2 = 32$ MB & 64 MB \\
fp8 & $4096^2 \times 1 = 16$ MB & 32 MB \\
\bottomrule
\end{tabular}
\caption{同一矩阵在不同精度下的内存占用}
\end{table}

\textbf{计算吞吐对比}（以 H100 SXM 为例）：
\begin{table}[H]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{dtype} & \textbf{峰值 TFLOP/s} & \textbf{相对 FP32 加速} \\
\midrule
FP32 & 67 & $1\times$ \\
TF32 (Tensor Core) & 495 & $\sim 7\times$ \\
BF16/FP16 (Tensor Core) & 990 & $\sim 15\times$ \\
FP8 (Tensor Core) & 1979 & $\sim 30\times$ \\
\bottomrule
\end{tabular}
\caption{H100 SXM 在不同数据类型下的理论峰值算力}
\end{table}

这解释了为什么混合精度训练不只是”省内存”——它同时带来了巨大的\textbf{计算加速}。从 fp32 切到 bf16，理论上可以获得约 15 倍的矩阵乘法加速。

\subsection{不同模型规模的内存占用对比}

下表展示了从 1B 到 70B 不同规模模型在使用 AdamW 混合精度训练时的静态内存需求（不含激活值）：

\begin{table}[H]
\centering
\begin{tabular}{lccccc}
\toprule
\textbf{模型规模} & \textbf{参数 (bf16)} & \textbf{梯度 (bf16)} & \textbf{主权重 (fp32)} & \textbf{Adam 状态 (fp32)} & \textbf{总计} \\
\midrule
1B & 2 GB & 2 GB & 4 GB & 8 GB & 16 GB \\
7B & 14 GB & 14 GB & 28 GB & 56 GB & 112 GB \\
13B & 26 GB & 26 GB & 52 GB & 104 GB & 208 GB \\
30B & 60 GB & 60 GB & 120 GB & 240 GB & 480 GB \\
70B & 140 GB & 140 GB & 280 GB & 560 GB & 1120 GB \\
\bottomrule
\end{tabular}
\caption{不同规模模型使用 AdamW 混合精度训练的静态内存需求}
\end{table}

\begin{knowledgebox}{快速心算规则}
对于 AdamW 混合精度训练，每个参数的静态内存约为 $2 + 2 + 4 + 4 + 4 = 16$ bytes。因此：
\begin{itemize}
    \item 1B 模型 $\approx$ 16 GB（1 张 H100 刚好）
    \item 7B 模型 $\approx$ 112 GB（至少 2 张 H100）
    \item 70B 模型 $\approx$ 1120 GB（至少 14 张 H100，纯装参数和状态）
\end{itemize}
这还没有算激活值！实际训练通常需要更多显存或使用模型并行。
\end{knowledgebox}

\subsection{GPU 内存层次与带宽}

理解 GPU 的内存层次对于优化训练性能至关重要。GPU 的内存系统是分层的，越靠近计算单元速度越快、容量越小：

\begin{table}[H]
\centering
\begin{tabular}{p{0.20\textwidth}p{0.14\textwidth}p{0.18\textwidth}p{0.30\textwidth}}
\toprule
\textbf{层级} & \textbf{容量} & \textbf{带宽} & \textbf{说明} \\
\midrule
寄存器 & KB 级 & --- & 每个线程私有 \\
共享内存/L1 & 128--228 KB/SM & $\sim$100 TB/s & SM 内共享，用户可编程 \\
L2 缓存 & 50--60 MB & $\sim$12 TB/s & 全 GPU 共享 \\
HBM (显存) & 80 GB & 3.35 TB/s & H100 SXM 规格 \\
NVLink & --- & 900 GB/s & GPU 间互联 \\
PCIe 5.0 & --- & 128 GB/s & CPU--GPU 互联 \\
\bottomrule
\end{tabular}
\caption{H100 SXM GPU 的内存层次（数值为近似值）}
\end{table}

\begin{importantbox}{Arithmetic Intensity 与带宽瓶颈}
一个操作是\textbf{计算瓶颈}（compute-bound）还是\textbf{带宽瓶颈}（memory-bound），取决于它的 \textbf{arithmetic intensity}（每字节数据传输执行的 FLOPs 数）：
$$
\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{Bytes accessed}}
$$
\begin{itemize}
    \item \textbf{矩阵乘法}：arithmetic intensity 高（$O(n)$ FLOPs per byte），通常是 compute-bound，能充分利用 Tensor Core。
    \item \textbf{逐元素操作}（如 ReLU、LayerNorm、softmax）：arithmetic intensity 低（$O(1)$ FLOPs per byte），通常是 memory-bound，瓶颈在显存带宽。
\end{itemize}
这就解释了为什么 Transformer 的矩阵乘法可以通过低精度获得近线性加速，而 LayerNorm 等操作的加速效果有限。
\end{importantbox}

\begin{knowledgebox}{Roofline 模型}
Roofline 模型是分析 GPU 性能的经典工具。它给出了给定 arithmetic intensity 下的理论性能上限：
$$
\text{Attainable FLOP/s} = \min\left(\text{Peak FLOP/s},\;\text{Bandwidth} \times \text{Arithmetic Intensity}\right)
$$
当 arithmetic intensity 低于某个拐点（= Peak FLOP/s / Bandwidth）时，性能受带宽限制；超过拐点后，性能受计算能力限制。对 H100 SXM bf16 而言，这个拐点约为 $990 \text{ TFLOP/s} / 3.35 \text{ TB/s} \approx 295$ FLOPs/Byte。
\end{knowledgebox}

\subsection{CPU 与 GPU 的内存位置}

默认情况下，PyTorch tensor 存在 CPU 内存中。要利用 GPU 的并行计算能力，必须显式把 tensor 迁移到 GPU 显存中。

\begin{table}[H]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{位置} & \textbf{特点} & \textbf{常见操作} \\
\midrule
CPU paged memory & 默认、可分页、适合常规处理 & 数据预处理、批生成 \\
CPU pinned memory & 页锁定，可异步拷贝到 GPU & 数据加载器预取 \\
GPU memory & 计算速度最快、容量最宝贵 & 参数、激活、临时计算结果 \\
\bottomrule
\end{tabular}
\caption{训练中常见的内存位置}
\end{table}

\begin{knowledgebox}{Pinned memory 的意义}
把 CPU 上的 batch 放到 pinned memory 后，可以用 \texttt{non\_blocking=True} 异步拷贝到 GPU。这样可以把“CPU 取下一批数据”和“GPU 处理当前批数据”重叠起来。
\end{knowledgebox}

\subsection{数据位置的局限}

CPU 内存、pinned memory 和 GPU 显存各有分工，但也各有上限。CPU 便宜但慢，GPU 快但贵，pinned memory 便于搬运但不是无限大。训练系统经常死在“所有东西都放 GPU”这种简单化思路上。真正的工程做法是让数据尽量在正确的位置停留尽量短的时间。

\subsection{Tensor storage, stride 和 view}

PyTorch tensor 不只是“值的集合”，它更像是：
\begin{itemize}
    \item 一块连续或者非连续的底层存储；
    \item 加上一组元数据，告诉你每个维度怎么映射到存储地址；
    \item 以及形状、dtype、device 等信息。
\end{itemize}

在二维张量里，`stride(0)` 表示跨行要跳多少元素，`stride(1)` 表示跨列要跳多少元素。对一个标准行主序矩阵来说，行 stride 通常等于列数，列 stride 通常为 1。

\begin{importantbox}{view 是免费的，copy 不是}
`view`、切片和转置很多时候都只是不同的“视图”，不会复制底层数据，因此几乎不耗额外内存。真正触发 `.contiguous()` 或显式拷贝时，才会同时增加内存和计算成本。
\end{importantbox}

一个容易踩坑的点是：不是所有视图都能再 `view` 成任意形状。若张量是非连续的，直接 `view` 常常会报错，这时要先调用 `.contiguous()`。

\subsection{View 的机制细节}

`view` 的前提是底层存储布局能被新的形状解释。如果张量是转置后的非连续布局，就不能直接靠 `view` 改形状。这就是为什么 `contiguous()` 不是多余操作，而是把“逻辑视图”重新落回“物理连续存储”的一步。

\begin{warningbox}{为什么这一点重要}
很多模型代码看起来只是“换了个形状”，实际上却偷偷复制了一大块张量。只要复制发生得不该发生，内存就会涨，性能也会掉。
\end{warningbox}

\subsection{切片、转置和掩码}

切片、列选择、转置等操作通常只返回新的视图而不是复制数据。对于 causal attention 这类场景，`triu` 很常见，因为它可以快速生成上三角矩阵，用来屏蔽未来 token 的信息。

\begin{knowledgebox}{causal mask 的直觉}
如果 \(M[i,j]=1\) 表示位置 \(i\) 对位置 \(j\) 有贡献，那么自回归语言模型要求未来位置不能泄露到过去，于是只保留上三角或下三角的合法连接。
\end{knowledgebox}

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.26\textwidth}p{0.32\textwidth}}
\toprule
\textbf{操作} & \textbf{是否复制数据} & \textbf{常见用途} \\
\midrule
切片 `x[:, 1]` & 通常不复制 & 选列、选 token、选 channel \\
转置 `transpose` & 不复制，但可能非连续 & 改变矩阵方向 \\
`contiguous()` & 会复制 & 把非连续视图落回连续内存 \\
`triu()` & 生成新张量 & causal mask、上三角约束 \\
\bottomrule
\end{tabular}
\caption{哪些操作免费、哪些会复制，是内存优化的基础判断}
\end{table}

\subsection{逐元素运算和矩阵乘法}

逐元素运算包括加法、乘法、平方、开方、`triu` 等。这类操作一般是线性的，FLOPs 规模约为 \(O(mn)\)。

真正的深度学习主角还是矩阵乘法：
\[
X \in \mathbb{R}^{B\times D}, \quad W \in \mathbb{R}^{D\times K}, \quad Y = XW \in \mathbb{R}^{B\times K}.
\]
这里的 FLOPs 约为
\[
2BDK,
\]
因为每个输出元素都需要一次乘法和一次加法。

\begin{table}[H]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{运算} & \textbf{复杂度} & \textbf{备注} \\
\midrule
逐元素加/乘 & \(O(mn)\) & 规模线性 \\
矩阵乘法 & \(O(mnp)\) & 通常是计算瓶颈 \\
反向传播 & 约为前向的 2 倍到 3 倍 & 依赖具体图结构 \\
\bottomrule
\end{tabular}
\caption{训练里最常见的计算量级}
\end{table}

当输入张量多出 batch 维和 sequence 维时，PyTorch 的矩阵乘法会自动在前面的维度上广播并批处理。也就是说，只要最后两维满足矩阵乘法规则，前面的维度会被视作并行的“外层循环”。

\subsection{内存碎片与 OOM 调试}

GPU 内存并不是一个可以任意使用的大池子。PyTorch 的 CUDA 内存分配器会缓存已释放的内存块以加速后续分配，但这也会导致\textbf{内存碎片化}——总可用内存看起来足够，但没有足够大的连续块来满足一次大的分配请求。

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.30\textwidth}p{0.32\textwidth}}
\toprule
\textbf{OOM 类型} & \textbf{表现} & \textbf{排查方法} \\
\midrule
真实 OOM & 显存确实不够 & 减小 batch size 或用模型并行 \\
碎片化 OOM & 总量够但连续块不够 & \texttt{torch.cuda.memory\_summary()} \\
泄漏式 OOM & 显存随步数缓慢增长 & 检查是否把 tensor 挂在 history 上 \\
激活值爆炸 & 某步突然 OOM & 检查输入序列长度异常值 \\
\bottomrule
\end{tabular}
\caption{常见的 GPU OOM 场景与排查方法}
\end{table}

\begin{warningbox}{常见的内存泄漏陷阱}
以下写法会意外保留计算图，导致内存持续增长：
\begin{itemize}
    \item 把 \texttt{loss} 直接 append 到 Python list 中（应使用 \texttt{loss.item()}）
    \item 在循环外部持有中间 tensor 的引用而不 detach
    \item 忘记调用 \texttt{optimizer.zero\_grad(set\_to\_none=True)}，导致梯度在多步间累积
\end{itemize}
\end{warningbox}

PyTorch 提供了一套实用的内存调试工具：

\begin{lstlisting}[caption={GPU 内存调试常用命令}]
# 查看当前显存使用
print(torch.cuda.memory_allocated() / 1e9, "GB allocated")
print(torch.cuda.memory_reserved() / 1e9, "GB reserved")

# 打印详细内存摘要
print(torch.cuda.memory_summary())

# 记录内存快照用于分析
torch.cuda.memory._record_memory_history()
# ... run some code ...
torch.cuda.memory._dump_snapshot("mem_snapshot.pickle")
\end{lstlisting}

\subsection{Einops 和维度命名}

原生 PyTorch 容易在 `-2`、`-1` 这种负索引里把维度写乱。`einops` 的价值在于把维度变成可读的名字。

\begin{knowledgebox}{维度命名的价值}
\begin{itemize}
    \item `batch seq hidden` 一眼就知道张量语义。
    \item `einsum` 可以把广播和求和写得更清楚。
    \item `reduce` 和 `rearrange` 让 reshape 逻辑更不容易出错。
\end{itemize}
\end{knowledgebox}

这讲里最常见的三种写法是：
\begin{itemize}
    \item `einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")`
    \item `reduce(x, "... hidden -> ...", "sum")`
    \item `rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)`
\end{itemize}

\section{Compute Accounting}

\subsection{FLOPs 和 FLOP/s}

FLOP 是一次浮点操作，比如一次加法或乘法。FLOP/s 则是硬件每秒能执行多少浮点操作。这里要避免两个缩写混淆：

\begin{itemize}
    \item \textbf{FLOPs}：算了多少浮点操作。
    \item \textbf{FLOP/s}：每秒能做多少浮点操作。
\end{itemize}

\begin{importantbox}{资源核算的核心问题}
训练时间约等于
\[
\frac{\text{总 FLOPs}}{\text{有效 FLOP/s}}.
\]
所以除了算法本身，我们还必须知道硬件的有效利用率。
\end{importantbox}

\subsection{线性模型的 FLOPs}

先从最简单的线性模型开始。若有 \(B\) 个样本、每个样本维度为 \(D\)、输出维度为 \(K\)，那么前向传播的矩阵乘法
\[
Y = XW
\]
大约需要
\[
2BDK
\]
个 FLOPs。

后向传播则需要对参数和中间激活分别做链式法则，整体大约是前向的两倍，所以线性层的完整训练步常用近似是：
\[
\text{forward} + \text{backward} \approx 6BDK.
\]

\begin{knowledgebox}{为什么是 6}
粗略地看：
\begin{itemize}
    \item 前向：一次乘法 + 一次加法，约 \(2BDK\)
    \item 反向：对参数和输入各再做一次类似的乘加，约 \(4BDK\)
\end{itemize}
合起来就是约 \(6BDK\)。
\end{knowledgebox}

\subsection{更一般的近似}

对于 Transformer 这类模型，虽然结构更复杂，但一阶近似下，训练 FLOPs 仍可以写成：
\[
6 \times (\text{token 数}) \times (\text{参数量}).
\]
这也是前面 napkin math 的来源。

\subsection{不同 GPU 代际的算力对比}

选择硬件时，不仅要看峰值 FLOP/s，还要看内存容量和带宽。不同代际 GPU 的关键参数对比如下：

\begin{table}[H]
\centering
\begin{tabular}{p{0.14\textwidth}p{0.13\textwidth}p{0.13\textwidth}p{0.14\textwidth}p{0.12\textwidth}p{0.12\textwidth}}
\toprule
\textbf{GPU} & \textbf{HBM} & \textbf{带宽} & \textbf{BF16 Peak} & \textbf{FP8 Peak} & \textbf{年份} \\
\midrule
V100 SXM & 32 GB & 900 GB/s & --- & --- & 2017 \\
A100 SXM & 80 GB & 2.0 TB/s & 312 TF & --- & 2020 \\
H100 SXM & 80 GB & 3.35 TB/s & 990 TF & 1979 TF & 2023 \\
H200 SXM & 141 GB & 4.8 TB/s & 990 TF & 1979 TF & 2024 \\
B200 SXM & 192 GB & 8.0 TB/s & 2250 TF & 4500 TF & 2025 \\
\bottomrule
\end{tabular}
\caption{NVIDIA 数据中心 GPU 关键参数对比（TF = TFLOP/s）}
\end{table}

\begin{knowledgebox}{为什么 H200 和 H100 算力相同但价值更高}
H200 和 H100 使用相同的 GPU 芯片（GH200），BF16 算力完全相同。H200 的优势在于将 HBM 从 80 GB 升级到 141 GB、带宽从 3.35 TB/s 升级到 4.8 TB/s。对于大模型推理（memory-bound 场景），更大的显存意味着可以加载更大的模型，更高的带宽意味着更快的 token 生成速度。
\end{knowledgebox}

\subsection{模型 FLOPs 利用率 MFU}

硬件规格表给的是“峰值”性能，但真实训练不可能一直跑满。于是引入：
\[
\text{MFU} = \frac{\text{actual FLOP/s}}{\text{promised FLOP/s}}.
\]

MFU 不是一个抽象学术指标，而是直接决定你有多少硬件预算被真正转化成训练速度。一般来说，MFU 大于 0.5 已经算不错，尤其当矩阵乘法占主导时更是如此。

\begin{warningbox}{为什么 float32 和 bfloat16 的 FLOP/s 不一样}
硬件规格里的“峰值 FLOP/s”强依赖数据类型。很多 GPU 对 bfloat16/float16 的吞吐量远高于 float32，所以同一块卡上，切换 dtype 会显著改变你的实际吞吐。
\end{warningbox}

\subsection{Worked example: 把 FLOPs 直接换成训练时间}

如果你知道总 FLOPs 和硬件的有效 FLOP/s，就可以直接估计 wall-clock time。以 70B 参数、15T tokens 的训练为例：
\[
6 \times 70\times 10^9 \times 15\times 10^{12}
\]
是总 FLOPs。若用 1024 张 H100，且把有效利用率设成 0.5，那么每天能提供的 FLOPs 就是：
\[
1024 \times 0.5 \times \text{H100 peak FLOP/s} \times 86400.
\]
把前者除以后者，就得到大约 144 天。

\begin{table}[H]
\centering
\begin{tabular}{p{0.24\textwidth}p{0.24\textwidth}p{0.23\textwidth}}
\toprule
\textbf{量} & \textbf{公式} & \textbf{含义} \\
\midrule
总 FLOPs & \(6NT\) & \(N\) 是参数量，\(T\) 是 token 数 \\
日 FLOPs & \(G \cdot \eta \cdot 86400\) & \(G\) 是峰值 FLOP/s，\(\eta\) 是 MFU \\
训练天数 & \(\text{总 FLOPs}/\text{日 FLOPs}\) & 直接换算 wall-clock \\
\bottomrule
\end{tabular}
\caption{把算力预算换成时间，是训练计划里最实用的一个公式}
\end{table}

\begin{warningbox}{这个估算的局限}
这个公式没有把通信、数据等待、checkpoint 开销、非矩阵乘法算子和 pipeline bubble 细节全部展开，所以它是预算级估算，不是秒级预测。
\end{warningbox}

\subsection{反向传播的 FLOPs}

以两层线性模型为例：
\[
x \xrightarrow{w_1} h_1 \xrightarrow{w_2} h_2 \rightarrow \text{loss}.
\]
前向已经需要两次矩阵乘法。反向传播时，对 \(w_2\)、\(h_1\)、\(w_1\) 以及输入的梯度分别做链式法则，整体 FLOPs 约为前向的两倍。因此整个训练步常见的经验法则仍然是：
\[
\text{train step FLOPs} \approx 6 \times (\text{tokens}) \times (\text{parameters}).
\]

\begin{importantbox}{这个近似什么时候够用}
只要你在做的是规模估算、预算估算、训练时长估算，这个近似就足够好。它不是为了证明定理，而是为了让你能快速判断“这个实验值不值得跑”。
\end{importantbox}

\subsection{Worked example: 两层线性网络的训练步}

假设输入 \(x\in\mathbb{R}^{B\times D}\)，第一层权重 \(w_1\in\mathbb{R}^{D\times D}\)，第二层权重 \(w_2\in\mathbb{R}^{D\times K}\)。那么：
\begin{itemize}
    \item 前向：\(xw_1\) 再乘 \(w_2\)
    \item 反向：对 \(w_2\)、\(h_1\)、\(w_1\) 求梯度
\end{itemize}

如果你把每个矩阵乘法都按 \(2\) FLOPs / 输出元素来算，最后就会得到“前向约 2 倍参数乘 token 数、反向约再来 4 倍”的总量级。这个例子说明：复杂模型的 FLOPs 近似，本质上是许多小矩阵乘法的叠加。

\begin{table}[H]
\centering
\begin{tabular}{p{0.2\textwidth}p{0.26\textwidth}p{0.22\textwidth}p{0.2\textwidth}}
\toprule
\textbf{阶段} & \textbf{主要操作} & \textbf{FLOPs 级别} & \textbf{含义} \\
\midrule
前向 & 两次 matmul & \(2BD(D+K)\) & 算预测 \\
反向 & 对参数和输入求导 & 约为前向的 2 倍 & 传播梯度 \\
训练步 & 前向 + 反向 & 约 \(6BD(D+K)\) & 完整更新一次 \\
\bottomrule
\end{tabular}
\caption{两层线性网络是理解训练 FLOPs 的最小实例}
\end{table}

\section{Models}

\subsection{\texttt{nn.Parameter}}

PyTorch 里模型参数通常以 `nn.Parameter` 的形式保存。它和普通 tensor 类似，但会自动被 `nn.Module` 注册为可训练参数，并出现在 `parameters()` 和 \texttt{state\_dict()} 里。

\begin{knowledgebox}{参数与普通 tensor 的差异}
`nn.Parameter` 的本质还是 tensor，但它明确表达了“这是要被训练的东西”。这会影响：
\begin{itemize}
    \item 自动求导
    \item 模型参数遍历
    \item checkpoint 保存/恢复
\end{itemize}
\end{knowledgebox}

\subsection{初始化：为什么要按输入维度缩放}

如果直接用高斯随机初始化，输出的尺度会随着输入维度增长而变大。一个简单的做法是除以 \(\sqrt{\text{input\_dim}}\)，这样可以让输出的方差大致稳定。

\begin{importantbox}{Xavier 的直觉}
初始化要尽量让前向传播和反向传播中的数值尺度都保持稳定。最简单的经验就是：参数规模要随输入维度做归一化。
\end{importantbox}

在更实际的实现里，通常会使用 truncated normal，把极端离群值裁掉，进一步减少训练不稳定的风险。

\subsection{自定义模块}

下面这类简单线性模块很关键，因为它展示了 `nn.Module` 的最小骨架。

\begin{lstlisting}
class Linear(nn.Module):
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.weight = nn.Parameter(
            torch.randn(input_dim, output_dim) / np.sqrt(input_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight
\end{lstlisting}

再往上堆叠多个层，就能得到一个简单的深度线性模型。重点不是”线性模型很强”，而是”\textbf{理解模块化结构如何把参数组织起来}”。

\subsection{Transformer 参数量估算}

理解模型参数量的具体构成，对于内存预算和训练规划至关重要。以一个标准的 decoder-only Transformer 为例：

\begin{table}[H]
\centering
\begin{tabular}{p{0.25\textwidth}p{0.30\textwidth}p{0.28\textwidth}}
\toprule
\textbf{组件} & \textbf{参数量} & \textbf{说明} \\
\midrule
Token embedding & $V \times d$ & $V$: 词表大小，$d$: 隐藏维度 \\
每层 QKV 投影 & $3 \times d \times d$ & 三个矩阵各 $d \times d$ \\
每层输出投影 & $d \times d$ & Attention 输出映射 \\
每层 FFN & $2 \times d \times 4d$ & 上投影 + 下投影 \\
每层 LayerNorm & $2 \times 2d$ & 两个 LN 各有 $\gamma, \beta$ \\
LM Head & 通常与 embedding 共享 & 不额外计数 \\
\midrule
\textbf{每层合计} & $\approx 12d^2$ & 忽略 LN 的 $O(d)$ 项 \\
\textbf{$L$ 层总计} & $\approx 12Ld^2 + Vd$ & \\
\bottomrule
\end{tabular}
\caption{Decoder-only Transformer 的参数量分解}
\end{table}

\textbf{Worked Example}：LLaMA-2 7B 的参数量验证

LLaMA-2 7B 的超参数为 $d = 4096$, $L = 32$, $V = 32000$, FFN 隐藏维度为 $11008$（使用 SwiGLU，实际为 $\frac{8}{3}d$ 上取整）。

\begin{itemize}
    \item Embedding: $32000 \times 4096 = 131M$
    \item 每层 Attention: $4 \times 4096^2 = 67M$（QKV + O）
    \item 每层 FFN (SwiGLU): $3 \times 4096 \times 11008 = 135M$（gate + up + down）
    \item 每层 LN: $2 \times 2 \times 4096 = 16K$（可忽略）
    \item 32 层合计: $32 \times (67M + 135M) = 6,464M$
    \item 最终 LN + embedding: $\approx 131M + 16K$
    \item \textbf{总计}: $\approx 6,738M \approx 6.7B$
\end{itemize}

这和”7B”的命名基本吻合（考虑到 RoPE、bias 等细节的差异）。

\begin{importantbox}{$12d^2$ 速算法}
对于标准 Transformer（FFN 倍率为 4），每层参数约为 $12d^2$。因此给定隐藏维度 $d$ 和层数 $L$，总参数量（不含 embedding）约为 $12Ld^2$。这个公式在做 napkin math 时非常方便：
\begin{itemize}
    \item $d=4096, L=32$: $12 \times 32 \times 4096^2 \approx 6.4B$
    \item $d=5120, L=40$: $12 \times 40 \times 5120^2 \approx 12.6B$
    \item $d=8192, L=80$: $12 \times 80 \times 8192^2 \approx 64.4B$
\end{itemize}
\end{importantbox}

\begin{knowledgebox}{\texttt{state\_dict} 里有什么}
\texttt{state\_dict()} 负责保存模型的状态。对一个模块来说，它通常包括：
\begin{itemize}
    \item 所有 `nn.Parameter`
    \item 有状态层的缓存
\end{itemize}
所以 checkpoint 不只是保存参数本身，而是保存整个训练状态。
\end{knowledgebox}

\subsection{模型参数、缓存和 checkpoint 的关系}

\texttt{state\_dict()} 的价值不只是“方便保存文件”。它把训练时需要恢复的所有重要状态显式列出来，避免你把“模型”误以为只是权重。对有些层来说，除了权重还有 running statistics、动量缓存或其他持久状态。若恢复时漏掉这些内容，模型行为就会和上次训练完全不是同一个轨迹。

\begin{table}[H]
\centering
\begin{tabular}{p{0.21\textwidth}p{0.27\textwidth}p{0.28\textwidth}}
\toprule
\textbf{对象} & \textbf{存什么} & \textbf{恢复时影响什么} \\
\midrule
Model state & 参数和模块缓存 & 前向输出、收敛路径 \\
Optimizer state & 动量、平方梯度等 & 更新方向、学习率适配 \\
Random seed & RNG 状态 & 数据顺序、dropout、初始化 \\
\bottomrule
\end{tabular}
\caption{checkpoint 不只是参数文件，而是一份训练状态快照}
\end{table}

\section{Training Loop and Best Practices}

\subsection{随机性与可复现}

训练里随机性来自很多地方：
\begin{itemize}
    \item 参数初始化
    \item dropout
    \item 数据顺序
\end{itemize}

因此调试时最好把三套随机种子一次性设好：
\begin{itemize}
    \item \texttt{torch.manual\_seed}
    \item \texttt{numpy.random.seed}
    \item \texttt{random.seed}
\end{itemize}

\begin{importantbox}{为什么要统一随机种子}
当实验结果不稳定时，最怕的是“同一个 bug 在不同 run 里表现不同”。先锁定随机性，才能把问题缩小到实现本身。
\end{importantbox}

\subsection{数据加载：把序列切成 batch}

语言模型的数据本质上是一串整数 token。最常见的做法是把它序列化到 numpy 数组，再用 `memmap` 懒加载。

\begin{lstlisting}
start_indices = torch.randint(len(data) - sequence_length, (batch_size,))
x = torch.tensor([data[start:start + sequence_length] for start in start_indices])
\end{lstlisting}

这段逻辑的语义很直接：
\begin{itemize}
    \item 随机采样若干起点；
    \item 从每个起点切出一段固定长度的 token 序列；
    \item 堆成一个 batch。
\end{itemize}

\begin{knowledgebox}{为什么要用 memmap}
真实数据集可能非常大，不能一次性全读进内存。`np.memmap` 允许只把当前访问到的部分映射到内存中，因此更适合大规模语言建模。
\end{knowledgebox}

\subsection{Pinned memory 和异步拷贝}

若 batch 先放在 pinned memory，再执行
\[
\texttt{x.to(device, non\_blocking=True)},
\]
CPU 到 GPU 的拷贝可以异步进行。这样就能把数据搬运和 GPU 计算重叠起来，减少空转。

\begin{warningbox}{数据管线也是性能瓶颈}
很多时候模型不慢，是因为数据喂不够快。只盯着算子优化而不看数据加载，常常会把训练吞吐的真正瓶颈漏掉。
\end{warningbox}

\subsection{优化器}

本讲先自己实现了两个最简单的优化器：
\begin{itemize}
    \item SGD
    \item AdaGrad
\end{itemize}

然后再把它们和更常见的优化器对应起来：
\begin{itemize}
    \item Momentum = SGD + 指数滑动平均
    \item AdaGrad = SGD + 累积平方梯度
    \item RMSProp = AdaGrad + 指数滑动平均的平方梯度
    \item Adam = RMSProp + Momentum
\end{itemize}

\begin{knowledgebox}{优化器的代价}
优化器不仅改变更新规则，还会显式增加状态量。比如 Adam 至少要保存两份额外状态，这就是为什么优化器内存常常比你想象得更贵。
\end{knowledgebox}

\subsection{Adam 的更新规则}

Adam 的核心是同时维护梯度的一阶矩估计（均值）和二阶矩估计（未中心化方差），并用偏差校正来补偿初始化时的偏差：

$$
m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \qquad \text{（一阶矩/动量）}
$$
$$
v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \qquad \text{（二阶矩/自适应学习率）}
$$
$$
\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t} \qquad \text{（偏差校正）}
$$
$$
\theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}
$$

其中 $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$ 是常见默认值。

\begin{knowledgebox}{AdamW vs Adam + L2 正则化——一个微妙但重要的区别}
经典 Adam + L2 正则化把权重衰减加在梯度上：
$$
g_t' = g_t + \lambda \theta_t, \qquad \theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t(g_t')}{\sqrt{\hat{v}_t(g_t')} + \epsilon}
$$
问题是：权重衰减项 $\lambda \theta_t$ 也被 Adam 的自适应学习率缩放了。对于梯度较大的参数，衰减效果被削弱；对于梯度较小的参数，衰减效果被放大。这并不是我们想要的行为。

AdamW（Loshchilov \& Hutter, 2019）把权重衰减解耦出来：
$$
\theta_{t+1} = (1 - \eta \lambda) \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}
$$
这样权重衰减不受自适应学习率的影响，对所有参数施加均匀的正则化。

\textbf{实践结论}：现代 LLM 训练几乎都使用 AdamW 而非 Adam + L2。PyTorch 中 \texttt{torch.optim.AdamW} 实现的就是解耦版本。
\end{knowledgebox}

\subsection{Optimizer 选择的实际含义}

在很小的模型上，SGD、AdaGrad、RMSProp 和 Adam 的差异可能看起来只是收敛速度不同；但在大模型上，它们还直接决定内存占用。一个 optimizer 的”更智能”，通常意味着更多状态、更复杂的更新逻辑，外加更高的调参成本。

\begin{table}[H]
\centering
\begin{tabular}{p{0.12\textwidth}p{0.15\textwidth}p{0.16\textwidth}p{0.16\textwidth}p{0.24\textwidth}}
\toprule
\textbf{优化器} & \textbf{额外状态} & \textbf{每参数内存} & \textbf{超参数} & \textbf{适用场景} \\
\midrule
SGD & 无 & 0 bytes & $\eta$ & 基线、教学 \\
SGD+Mom. & $m$ & 4 bytes (fp32) & $\eta, \beta$ & 计算机视觉 \\
AdaGrad & $\sum g^2$ & 4 bytes & $\eta, \epsilon$ & 稀疏特征 \\
RMSProp & $v$ & 4 bytes & $\eta, \beta_2, \epsilon$ & 非平稳目标 \\
Adam & $m, v$ & 8 bytes & $\eta, \beta_1, \beta_2, \epsilon$ & 通用 \\
AdamW & $m, v$ & 8 bytes & $\eta, \beta_1, \beta_2, \epsilon, \lambda$ & LLM 训练标配 \\
\bottomrule
\end{tabular}
\caption{优化器对比：状态数量直接决定内存成本}
\end{table}

\begin{warningbox}{AdaGrad 的致命缺陷}
AdaGrad 累积的是梯度平方和 $\sum_{i=1}^{t} g_i^2$，这个值只增不减。随着训练进行，自适应学习率 $\eta / \sqrt{\sum g^2 + \epsilon}$ 会单调下降趋近于零，最终导致训练停滞。RMSProp 和 Adam 通过指数滑动平均（而非累积求和）解决了这个问题。
\end{warningbox}

\subsection{训练时的内存核算}

对一个普通的前馈模型，训练时要同时考虑：

\begin{table}[H]
\centering
\begin{tabular}{lcc}
\toprule
\textbf{部分} & \textbf{数量级} & \textbf{说明} \\
\midrule
参数 & \(N_p\) & 模型权重 \\
激活值 & \(N_a\) & 前向过程中保存的中间结果 \\
梯度 & \(N_p\) & 每个参数对应一个梯度 \\
优化器状态 & \(N_p\) 或更多 & 取决于算法 \\
\bottomrule
\end{tabular}
\caption{训练步的常见内存组成}
\end{table}

若用 float32 朴素训练，则总内存可以粗略写成：
\[
4 \times (N_p + N_a + N_p + N_p) = 4(3N_p + N_a).
\]
这也是为什么“模型参数看起来装得下”，但“训练时还是爆显存”。

\subsection{训练循环}

标准训练循环可以压缩成四步：

\begin{enumerate}
    \item 前向：算预测值和 loss
    \item 反向：`loss.backward()`
    \item 更新：`optimizer.step()`
    \item 清零：\texttt{optimizer.zero\_grad(set\_to\_none=True)}
\end{enumerate}

\begin{lstlisting}
for t in range(num_train_steps):
    x, y = get_batch(B=B)
    pred_y = model(x)
    loss = F.mse_loss(pred_y, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
\end{lstlisting}

\begin{importantbox}{训练循环的本质}
训练循环不是”写起来麻烦”，而是”把前向、反向、更新三种状态转换显式地串起来”。一旦逻辑清楚了，调参和定位 bug 都会容易很多。
\end{importantbox}

\subsection{梯度累积：用小 batch 模拟大 batch}

当 GPU 显存不足以容纳理想的 batch size 时，\textbf{梯度累积}（gradient accumulation）是最简单的解决方案。其原理是：连续做多次前向+反向但不执行 \texttt{optimizer.step()}，让梯度在多个 micro-batch 上累积，最后一次性更新参数。

\begin{lstlisting}[caption={梯度累积的标准实现}]
accumulation_steps = 4  # effective batch = micro_batch * 4

for step in range(num_train_steps):
    for micro_step in range(accumulation_steps):
        x, y = get_batch(B=micro_batch_size)
        pred = model(x)
        loss = loss_fn(pred, y) / accumulation_steps  # scale loss
        loss.backward()  # gradients accumulate

    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
\end{lstlisting}

\begin{warningbox}{梯度累积的两个常见错误}
\begin{itemize}
    \item \textbf{忘记缩放 loss}：如果不除以 \texttt{accumulation\_steps}，累积的梯度会比真正大 batch 的梯度大 $k$ 倍（$k$ 为累积步数），等效于学习率放大了 $k$ 倍。
    \item \textbf{与 BatchNorm 的冲突}：BatchNorm 的统计量是在每个 micro-batch 上独立计算的，和真正的大 batch 行为不同。LLM 通常使用 LayerNorm，不受此影响。
\end{itemize}
\end{warningbox}

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.14\textwidth}p{0.14\textwidth}p{0.14\textwidth}p{0.16\textwidth}}
\toprule
\textbf{方法} & \textbf{Micro BS} & \textbf{累积步数} & \textbf{有效 BS} & \textbf{显存占用} \\
\midrule
直接大 batch & 64 & 1 & 64 & 极高 \\
梯度累积 & 16 & 4 & 64 & 中等 \\
梯度累积 & 4 & 16 & 64 & 较低 \\
梯度累积 & 1 & 64 & 64 & 最低 \\
\bottomrule
\end{tabular}
\caption{梯度累积在数学上等价于大 batch，但显存和吞吐有差异}
\end{table}

\subsection{学习率调度}

现代 LLM 训练几乎都使用学习率调度器（learning rate scheduler），而不是固定学习率。最常见的策略是 \textbf{warmup + cosine decay}：

\begin{enumerate}
    \item \textbf{Warmup 阶段}：从接近 0 的学习率线性增长到峰值学习率，通常持续总步数的 1\%--5\%。
    \item \textbf{Cosine decay 阶段}：学习率按余弦函数从峰值衰减到一个很小的最终值（通常为峰值的 1/10 或更小）。
\end{enumerate}

$$
\eta_t = \begin{cases}
\eta_{\max} \cdot \frac{t}{T_{\text{warmup}}} & \text{if } t < T_{\text{warmup}} \\[6pt]
\eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{t - T_{\text{warmup}}}{T_{\text{total}} - T_{\text{warmup}}} \cdot \pi\right)\right) & \text{otherwise}
\end{cases}
$$

\begin{knowledgebox}{为什么需要 warmup}
训练初期，Adam 的二阶矩 $v_t$ 尚未充分预热（bias correction 后的估计仍不稳定），此时使用大学习率容易导致梯度爆炸或训练发散。Warmup 让模型在参数空间中先”小步试探”，等优化器状态稳定后再加大步长。

经验上，warmup 步数通常设为总步数的 1\%--2\%，或者 2000 步左右。
\end{knowledgebox}

\subsection{梯度裁剪}

梯度裁剪（gradient clipping）是防止训练不稳定的重要工具。最常见的做法是\textbf{全局梯度范数裁剪}：

$$
\hat{g} = \begin{cases}
g & \text{if } \|g\| \leq C \\
C \cdot \frac{g}{\|g\|} & \text{if } \|g\| > C
\end{cases}
$$

其中 $C$ 是裁剪阈值（通常设为 1.0），$\|g\|$ 是所有参数梯度拼接后的 L2 范数。

\begin{lstlisting}[caption={梯度裁剪的标准用法}]
loss.backward()
# Clip gradients before optimizer step
grad_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(), max_norm=1.0
)
optimizer.step()
optimizer.zero_grad(set_to_none=True)

# Monitor grad_norm for training stability
if grad_norm > 10.0:
    print(f”Warning: large gradient norm {grad_norm:.2f}”)
\end{lstlisting}

\begin{importantbox}{梯度范数是训练健康的晴雨表}
监控梯度范数可以帮助诊断训练问题：
\begin{itemize}
    \item 梯度范数持续增大 $\to$ 可能训练发散，需要降低学习率
    \item 梯度范数突然飙升 $\to$ 可能遇到异常数据或数值不稳定
    \item 梯度范数趋近于零 $\to$ 可能学习率太小或梯度消失
    \item 梯度范数频繁被裁剪 $\to$ 裁剪阈值可能设得太低
\end{itemize}
\end{importantbox}

\subsection{训练循环的状态机}

训练步可以看成一个状态机的四个阶段：\textbf{Batch ready $\to$ Forward $\to$ Backward $\to$ Optimizer step}，然后循环。真正出问题的地方通常不是其中某一步本身，而是状态切换时忘了清空梯度、忘了保存 optimizer state，或者把数据放在了错误的 device 上。

\begin{table}[H]
\centering
\begin{tabular}{p{0.18\textwidth}p{0.28\textwidth}p{0.24\textwidth}p{0.2\textwidth}}
\toprule
\textbf{状态} & \textbf{输入} & \textbf{输出} & \textbf{常见 bug} \\
\midrule
Batch ready & tokenizer / data loader & batch tensor & 放错 device、shape 对不上 \\
Forward & batch + parameters & prediction / loss & dtype 不匹配、mask 错误 \\
    Backward & loss & gradients & 忘记 \texttt{requires\_grad}、梯度爆炸 \\
    Step & gradients + optimizer state & 更新后的 parameters & 忘记 \texttt{zero\_grad}、checkpoint 丢 optimizer \\
\bottomrule
\end{tabular}
\caption{训练循环里每个状态都对应一类典型 bug}
\end{table}

\subsection{训练监控与日志}

大规模训练中，监控指标是及早发现问题的唯一手段。一个完整的训练监控系统通常需要记录以下信息：

\begin{table}[H]
\centering
\begin{tabular}{p{0.20\textwidth}p{0.25\textwidth}p{0.35\textwidth}}
\toprule
\textbf{指标} & \textbf{正常行为} & \textbf{异常信号} \\
\midrule
Training loss & 单调下降（整体趋势） & 突然跳升、停滞不下、NaN \\
Learning rate & 按调度器变化 & 始终为 0 或异常大 \\
梯度范数 & 稳定在某个范围内 & 持续增大、突然飙升 \\
GPU 利用率 & 接近 100\% & 大幅波动或持续低于 50\% \\
显存占用 & 稳定不变 & 随步数缓慢增长（泄漏） \\
吞吐量 (tokens/s) & 稳定 & 周期性下降（可能是 checkpoint I/O） \\
\bottomrule
\end{tabular}
\caption{训练监控的核心指标与异常诊断}
\end{table}

\begin{lstlisting}[caption={训练循环中的日志记录示例}]
import wandb  # or tensorboard

wandb.init(project="cs336-lecture02")

for step in range(num_train_steps):
    x, y = get_batch(B=B)
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()

    grad_norm = torch.nn.utils.clip_grad_norm_(
        model.parameters(), max_norm=1.0
    )

    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    if step % log_interval == 0:
        wandb.log({
            "loss": loss.item(),
            "grad_norm": grad_norm.item(),
            "lr": optimizer.param_groups{[}0{]}{[}'lr'{]},
            "gpu_mem_gb": torch.cuda.memory_allocated() / 1e9,
        }, step=step)
\end{lstlisting}

\begin{knowledgebox}{Loss spike 的常见原因}
训练过程中偶尔出现的 loss spike（突然跳升）可能由以下原因引起：
\begin{itemize}
    \item \textbf{异常数据}：某个 batch 包含异常长的序列或罕见 token 组合
    \item \textbf{学习率过大}：尤其在 warmup 结束附近
    \item \textbf{数值不稳定}：梯度中出现 inf 或 nan
    \item \textbf{数据重复}：训练数据中存在大量重复导致过拟合后跳出
\end{itemize}
轻微的 loss spike 通常可以自行恢复，但如果 loss 持续不下降或出现 NaN，通常需要从上一个 checkpoint 回滚并降低学习率。
\end{knowledgebox}

\subsection{Checkpointing}

训练语言模型往往很久，而且很容易中断（硬件故障、抢占、OOM）。所以要定期保存 checkpoint。

\begin{lstlisting}[caption={保存 checkpoint 的标准写法}]
checkpoint = {
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'step': step,
    'loss': loss.item(),
    'rng_state': torch.random.get_rng_state(),
}
torch.save(checkpoint, f'checkpoint_step_{step}.pt')
\end{lstlisting}

\begin{lstlisting}[caption={恢复 checkpoint 的标准写法}]
checkpoint = torch.load(f'checkpoint_step_{step}.pt')
model.load_state_dict(checkpoint{[}'model_state_dict'{]})
optimizer.load_state_dict(checkpoint{[}'optimizer_state_dict'{]})
torch.random.set_rng_state(checkpoint{[}'rng_state'{]})
start_step = checkpoint{[}'step'{]}
\end{lstlisting}

\begin{warningbox}{只存模型不够}
如果只存 model，不存 optimizer，那么恢复后学习率状态、动量累积、平方梯度缓存都没了，训练行为会和之前不一样。尤其是 Adam 系列优化器，$m_t$ 和 $v_t$ 需要大量 step 才能预热到稳定状态，丢失后相当于重新开始预热。
\end{warningbox}

\subsubsection{Checkpoint 频率与存储成本}

Checkpoint 的频率需要在”恢复代价”和”存储成本”之间权衡：

\begin{table}[H]
\centering
\begin{tabular}{p{0.18\textwidth}p{0.28\textwidth}p{0.35\textwidth}}
\toprule
\textbf{策略} & \textbf{频率} & \textbf{权衡} \\
\midrule
保守 & 每 1000 步 & 存储开销小，但中断后最多浪费 1000 步 \\
激进 & 每 100 步 & 恢复损失小，但存储快速增长 \\
混合 & 近期密集 + 远期稀疏 & 保留最近 5 个 + 每 N 步一个永久 \\
\bottomrule
\end{tabular}
\caption{checkpoint 频率策略的常见选择}
\end{table}

以 7B 模型为例，一个完整 checkpoint（模型 + 优化器 + RNG）大约需要：
$$
\underbrace{7 \times 10^9 \times 2}_{\text{bf16 参数}} + \underbrace{7 \times 10^9 \times (4+4+4)}_{\text{fp32 主权重 + } m + v} \approx 14 + 84 = 98 \text{ GB}
$$

如果每 100 步存一个、保留最近 10 个，光 checkpoint 就需要约 1 TB 存储。

\subsubsection{Checkpoint 的异步保存}

大规模训练中，同步保存 checkpoint 会阻塞训练。常见做法是把 checkpoint 数据先拷贝到 CPU 内存（或 pinned memory），然后在后台线程中异步写入磁盘，同时 GPU 继续下一步训练。

\begin{knowledgebox}{分布式 checkpoint}
当模型被切分到多张 GPU 上（模型并行/数据并行）时，每张卡只保存自己的那部分状态。恢复时需要确保切分方式完全一致，否则状态会对不上。PyTorch 的 \texttt{DistributedStateDictSaveLoad} 和 DeepSpeed 的 checkpoint 引擎都提供了自动化的分布式 checkpoint 管理。
\end{knowledgebox}

\subsubsection{Checkpoint 格式选择}

不同的 checkpoint 格式有不同的权衡：

\begin{table}[H]
\centering
\begin{tabular}{p{0.17\textwidth}p{0.22\textwidth}p{0.22\textwidth}p{0.22\textwidth}}
\toprule
\textbf{格式} & \textbf{优点} & \textbf{缺点} & \textbf{适用场景} \\
\midrule
\texttt{torch.save} (pickle) & PyTorch 原生、支持任意对象 & 安全风险（pickle 反序列化）、不支持部分加载 & 训练中间 checkpoint \\
\texttt{safetensors} & 安全、支持 mmap 部分加载、跨框架 & 只能存 tensor 字典 & 模型发布与共享 \\
\texttt{torch.distributed .checkpoint} & 支持分布式保存/加载、resharding & API 较新、生态不够成熟 & 大规模分布式训练 \\
\bottomrule
\end{tabular}
\caption{常见 checkpoint 格式的对比}
\end{table}

\begin{warningbox}{pickle 的安全风险}
\texttt{torch.save} 使用 Python 的 pickle 协议，这意味着加载一个不受信任的 checkpoint 文件可能执行任意代码。在加载来自互联网的模型权重时，优先选择 \texttt{safetensors} 格式，或者使用 \texttt{torch.load(..., weights\_only=True)} 来限制反序列化行为。
\end{warningbox}

\subsubsection{长训练的 Checkpoint 最佳实践}

对于持续数周甚至数月的大规模训练，checkpoint 策略需要更加系统化：

\begin{enumerate}
    \item \textbf{滚动保留}：保留最近 $k$ 个 checkpoint（如最近 5 个），自动删除更早的，防止存储爆炸。
    \item \textbf{里程碑永久保存}：在关键节点（如每 10\% 训练进度）永久保存一个 checkpoint，用于后续分析。
    \item \textbf{保存训练指标}：除了模型和优化器状态，还要保存 loss 曲线、学习率、梯度范数等指标，方便恢复后验证一致性。
    \item \textbf{预留时间}：在集群抢占或定时任务到期前，提前触发一次 checkpoint（如收到 SIGTERM 信号时）。
    \item \textbf{校验完整性}：保存后立即验证 checkpoint 文件的完整性（如 checksum），避免写入中断导致的损坏。
\end{enumerate}

\begin{knowledgebox}{Checkpoint 与训练的 wall-clock 开销}
以 7B 模型为例，一个完整 checkpoint 约 98 GB。在 NVMe SSD（写入 3 GB/s）上，同步保存需要约 33 秒。如果每 100 步存一次，且每步约 2 秒，那么 checkpoint 开销约为 $33 / (100 \times 2) \approx 16\%$。使用异步保存可以将这个开销降到接近零。
\end{knowledgebox}

\subsection{Checkpoint 调试清单}

Checkpoint 往往不是“能不能存”这么简单，而是“恢复后是不是同一个实验”。如果恢复后的 loss 曲线明显不一致，通常优先查下面几项：

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.28\textwidth}p{0.26\textwidth}}
\toprule
\textbf{现象} & \textbf{优先检查} & \textbf{常见原因} \\
\midrule
恢复后 loss 跳变 & optimizer state、学习率调度 & 只存了 model 没存 optimizer \\
训练变慢 & device / dtype / data loader & batch 没放到 GPU 或 pinned memory \\
梯度不更新 & requires\_grad / zero\_grad & 参数被冻住或梯度没清 \\
结果不可复现 & RNG、数据顺序、dropout & 随机种子没统一 \\
\bottomrule
\end{tabular}
\caption{checkpoint 的问题通常要从训练状态而不是单纯参数开始排查}
\end{table}

\subsection{混合精度训练}

\subsubsection{混合精度的核心思想}

混合精度训练（Mixed Precision Training）的核心思想是：\textbf{用低精度做计算密集的部分（前向、反向的矩阵乘法），用高精度做数值敏感的部分（参数更新、loss 计算）}。

一个典型的混合精度训练流程：
\begin{enumerate}
    \item 维护一份 \textbf{fp32 主权重}（master weights）
    \item 每步开始时，将主权重转换为 bf16 副本
    \item 用 bf16 做前向和反向传播（矩阵乘法在 Tensor Core 上加速）
    \item 得到 bf16 梯度后，转换回 fp32
    \item 在 fp32 精度下更新主权重
\end{enumerate}

\begin{importantbox}{混合精度的三重好处}
\begin{enumerate}
    \item \textbf{计算加速}：bf16 矩阵乘法在 Tensor Core 上比 fp32 快约 $8\times$--$16\times$
    \item \textbf{内存节省}：激活值用 bf16 存储，内存减半
    \item \textbf{带宽节省}：bf16 数据传输量减半，减少 memory bandwidth 瓶颈
\end{enumerate}
\end{importantbox}

\subsubsection{PyTorch AMP 实现}

PyTorch 的 Automatic Mixed Precision (AMP) 通过 \texttt{torch.cuda.amp} 提供了自动化的混合精度支持：

\begin{lstlisting}[caption={使用 PyTorch AMP 的训练循环}]
scaler = torch.amp.GradScaler()  # only needed for fp16

for x, y in dataloader:
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        pred = model(x)
        loss = loss_fn(pred, y)

    # bf16 does not need scaler, but fp16 does
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
\end{lstlisting}

\texttt{autocast} 上下文管理器会自动决定哪些操作用低精度、哪些保留 fp32。一般来说：
\begin{itemize}
    \item \textbf{用 bf16/fp16}：矩阵乘法、卷积、线性层
    \item \textbf{保留 fp32}：softmax、layernorm、loss 计算、累加操作
\end{itemize}

\begin{warningbox}{Loss Scaling：fp16 的必要补丁}
使用 fp16 时，由于动态范围只有 $10^{\pm 5}$，小梯度容易 underflow 为零。\textbf{Loss scaling} 通过在反向传播前将 loss 乘以一个大数（如 $2^{16}$），使梯度值被放大到 fp16 的可表示范围内，然后在参数更新前除回来。

\texttt{GradScaler} 还会动态调整 scale 因子：如果发现梯度中有 inf/nan（overflow），就减小 scale 并跳过该步更新。

\textbf{使用 bf16 时不需要 loss scaling}，因为 bf16 的动态范围和 fp32 一样大。这是 bf16 在实践中更受欢迎的另一个原因。
\end{warningbox}

\subsubsection{混合精度下的完整内存分析}

以 $N$ 参数模型为例，对比纯 fp32 训练和 bf16 混合精度训练的内存：

\begin{table}[H]
\centering
\begin{tabular}{p{0.28\textwidth}p{0.16\textwidth}p{0.16\textwidth}p{0.22\textwidth}}
\toprule
\textbf{组件} & \textbf{纯 FP32} & \textbf{混合精度} & \textbf{说明} \\
\midrule
模型参数 & $4N$ & $2N$ (bf16) & 前向用低精度 \\
梯度 & $4N$ & $2N$ (bf16) & 反向产生低精度梯度 \\
主权重 & --- & $4N$ (fp32) & 混合精度额外持有 \\
Adam $m_t$ & $4N$ & $4N$ & 始终 fp32 \\
Adam $v_t$ & $4N$ & $4N$ & 始终 fp32 \\
\midrule
\textbf{静态总计} & $\mathbf{16N}$ & $\mathbf{16N}$ & 静态内存相当 \\
激活值 & $4 \cdot A$ & $2 \cdot A$ & \textbf{激活减半是关键} \\
\bottomrule
\end{tabular}
\caption{混合精度的主要内存收益来自激活值减半和计算加速}
\end{table}

关键洞察：混合精度在静态内存（参数+优化器）上并没有显著节省，因为需要额外维护 fp32 主权重。\textbf{真正的收益在于}：
\begin{enumerate}
    \item 激活值内存减半（对于大 batch size 和长序列影响巨大）
    \item Tensor Core 的计算加速（$8\times$--$16\times$）
    \item 内存带宽减半
\end{enumerate}

\subsubsection{FP8 训练：下一代低精度前沿}

FP8 是 H100 及更新 GPU 支持的 8-bit 浮点格式。它有两种变体，各自针对不同的使用场景：

\begin{table}[H]
\centering
\begin{tabular}{p{0.12\textwidth}p{0.12\textwidth}p{0.14\textwidth}p{0.14\textwidth}p{0.30\textwidth}}
\toprule
\textbf{格式} & \textbf{指数位} & \textbf{尾数位} & \textbf{动态范围} & \textbf{设计目的} \\
\midrule
E4M3 & 4 & 3 & $\sim 10^{\pm 2}$ & 前向权重和激活（需要更高精度） \\
E5M2 & 5 & 2 & $\sim 10^{\pm 5}$ & 反向梯度（需要更大范围） \\
\bottomrule
\end{tabular}
\caption{两种 FP8 变体的设计权衡}
\end{table}

FP8 训练的关键挑战是动态范围非常有限，因此需要 \textbf{per-tensor scaling}：为每个 tensor 维护一个缩放因子，使其值域映射到 FP8 的可表示范围内。这增加了实现复杂度，但在 H100 上可以获得相比 BF16 约 $2\times$ 的额外加速。

\begin{warningbox}{FP8 训练的成熟度}
截至 2025 年，FP8 训练仍处于快速发展期。主要挑战包括：
\begin{itemize}
    \item 缩放因子的选择策略（延迟缩放 vs 即时缩放）
    \item 某些层（如 attention softmax、LayerNorm）仍需更高精度
    \item 不同框架（NVIDIA TransformerEngine、MS-AMP）的实现差异
    \item 训练稳定性在超大规模模型上的验证仍不充分
\end{itemize}
FP8 在推理中的应用已相当成熟，但在训练中仍需谨慎评估。
\end{warningbox}

\subsection{Activation Checkpointing（梯度检查点）}

\begin{importantbox}{Activation Checkpointing 的核心思想}
反向传播需要前向过程中保存的激活值来计算梯度。对于深层网络，所有层的激活值同时驻留在 GPU 内存中，内存开销巨大。

\textbf{Activation checkpointing}（也叫 gradient checkpointing 或 rematerialization）的策略是：只保存部分层的激活值（称为 checkpoint），丢弃其余层的激活值。反向传播到某一层时，从最近的 checkpoint 重新做前向计算来恢复所需的激活值。

这是一个经典的\textbf{时间换空间}权衡。
\end{importantbox}

假设模型有 $L$ 层，每层激活值占 $a$ 内存：

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.22\textwidth}p{0.22\textwidth}p{0.18\textwidth}}
\toprule
\textbf{策略} & \textbf{激活值内存} & \textbf{额外计算} & \textbf{适用场景} \\
\midrule
不做 checkpointing & $L \cdot a$ & 0 & 内存充足 \\
每 $\sqrt{L}$ 层 checkpoint & $\sqrt{L} \cdot a$ & $\sim 33\%$ 前向 & 经典权衡 \\
每层 checkpoint & $a$ (仅 1 层) & $\sim 100\%$ 前向 & 极端内存受限 \\
\bottomrule
\end{tabular}
\caption{Activation checkpointing 的不同策略}
\end{table}

PyTorch 提供了 \texttt{torch.utils.checkpoint.checkpoint} 函数来实现这一功能：

\begin{lstlisting}[caption={使用 activation checkpointing 的示例}]
from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        # ... transformer block logic ...
        return output

class Model(nn.Module):
    def __init__(self, num_layers):
        super().__init__()
        self.layers = nn.ModuleList(
            {[}TransformerBlock() for _ in range(num_layers){]}
        )

    def forward(self, x):
        for layer in self.layers:
            # Wrap each layer with checkpointing
            x = checkpoint(layer, x, use_reentrant=False)
        return x
\end{lstlisting}

\begin{warningbox}{Activation Checkpointing 的陷阱}
\begin{itemize}
    \item \textbf{计算开销}：每个被 checkpoint 的段在反向时会重新执行前向计算，总计算量增加约 33\%（如果每 $\sqrt{L}$ 层 checkpoint）。
    \item \textbf{不可 checkpoint 的操作}：含有副作用（如 dropout 的随机状态）的层需要特殊处理，否则重新计算时行为会不一致。
    \item \textbf{与 AMP 的交互}：在混合精度训练中使用 checkpointing 时，需要确保 autocast 上下文在 checkpoint 段内正确传播。
\end{itemize}
\end{warningbox}

\begin{knowledgebox}{Activation Checkpointing 的内存收益实例}
以 70B 模型（约 80 层 Transformer）为例，假设每层激活值约占 200 MB（取决于 batch size 和序列长度）：
\begin{itemize}
    \item 不做 checkpointing：$80 \times 200 = 16,000$ MB $= 16$ GB
    \item 每 $\sqrt{80} \approx 9$ 层 checkpoint：$9 \times 200 = 1,800$ MB $\approx 1.8$ GB
    \item 代价：约 33\% 额外前向计算
\end{itemize}
这意味着可以节省约 14 GB 激活值内存，代价是每步训练时间增加约 33\%。在 GPU 内存紧张时，这个权衡几乎总是值得的。
\end{knowledgebox}

\subsection{混合精度 + Activation Checkpointing 的协同}

在实际大规模训练中，混合精度和 activation checkpointing 通常同时使用，形成一套完整的内存优化方案：

\begin{table}[H]
\centering
\begin{tabular}{p{0.30\textwidth}p{0.20\textwidth}p{0.20\textwidth}p{0.15\textwidth}}
\toprule
\textbf{优化技术} & \textbf{内存节省} & \textbf{计算开销} & \textbf{实现难度} \\
\midrule
混合精度 (bf16) & 激活减半 & 无（反而加速） & 低 \\
Activation checkpointing & 激活减 $\sqrt{L}\times$ & +33\% 前向 & 中 \\
梯度累积 & 等效大 batch & 无额外开销 & 低 \\
模型并行 & 线性切分 & 通信开销 & 高 \\
\bottomrule
\end{tabular}
\caption{常见内存优化技术的对比}
\end{table}

\section{总结与延伸}

\subsection{本讲知识体系}

本讲把”训练一个模型”拆成了可操作的最小单元：
\begin{itemize}
    \item 张量如何占内存——从 dtype、numel 到 GPU 内存层次；
    \item 矩阵乘法如何计算 FLOPs——从 $2BDK$ 到 $6ND$ 规则；
    \item 参数如何封装进 \texttt{nn.Module}——从 \texttt{nn.Parameter} 到 \texttt{state\_dict}；
    \item 数据如何切 batch——memmap、pinned memory、异步拷贝；
    \item 优化器、训练循环、checkpoint 和混合精度如何协同工作。
\end{itemize}

\subsection{核心公式速查}

\begin{table}[H]
\centering
\begin{tabular}{p{0.28\textwidth}p{0.32\textwidth}p{0.28\textwidth}}
\toprule
\textbf{场景} & \textbf{公式} & \textbf{用途} \\
\midrule
张量内存 & $\text{numel} \times \text{element\_size}$ & 单个 tensor 内存估算 \\
矩阵乘法 FLOPs & $2BDK$ & 单次 matmul 计算量 \\
训练总 FLOPs & $6ND$ & 完整训练计算预算 \\
训练时间 & $C / (\text{GPUs} \times \eta \times \text{Peak})$ & Wall-clock 估算 \\
AdamW 静态内存 & $16N$ bytes & 混合精度下每参数 \\
Transformer 每层参数 & $\approx 12d^2$ & 模型规模速算 \\
Activation checkpointing & $\sqrt{L} \cdot a$ 内存，+33\% 计算 & 内存优化评估 \\
\bottomrule
\end{tabular}
\caption{本讲核心公式速查表}
\end{table}

\subsection{从本讲到后续课程的衔接}

\begin{table}[H]
\centering
\begin{tabular}{p{0.22\textwidth}p{0.32\textwidth}p{0.30\textwidth}}
\toprule
\textbf{本讲内容} & \textbf{后续延伸} & \textbf{对应课程} \\
\midrule
内存核算 & 模型并行、张量并行、流水线并行 & Lecture 5 (GPUs) \\
FLOPs 估算 & Scaling laws、compute-optimal 训练 & Lecture 4 (Scaling) \\
\texttt{nn.Module} & 完整 Transformer 实现 & Assignment 1 \\
混合精度 & 分布式混合精度、FSDP & Lecture 5 (GPUs) \\
Checkpoint & 分布式 checkpoint、容错训练 & 高级话题 \\
\bottomrule
\end{tabular}
\caption{本讲内容在课程体系中的位置}
\end{table}

\begin{importantbox}{最后要记住的一句话}
深度学习系统不是”把模型写出来”就结束了，而是要持续回答三个问题：
\begin{enumerate}
    \item 这个东西占多少内存？
    \item 这个东西消耗多少 FLOPs？
    \item 这个东西能不能稳定地、可复现地训练起来？
\end{enumerate}
\end{importantbox}

做完 Assignment 1 之后，这些概念会变得更具体，也会更牢固，因为你会亲手把它们用在真正的 Transformer 上。

\subsection{本章小结}
这一讲把 PyTorch 中真正决定训练是否可行的三件事放到了一起：参数与状态的内存核算、前向反向的 FLOPs 规模，以及 mixed precision / checkpoint 等工程优化。真正有价值的 takeaway 不是某个公式，而是把“模型实现”提升成“系统预算”的习惯。

\subsection{拓展阅读}

\begin{itemize}
    \item \textbf{混合精度训练原始论文}：Micikevicius et al., ``Mixed Precision Training'' (2018). 提出了 loss scaling 和 master weights 的混合精度框架。
    \item \textbf{AdamW 论文}：Loshchilov \& Hutter, ``Decoupled Weight Decay Regularization'' (2019). 阐明了解耦权重衰减与 L2 正则化的区别。
    \item \textbf{Activation Checkpointing 原始论文}：Chen et al., ``Training Deep Nets with Sublinear Memory Cost'' (2016). 提出了 $\sqrt{L}$ 策略的理论分析。
    \item \textbf{FP8 训练}：Micikevicius et al., ``FP8 Formats for Deep Learning'' (2022). NVIDIA 提出的 E4M3/E5M2 FP8 标准。
    \item \textbf{PyTorch 内部机制}：Edward Z. Yang 的博客 ``PyTorch internals''. 详细讲解了 tensor storage、autograd 和 dispatch 机制。
    \item \textbf{Scaling Laws}：Kaplan et al., ``Scaling Laws for Neural Language Models'' (2020). 建立了参数量、数据量、计算量与 loss 之间的幂律关系。
    \item \textbf{Chinchilla}：Hoffmann et al., ``Training Compute-Optimal Large Language Models'' (2022). 给出了更精确的 compute-optimal 训练配方。
\end{itemize}

\end{document}
