CS336 Lecture Notes 5

Warning

本文为本人学习相关开源课程过程中，整理的个人学习笔记及作业解答，核心目的仅用于记录个人学习轨迹、巩固所学知识、梳理学习思路，全程为个人自主学习使用，不具备任何商业用途，也不构成任何形式的课程辅导或标准答案参考。

需特别说明的是，由于本人学习进度及知识储备有限，笔记内容及作业解答中可能存在大量纰漏、思路偏差甚至错误，仅代表本人当时的学习理解，不具备权威性和准确性。

在此郑重提醒：请勿将本文中的任何作业解答复制粘贴，作为自身所修课程的提交答案。任何因抄袭本文内容导致的课程成绩问题、学术诚信问题，均由抄袭者自行承担全部责任，本人不承担任何相关连带责任。

同时，本文所分享的内容均基于开源课程的公开内容整理，尊重原课程创作者的知识产权，若涉及相关内容的版权问题，请及时联系本人，本人将第一时间进行调整或删除。

感谢各位读者的理解与支持，也欢迎大家针对笔记及解答中的问题提出宝贵建议，共同交流学习、共同进步。

Important

课程网站: https://cs336.stanford.edu/
Lec0_ 资料: lecture_0x.py
Lec0_ 资料: lecture_0x.pdf

Scaling laws (Part I)#

动机：为什么要研究 Scaling Laws？#

想象这样一个场景：朋友给你一万张 B200 GPU 一个月的使用权，让你构建一个优秀的开源语言模型。你会怎么做？

搭建基础设施团队和分布式训练框架
准备高质量的预训练数据集
运行一个大模型（但选哪个模型？）

核心问题：Wide or deep? How many heads? Which nonlinearity?

Scaling Laws 提供了一种简单、可预测的规则来理解语言模型的行为。

传统方法：在大模型上调参（昂贵） 新方法：在小模型上调参，然后外推到大模型

Part 1: 数据扩展定律 (Data Scaling Laws)#

数据与性能的关系#

核心发现：Loss 和数据集大小在 log-log 图上呈线性关系，即”幂律”关系。

从 Kaplan+ 2020 的研究可以看出，语言模型的测试损失与数据集大小呈幂律关系： $\text{Loss} \propto n^{-\alpha}$

理论基础：为什么会出现幂律？#

简单示例：均值估计

输入： $x_1, \ldots, x_n \sim \mathcal{N}(\mu, \sigma^2)$ ，估计均值 $\hat{\mu} = \frac{1}{n}\sum_i x_i$

误差： $\mathbb{E}[(\hat{\mu} - \mu)^2] = \frac{\sigma^2}{n}$

取对数后： $\log(\text{Error}) = -\log n + 2\log\sigma$

这是一个标准的 scaling law 形式。

非参数学习的扩展定律

在 $d$ 维空间中，非参数学习的误差为 $\text{Error} = n^{-1/d}$ ，这意味着斜率与数据的内在维度相关。

内在维度理论 (Bahri 2021)：scaling law 的斜率 $\alpha$ 与数据的内在维度密切相关。

数据重复的影响#

实际场景中数据是有限的，重复数据如何影响扩展定律？

$D'$ = 有效数据量
$U_d$ = 唯一 token 数量
$R_d$ = 重复次数

重复数据会降低有效数据量的价值。

数据选择与分布偏移#

数据组成如何影响性能？数据组成影响的是 offset（偏移量），而不是 slope（斜率）。这些”分布偏移”扩展定律可以帮助我们理解收集多样化数据的重要性。

Part 2: 模型工程中的扩展定律#

1. 架构选择：Transformer vs LSTM#

通过在小模型上建立 scaling law，可以在不花费巨额成本的情况下比较不同架构。

2. 优化器选择：Adam vs SGD#

Hestness+ 2017 的研究表明，不同优化器遵循不同的扩展曲线。

3. 深度与宽度#

1层 vs 2层：差异巨大
更多层数在 $10^7$ 参数以下收益递减
超参数（如 aspect ratio）是否依赖于规模？

重要发现：并非所有参数都是等价的！Embedding 层参数的行为与其他参数不同。

4. 批量大小：Critical Batch Size#

批量大小超过某一点后会有强烈的边际收益递减。

Critical Batch Size 定义：

选择目标损失，测量所需的步数 $S$ 和样本数 $E$
扫描不同批量大小
曲线遵循： $B_{\text{crit}} = \frac{E_{\min}}{S_{\min}}$

关键发现：损失目标越小，临界批量越大。

5. 学习率：muP 和规模感知的 LR 选择#

如果简单地扩展模型，最优学习率会依赖于规模。需要使用规模感知的初始化和学习率缩放策略。

Part 3: 模型与数据的联合扩展定律#

核心问题：更多数据还是更大模型？#

Rosenfeld+ 2020： $\text{Error} = n^{-\alpha} + m^{-\beta} + C$

Kaplan+ 2020： $\text{Error} = m^{-\alpha} + n^{-1/\beta}$

这些公式能够很好地拟合模型-数据联合误差。

Chinchilla vs Kaplan：计算最优权衡#

Kaplan 结论： $N_{\text{opt}} = C^{0.73}$ , $D_{\text{opt}} = C^{0.27}$ （每参数 token 数随计算量增加而减少）

Chinchilla 方法：

Method 1：最小值法 - 所有训练曲线的最小值形成幂律
Method 2：IsoFLOPS - 选择一系列 FLOP 预算，变化参数数量，取凸形状的最小值
Method 3：联合拟合 - 在 size-data 网格上运行多个模型

差异原因：

Kaplan 从参数计数中移除了最后一层参数
在非常小的计算预算上，warmup 过高
非 embedding 参数与总参数选择的影响

重要结论：训练最优 ≠ 实际最优#

Chinchilla 告诉你在固定训练计算量下如何获得最佳模型，但实际部署中大部分计算是推理。

实际训练中的 tokens/param 比例：

模型	Tokens/Param
GPT-3	2
Chinchilla	20
LLaMA 65B	22
Llama 2 70B	29
Mistral 7B	110
Llama 3 70B	215

结论：预期使用越多，越值得支付前期过度训练的成本。

Scaling Law 设计流程#

训练几个较小的模型
建立 scaling law（如 Adam vs SGD 扩展定律）
根据 scaling law 预测选择最优超参数

总结#

类型	核心发现	应用
数据扩展	Log-linear 关系，有清晰的理论基础	数据收集/筛选
模型扩展	大幅降低训练成本	架构、优化器、超参数选择
联合扩展	计算与数据/模型权衡	资源分配决策

关键洞察：大模型上超参数的影响可以在训练前通过小模型预测！

Inference#

推理工作负载概述#

推理场景：

实际使用（聊天机器人、代码补全、Agent、批处理）
模型评估（指令遵循等）
强化学习（采样多个生成，然后评分）

效率的重要性：训练是一次性成本，推理重复多次

OpenAI 每天处理约 8.6T tokens
DeepSeek v4 训练用了 32T tokens

性能指标：

指标	定义	适用场景
TTFT (Time-to-first-token)	用户等待第一个 token 的时间	交互式应用
Latency (seconds/token)	单个查询的 token 生成速度	交互式应用
Throughput (tokens/second)	多个查询的 token 生成速度	批处理

开源推理框架：

vLLM：Berkeley 开发，首创 PagedAttention
SGLang：Berkeley 开发，首创 RadixAttention，适合 Agent 场景
TensorRT-LLM：NVIDIA 开发，GPU 高度优化
llama.cpp：纯 C++，支持 CPU 推理，本地运行

算术强度回顾#

定义：算术强度 = FLOPs / 字节传输量（越高越好）

以矩阵乘法 X (B x D) @ W (D x F) 为例：

FLOPs： $2 \cdot B \cdot D \cdot F$
字节传输： $2BD + 2DF + 2BF$ （bf16 每个 number 2 bytes）
算术强度 $\approx B$ （当 $B \ll D, F$ 时）

H100 参数：

FLOPs/s：989T
内存带宽：3.35T bytes/s
加速器强度： $295$

结论：

算术强度 > 295 → 计算受限（好）
术强度 < 295 → 内存受限（坏）
极端情况 B=1 → 算术强度=1 → 内存受限

推理的算术强度分析#

朴素推理问题：生成 T tokens 需要 $O(T^3)$ FLOPs（每次前向传播 $O(T^2)$ ）

解决方案：KV Cache

存储每个 sequence、token、layer、head 的 H 维向量
避免重复计算 prefix

推理的两个阶段：

Prefill：编码 prompt（可并行，类似训练）
Generation：逐个生成新 token（顺序执行）

MLP 层分析#

FLOPs： $6 \cdot B \cdot T \cdot D \cdot F$ 字节传输： $4BTD + 4BTF + 6DF$ 算术强度 $\approx B \cdot T$

Attention 层分析（FlashAttention）#

FLOPs： $4 \cdot B \cdot S \cdot T \cdot D$ （S 为已生成 tokens，T 为待生成）字节传输： $4BSD + 4BTD$ 算术强度 $= \frac{S \cdot T}{S + T}$

阶段	MLP 算术强度	Attention 术强度
Prefill	$B \cdot S$	$S/2$
Generation	$B$	$<1$

关键洞察：

Prefill 可达到计算受限（批量足够大）
Generation 是内存受限（Attention 无法通过 batching 改善）
Attention 算术强度不依赖 B（每个 sequence 有独立的 KV cache）

延迟与吞吐量权衡#

假设：推理是内存受限，延迟由内存 IO 决定

$\text{Latency} = \frac{\text{Memory}}{\text{Memory Bandwidth}}$

$\text{Throughput} = \frac{B}{\text{Latency}}$

批量大小影响：

更小批量：更好的延迟，更差的吞吐量
更大批量：更好的吞吐量，更差的延迟

TTFT 优化：

Prefill 用小批量（更快 TTFT）
Generation 用大批量（更高吞吐量）

减少 KV Cache 大小（有损方法）#

Grouped-Query Attention (GQA)#

N 个 query heads，但只有 K 个 key/value heads
MHA：K = N
MQA：K = 1
GQA：K 在中间

效果：KV cache 减少 $N/K$ 倍，延迟/吞吐量改善

Multi-Head Latent Attention (MLA)#

存储 compressed vector $c$ （C 维）而非 K, V（N*H 维）
DeepSeek v2：16384 → 512 维
MLA 比 MHA 更准确且更便宜

Cross-Layer Attention (CLA)#

跨层共享 KV（类似 GQA 跨 head 共享）
改善 accuracy-KV cache Pareto frontier

Local (Sliding Window) Attention#

只看局部上下文
KV cache 与序列长度无关
问题：可能损害 accuracy
解决：混合 local + global attention

DeepSeek v4 Attention#

支持 1M context length：

CSA：每 m tokens 压缩成 1 个
DSA：选择 top k
HCA：进一步压缩

量化#

精度格式对比：

格式	字节数	范围	用途
fp32	4	-	训练参数和优化器状态
bf16	2	-	推理默认
fp8	1	[-240, 240]	H100 e4m3，可训练
int8	1	[-128, 127]	推理，比 fp8 便宜
int4	0.5	[-8, 7]	最便宜，精度最低

量化方法：

QAT：训练时模拟量化误差
PTQ：训练后量化，更便宜
GPTQ：用 Hessian 信息补偿量化误差
AWQ：根据 activation 选择哪些 weight 保持高精度

模型剪枝 + 知识蒸馏#

算法：

在小数据集上识别重要的 layer/head/hidden dimension
移除不重要的部分
用原模型蒸馏修复剪枝模型

推投机采样（无损方法）#

核心思想：利用”检查比生成快”的不对称性

流程：

用便宜的 draft model 猜测几个 tokens
用 target model 并行验证
接受看起来正确的 tokens

数学保证：精确采样自 target model（修改的拒绝采样）

实践配置：

Target 70B + Draft 8B
Target 8B + Draft 1B

扩展方法：

Medusa：draft model 并行生成多个 tokens
EAGLE：draft model 使用 target model 的高层特征

动态工作负载处理#

Continuous Batching#

问题：请求不同时到达、不同长度 → ragged array

解决方案：迭代级调度

逐步解码
新请求到达时立即加入 batch

Selective Batching：

Attention：每个 sequence 单独处理
Non-attention：concatenate 所有 sequences

PagedAttention#

问题：KV cache 内存碎片化

内部碎片：分配最大长度但实际生成较少
外部碎片：sections 之间的未用空间

解决方案：借鉴操作系统 paging

将 KV cache 分成非连续 blocks
共享 prefix，copy-on-write

共享场景：

共享 system prompt
同一 prompt 采样多个响应

vLLM 其他优化：

Kernel 融合（减少 launch overhead）
FlashAttention、FlashDecoding
CUDA graphs

总结#

类别	技术	特点
新架构	GQA, MLA, CLA, Local Attention	减少 KV cache
量化	QAT, PTQ, AWQ	降低精度，减少内存
剪枝	Pruning + Distillation	移除不重要部分
无损加速	Speculative Sampling	精确采样，利用检查/生成不对称性
系统优化	Continuous Batching, PagedAttention	处理动态工作负载