【CS336】导言

CS336: Language Modeling from Scratch 是目前 AI 圈非常硬核的一门课，核心目标就一句话：

不调用 HuggingFace 黑盒，自己从零把一个现代 LLM 做出来。

它和很多"讲理论"的课程不一样，CS336 的作业更像：

• "自己实现 mini GPT"
• "自己写 FlashAttention"
• "自己做数据清洗 pipeline"
• "自己做 RLHF / reasoning"

所以这课在工程圈特别火。 (Stanford CS336)

---

整体作业结构

CS336 一共有 5 个 Assignment：

作业	核心主题	本质
A1	Basics	从零实现 Transformer
A2	Systems	GPU / Triton / FlashAttention
A3	Scaling	Scaling Law
A4	Data	预训练数据工程
A5	Alignment	SFT + RLHF / RL

其实就是：

"完整复刻 OpenAI / Anthropic / DeepMind 的 LLM 技术栈。"

---

Assignment 1：Basics（最重要）

这是整个课程灵魂。

很多人做到这里已经被虐麻了。

核心要求：

• tokenizer
• transformer
• attention
• optimizer
• training loop
• loss
• sampling
• inference

全部自己写。

官方描述：

"Implement all components necessary to train a standard Transformer language model." (Stanford CS336)

---

A1 本质上在干什么

你最终会得到：

model = GPT(...)
loss = model(x)

loss.backward()
optimizer.step()

然后真正训练一个小型 GPT。

---

A1 主要模块

---

1. BPE Tokenizer

这一部分是：

"ChatGPT 为什么知道 token？"

你会实现：

• vocab merge
• token frequency
• encode
• decode

本质是：

h e l l o
↓
hello

不断合并高频 token。

这部分会第一次理解：

GPT 其实不是理解"词"，而是理解 token。

---

2. Transformer

这是核心中的核心。

你会自己写：

• Embedding
• Multi-head Attention
• RMSNorm / LayerNorm
• MLP
• Residual
• RoPE
• Causal Mask

比如 Attention：

\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

你会第一次真正理解：

为什么 attention 会爆显存。

---

3. Causal LM

你会实现：

输入:
I love

标签:
love AI

本质是 next-token prediction。

损失函数：

L=-\sum_i y_i \log p_i

也就是 Cross Entropy。

---

4. AdamW

他们要求你自己实现 optimizer。

不是：

torch.optim.AdamW

而是自己写：

• momentum
• variance
• bias correction
• weight decay

Adam 更新：

\theta_t=\theta_{t-1}-\alpha\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}

这一部分会真正理解：

为什么 optimizer 是"训练稳定性"的核心。

---

Assignment 2：Systems（最难）

这是整门课最变态的部分。

很多 ML 工程师都做不完。

核心：

• GPU profiling
• Triton
• FlashAttention2
• Distributed Training
• Memory optimization

官方要求里明确提到：

自己实现 FlashAttention2。 (Stanford CS336)

---

为什么 A2 很恐怖

因为这已经不是"深度学习"了。

而是：

CUDA / GPU Kernel Engineering

你会接触：

• SRAM
• HBM
• memory bandwidth
• arithmetic intensity
• kernel fusion

这已经是：

OpenAI infra engineer 干的事。

---

FlashAttention 的核心思想

普通 attention：

QK^T
↓
生成巨大 attention matrix
↓
softmax

显存爆炸。

FlashAttention：

• 不存完整矩阵
• 分块计算
• online softmax

所以：

O(n² memory)
↓
O(n)

这是现代 LLM 能训练长上下文的关键之一。

---

Assignment 3：Scaling

这一部分开始接近：

"大模型公司真正的研究"

核心：

• scaling law
• compute optimal
• parameter scaling
• data scaling

经典 scaling：

L(N)=\left(\frac{N_c}{N}\right)^{\alpha_N}

你会理解：

为什么：

• 7B
• 70B
• 671B

不是随便拍脑袋。

而是 scaling law 推出来的。

---

Assignment 4：Data

这一部分非常工业级。

内容：

• Common Crawl
• dedup
• filtering
• quality scoring
• data mixture

本质：

"垃圾数据会毁掉模型。"

这一部分会第一次理解：

为什么 OpenAI 数据壁垒极强。

因为：

模型 architecture 已经卷不动了。

现在拼的是：

数据质量

---

Assignment 5：Alignment / RL

这一部分开始接近：

ChatGPT 的后训练阶段

包括：

• SFT
• RLHF
• reasoning RL
• reward model

官方描述：

用 RL 训练模型做数学推理。 (Stanford CS336)

---

为什么 CS336 这么火

因为它不是：

"介绍 Transformer"

而是：

"真正教你做一个 LLM 公司"

它覆盖了：

方向	是否覆盖
Transformer	✓
Training	✓
GPU优化	✓
分布式	✓
数据工程	✓
Scaling Law	✓
RLHF	✓

所以很多人认为：

这是目前最接近 OpenAI 内部工程体系的公开课程之一。 (Reddit)

---

这门课适合谁

适合：

• AI infra
• 大模型工程
• LLM research engineer
• 推理优化
• Agent infra

不太适合：

• 纯 beginner
• 只想调 API
• 只做 prompt engineering

因为它代码量极大。

Reddit 上很多人都提到：

"assignment huge"

"very implementation-heavy" (Reddit)

---

如果你想真正学懂，推荐路线

建议：

第一阶段

先看：

• Karpathy nanoGPT
• build-nanogpt
• Raschka《Build a LLM》

---

第二阶段

做 CS336 A1

这是质变点。

---

第三阶段

重点吃透：

• FlashAttention
• Triton
• Distributed

这部分是现在 AI 工程师最值钱的技能。

---

一句话总结

CS336 的作业本质是：

"从零复刻一个现代大模型公司的核心技术栈。"

它不是"学会调用 LLM"。

而是：

学会制造 LLM。