A1 (Basics): Building a Transformer LM

0. 总目标

从零实现训练一个标准 Transformer 语言模型所需的所有组件，并在数据集上完成训练、生成、评测、再做更大数据集实验与排行榜提交。

包含：

• 实现：BPE tokenizer、Transformer LM、cross-entropy + AdamW、训练循环（含保存/加载状态）
• 跑通：在 TinyStories 上训练 tokenizer → 把文本转 ID → 训练 LM → 生成与算 perplexity → 在 OpenWebText 上训练并提交 leaderboard

1) Tokenizer 线（BPE + 编码/解码 + 吞吐/压缩实验）

1.1 Unicode / UTF-8 基础小题（为 byte-level tokenizer 铺垫）

• 理解 Unicode 字符（如 chr(0)）、显示/打印差异、以及字符进入文本会发生什么（短答）
• 比较 UTF-8 vs UTF-16/32 的理由；找一个“错误 UTF-8 解码函数”的反例；给出无法解码的 2-byte 序列（短答）

1.2 训练一个 byte-level BPE（核心编程任务之一）

• 写 train_bpe(input_path, vocab_size, special_tokens)：输出
  • vocab: dict[int, bytes]（token_id → bytes）
    • merges: list[tuple[bytes, bytes]]（按产生顺序记录 merge）
      并通过单元测试
      • 在 TinyStories 上训练 10k vocab 的 tokenizer：记录耗时/内存、最长 token 是否合理、profile 瓶颈（短答）
• 在 OpenWebText 上训练 32k vocab 的 tokenizer，并对比两者 tokenizer 的差异（短答）

1.3 实现 Tokenizer 类（编码/解码 + 流式 tokenization）

• 实现 Tokenizer(vocab, merges, special_tokens)，支持：
  • encode(text)->list[int]
    • decode(ids)->str（非法字节用 Unicode replacement character 处理）
    • encode_iterable(iterable)->Iterator[int]（用于大文件内存友好编码）
    • from_files(...)（从磁盘加载 vocab+merges）
      并通过测试

1.4 Tokenizer 实验题（压缩率、错配 tokenization、吞吐估算）

• TinyStories/OWT 各抽样 10 文档：比较两种 tokenizer 的 压缩率 bytes/token
• 用 TinyStories tokenizer 去 tokenize OWT：压缩率变化/现象描述
• 估算 tokenizer 吞吐（bytes/s）并推算 tokenize 825GB（The Pile）的耗时
• 把训练/验证集编码成 token IDs（建议 uint16）并解释为什么合适

2) 模型线（从零搭 Transformer LM：模块 → Block → LM + 资源核算）

2.1 基础模块（自己写 nn.Module 的“替身”）

• 实现 Linear（无 bias），不能用 nn.Linear / F.linear，并通过测试
• 实现 Embedding（lookup），不能用 nn.Embedding / F.embedding，并通过测试
• 实现 RMSNorm（注意 upcast float32 再 downcast），并通过测试
• 实现 SwiGLU 前馈网络（SiLU + GLU gating，且 d_ff 取约 \(\frac{8}{3} d_{model}\) 并对齐硬件倍数），并通过测试
• 进一步还要把 RoPE、因果多头自注意力等组件补齐（文档后续章节继续展开）

2.2 拼装 Transformer Block 与整机 LM

• 实现 pre-norm Transformer block：RMSNorm → (MHA/FF) → 残差，按公式 \(y = x + \text{MHA}(\text{RMSNorm}(x))\) 等组装，并通过测试
• 实现 Transformer LM：Embedding → 多层 block → 输出归一化与线性投影得到 logits，并通过测试

2.3 资源/复杂度核算（FLOPs & 参数量）

• 做 “Transformer LM resource accounting”：
  • GPT-2 XL 形状模型：参数量、仅加载模型所需内存（短答）
    • 列出前向中主要矩阵乘与 FLOPs，总 FLOPs
    • 哪些部分 FLOPs 最大
    • 对 GPT-2 small/medium/large 重复并分析比例变化

3) 训练与优化线（loss、AdamW、LR schedule、梯度裁剪）

这一部分把“能训练”补齐成“训练得对、训得稳”。

• 训练侧需要实现 cross-entropy loss + AdamW（A1 Overview 明确列为必须实现）
• 实现 cosine learning rate schedule + warmup（有明确公式与测试）
• 还包括 梯度裁剪等训练稳定性机制（同一章节继续往后）

4) 工程化训练脚本（大文件加载、断点续训、可复现实验）

4.1 Checkpointing（断点续训必备）

• 实现 save_checkpoint(model, optimizer, iteration, out) 与 load_checkpoint(src, model, optimizer)->iteration：保存/恢复模型权重、优化器状态、迭代步数，并通过测试

4.2 把所有东西“跑成一条龙”（训练主程序）

• 写训练脚本，至少支持：
  • 可配置的模型/优化器超参
    • 用 np.memmap 内存友好加载大规模 train/val token IDs
    • 定期保存 checkpoint 到用户路径

5) 实验线（TinyStories → OpenWebText → Leaderboard）

• OWT 主实验：在 OpenWebText 上用与 TinyStories 相同架构与训练步数训练，交付 loss learning curve，并解释 loss 差异；再给出生成样例并解释为何同算力下质量更差
• Leaderboard（开放式改进）： * 在给定规则下（数据只用课程提供的 OWT；单次运行 ≤ 1.5 小时 H100 等）训练并提交
  • 交付最终 val loss、learning curve（含墙钟时间轴<1.5h）与做了哪些改动的说明