NJU LLM 从零到一 · 通勤 Handbook

1. 课程一览

Why this course

一句话

不是调 API 用模型，而是亲手实现类 LLaMA 的 decoder-only 组件，最后拼成完整 Transformer。

主讲: 徐经纬（jingweix@nju.edu.cn）
联合创始: 黄云鹏
评分: 5 次编程作业 80% + 综合大作业 20%
提交: GitHub Classroom → main push；看 score-feedback
环境: Python ≥ 3.10；推荐 GPU + Docker（Nvidia PyTorch 镜像）
本地仓: ~/Documents/nju-llm/{a0..a5}

能力目标（通勤时默念）

理论基础：理解 LLM 各组件在算什么
实现能力：PyTorch Module / functional / shape / dtype
工程习惯：seed、初始化、边界条件、可复现测试
前沿接口：LoRA、MoE、RoPE、Flash 风格 Attention、KVCache

课程反复强调：优先读论文 + 官方文档 + 源码，少依赖浅层博客。作业约定常与 HF 实现略有差异——以 blog TODO 为准。

2. 模块地图与依赖

一张图串起来

实现主线

A0 反向传播→ A1 多头 Matmul→ A2 Norm/Emb/RoPE→ A3 MLP+LoRA+MoE→ A4 Attention→ A5 Transformer

后面作业会 import 前面模块；跳着做从 A3 起几乎不可行。

拼装时谁依赖谁

A1 形状直觉 ─────────────────────┐
                                  ▼
A2 GroupRMSNorm ──► A4 QK-Norm ──► A5 Layer
A2 Vocab/ParallelEmb ─────────────► A5 Block
A2 RoPE (optional) ───────────────► A5 位置编码

A3 DenseMLP+LoRA ──► SparseMLP ───► A5 Layer.mlp
A4 OfflineAttn ────► Online (opt) ► A5 Layer.attn
A5 KVCache ───────────────────────► 推理 decoding

分仓提醒：a0–a5 是独立 Git 仓库。A2 写完要通过的模块，记得同步拷贝到 a3/a4/a5，否则 A5 会 import 到空壳。

3. 通勤 / 工位工作流

路上复习 · 工位输出

🚇 通勤（手机）

只做：读 shape、默写数据流、勾清单
问自己：输入/输出维？哪些是 Parameter？
不要：在手机上 debug 复杂实现
产出：脑子里一张 forward 图 + 待办 3 条

💻 图书馆 / 工位

最小闭环：写一点 → pytest test_toy.py
有 Docker 再跑 test_with_ref.py
push → 看 score-feedback/score.md
通过的模块立刻同步到后续仓库

手写 vs AI vibe coding

自己主导：数据流、shape、seed 约定、LoRA/MoE 分支逻辑
AI 辅助：查 API、读报错、润色注释
别整文件生成：作业约束与 HF 不同，黑盒改到崩溃很常见

判断标准：关掉 AI，你能否在纸上写出 shape 变换？能 → 自己写主干；不能 → 先画图再写。

每个 Task 的固定节奏（15–45min 切片）

读 blog TODO + dtype/device/seed 约束
先 empty Parameter / 骨架 forward
reset_parameters 或核心公式
toy 测例 → 边界（rank=0、LoRA rank=0…）
提交 / 同步

4. A0–A5 速查卡

点开看要点

DONE A0 · Onboarding

熟悉环境、提交流程、手动 vs 自动梯度。通常无需写业务代码。

线性层：Y = X @ W，loss = sum(Y²)
手动 backward：dW = Xᵀ @ dY，dX = dY @ Wᵀ（分母布局，梯度 shape = 变量 shape）
命令：pytest test_toy.py；可选 Docker + ref
提交：只 push 必要源码到 main，硬 DDL 后自动拒收

Blog: A0 Onboarding · 本地: a0/

DONE A1 · Matmul

为后面多头 / 稀疏计算打形状基础。

T1 多头 matmul：A[b,s,h]×W[h,e] → 拆 nh → O[b,s,nh,e]；禁止 for，用 einsum/bmm
T2 importance：top_p ∧ top_k 筛 token → O[t,nh,e]
T3 有 grad_output 才算 dA/dW；注意 autograd 副作用

A2 = A1.view(b,s,nh,hd); W2 = W1.view(nh,hd,e)
O2 = einsum("bshd,hde->bshe", A2, W2)

本地: a1/src/functional.py · 你这边主干已实现

WIP A2 · Norm & Embedding

T1 GroupRMSNorm：沿 hidden 分组做 RMS；γ 可学习；高精度计算，输出 dtype/device 跟 X
T2 VocabEmbedding：ID → 查表 T[v,e]；normal init + seed
T3 ParallelVocabEmbedding：词表按 rank 分片；区间外置 0；seed=base+rank
Opt T4 RoPE：NTK-aware；buffer 非 persistent；dynamic 外推

A4 QK-Norm、A5 Emb/RoPE 都依赖这里。必修优先，RoPE 可后置。

本地: a2/src/modeling/{norm,vocab_emb,pos_emb}.py

NOW A3 · Modeling MLP

当前主战场。见下方「A3 深潜」。

Dense GLU-MLP + 初始化
LoRA adapters（rank=0 可关）
SparseMLP = MoE（多个 Dense expert + gate top-k）

本地: a3/src/modeling/mlp.py

HARD A4 · Attention

Offline Sliding-Window Attn：mask（full/causal/window）、scale、temp/cap/clip/dropout、QK-Norm
Pack：Q_K_V / Q_KV / QKV；Layout：BSHD / SBHD / THD+cu_seqlens
GQA：kv heads repeat_interleave；矩形 Q/K 用 bottom-right causal
Opt Online：分块 + lse 在线合并（Flash 思想）

建议工期 4–6 天；时间紧只保 Offline 满分。

LATER A5 · Transformer Block

KVCache：set/get/append/reset（推理 prefilling + decoding）
DecoderLayer：PreNorm→Attn(+RoPE)→残差→PreNorm→MLP→残差
DecoderBlock：Emb → N×Layer → FinalNorm → LMHead；lm_head_tied；参数量统计
seed 偏移严格按作业附表（按 layer_idx）

本地 a5 中 attention/pos_emb 似有实现痕迹；norm/mlp/transformer 仍需补齐。

5. A3 深潜（当前）

通勤默写版

Dense GLU-MLP 在算什么

X: [b, s, d]
up   = X @ W_up      # W_up:   [d, ffh]
gate = φ(X @ W_gate) # W_gate: [d, ffh]
h    = up ⊙ gate     # 逐元乘
O    = h @ W_down    # W_down: [ffh, d]
# O: [b, s, d]，dtype/device 与 X 一致
# 无 bias

φ 由 activation_type 决定：SILU/GELU/RELU/SIGMOID/BILINEAR…

Learnables（先记清单）

基座：W_up, W_gate, W_down
LoRA（rank>0）：低秩 A、B + dropout；α 缩放；α 默认 = rank
Sparse：多个 expert（各自是 DenseMLP）+ gate 矩阵 G[h, ne]（float32）

reset_parameters 心智模型

不是每个 weight 对外调一次 reset；而是：

__init__:
  创建 empty Parameter
  调用一次 self.reset_parameters()  # 一次搞定全部

reset_parameters:
  对每个权重：manual_seed → init_ 填数

SIGMOID / BILINEAR → Xavier normal
其余 ReLU 族 → Kaiming normal，mode=fan_in
seed：up +1，gate +2，down +3（相对 init_base_seed）

你当前存法是 [in, out]（配合 X@W）。PyTorch kaiming 默认按 Linear 的 [out, in] 理解二维权重。写完用 std ≈ √(2/fan_in) 自检 fan 维是否正确。

你的代码进度（本地快照）

✅ args → self.xxx
✅ empty：W_up[d,ffh], W_gate[d,ffh], W_down[ffh,d]
🟡 reset_parameters：已选 xavier/kaiming 分支，待补 seed + 三次 init_ + 去掉 raise
⬜ __init__ 末尾调用 reset_parameters()
⬜ forward（GLU）
⬜ LoRA 分支
⬜ SparseMLP

LoRA（Part2）通勤记忆

ΔW ≈ (α / r) · A @ B   # 低秩
# 作业简化：整段 MLP 只挂一次 LoRA（≠ 每个投影各一个）
# rank==0 → 完全跳过 LoRA 逻辑
# A,B：uniform init；seed = lora_init_base_seed +1/+2
# dropout 用 lora_dropout_seed 保证可复现

SparseMLP / MoE（Task2）

G: [d, ne] → logits → softmax → 每 token 概率 P
取 top-k experts → 对这 k 个概率 renormalize
每个 expert ≈ 小 DenseMLP（ffh_size = e = ffh//ne）
本地只实例化 rank 分片上的 experts；非本地路由输出 0（partial sum）
G 的 dtype 固定 float32；expert 的 seed 按全局 expert index 偏移

A3 工位下一刀（别贪多）

补完 base 三权重的 reset（seed + init_）
self.reset_parameters() 挂上；能实例化
写 forward（先无 LoRA）
再开 LoRA，再开 Sparse

6. 工程模式速记

各作业通用

Module 模板

class Foo(nn.Module):
  def __init__(...):
    super().__init__()
    # 存超参
    # empty Parameter / 子模块
    self.reset_parameters()

  def reset_parameters(self):
    # seed + init_

  def forward(self, x):
    # 计算；对齐 dtype/device
    return out

几乎每题都考

输出 dtype/device 跟输入（参数可以不同）
可学习参数用 nn.Parameter
初始化要 可复现 seed
边界：None、0、整除 assert
禁止无意义 Python for 扫 batch（用张量 API）

调试命令

cd ~/Documents/nju-llm/a3
pip install -r requirements.txt   # 各作业可能不同
pytest test_toy.py -q
# Docker 有 ref 时：
# bash run_docker.sh && pytest test_with_ref.py

7. Shape 速查表

横滑查看

模块	输入	关键参数	输出
A1 multi-head matmul	A[b,s,h], W[h,e]	nh, hd=h/nh	[b,s,nh,e]
A1 importance	同上 + P[b,s]	top_p, top_k	O[t,nh,e]
GroupRMSNorm	X[b,s,h]	group_size, γ	[b,s,h]
VocabEmbedding	I[b,s] long	T[v,e]	E[b,s,e]
DenseMLP	X[b,s,d]	W_up/gate/down	[b,s,d]
Attn Q/K/V	Q[b,sq,hq,hd]	K/V[b,skv,hkv,hd]	O 同 Q layout
DecoderBlock	token id [b,s]	L layers	logits [b,s,vocab]

8. 高频踩坑

扣分重灾区

seed 只设一次：多个矩阵共享同一随机流 → 测例挂
init 分支写反：SIGMOID/BILINEAR 才 Xavier，其余 Kaiming
fan_in 与 weight layout：[in,out] vs [out,in] 搞混 → std 不对
输出 device/dtype 漂了：中间用 float32 算可以，最后要 cast 回输入
LoRA rank=0 没短路：仍创建/计算 A、B
lora_alpha is None：应视为 = rank
MoE 忘 renorm：top-k 概率未重新归一化
Attention 矩形 mask：不是 top-left，作业要求 bottom-right
改了 __init__.py 乱导出：导入路径炸
分仓未同步：a3 写完 a5 还是 stub

9. 可勾选清单

存在本机 localStorage

A3 · Dense 本周

empty 三权重 shape 正确 reset：三条 seed 偏移 +1/+2/+3 reset：xavier / kaiming 分支 __init__ 调用 reset_parameters() 能实例化且 W.std() 非零 forward GLU 无 LoRA 过 toy LoRA rank=0 / >0 都过 SparseMLP top-k + 本地 expert push 并查看 score-feedback

通勤 10 分钟自测

默写 Dense GLU 四步与 shape 说出为何 Kaiming 用 fan_in 说出 reset 调用几次、在哪调说出 LoRA rank=0 时行为画出 A2→A5 依赖箭头想起 bottom-right causal 是 A4 的事

勾选状态保存在浏览器本地；换浏览器会清空。

10. 链接与资料

需要联网

手机怎么打开这份 Handbook

Mac 上用 AirDrop / iCloud / 微信文件传输助手，把整个 handbook 文件夹发到手机
或用同一 iCloud Drive / 坚果云同步 Documents/nju-llm/handbook
手机 Safari / Chrome 打开 index.html
可选：分享菜单 → 添加到主屏幕，当「小 App」用

纯静态单文件，无后端；离线可看已打开缓存的页面。搜索与勾选依赖浏览器本地存储。

大语言模型基础
从零到一实现之路

目录

1. 课程一览

一句话

能力目标（通勤时默念）

2. 模块地图与依赖

实现主线

拼装时谁依赖谁

3. 通勤 / 工位工作流

🚇 通勤（手机）

💻 图书馆 / 工位

手写 vs AI vibe coding

每个 Task 的固定节奏（15–45min 切片）

4. A0–A5 速查卡

5. A3 深潜（当前）

Dense GLU-MLP 在算什么

Learnables（先记清单）

reset_parameters 心智模型

你的代码进度（本地快照）

LoRA（Part2）通勤记忆

SparseMLP / MoE（Task2）

A3 工位下一刀（别贪多）

6. 工程模式速记

Module 模板

几乎每题都考

调试命令

7. Shape 速查表

8. 高频踩坑

9. 可勾选清单

A3 · Dense 本周

通勤 10 分钟自测

10. 链接与资料

手机怎么打开这份 Handbook

大语言模型基础从零到一实现之路

目录

1. 课程一览

一句话

能力目标（通勤时默念）

2. 模块地图与依赖

实现主线

拼装时谁依赖谁

3. 通勤 / 工位工作流

🚇 通勤（手机）

💻 图书馆 / 工位

手写 vs AI vibe coding

每个 Task 的固定节奏（15–45min 切片）

4. A0–A5 速查卡

5. A3 深潜（当前）

Dense GLU-MLP 在算什么

Learnables（先记清单）

reset_parameters 心智模型

你的代码进度（本地快照）

LoRA（Part2）通勤记忆

SparseMLP / MoE（Task2）

A3 工位下一刀（别贪多）

6. 工程模式速记

Module 模板

几乎每题都考

调试命令

7. Shape 速查表

8. 高频踩坑

9. 可勾选清单

A3 · Dense 本周

通勤 10 分钟自测

10. 链接与资料

手机怎么打开这份 Handbook

大语言模型基础
从零到一实现之路