Stanford CS336 Language Models from Scratch-Assignment 1 OVerview

前些天计划看Stanford 新开的 CS336，顺便记录下，发到这里，大家一起交流。

1. Assignment 1 概览

核心目标：从零开始构建训练标准 Transformer 语言模型（LM）所需的所有组件。

作业核心任务 从底层开始实现以下四个核心部分：

字节对编码（BPE）分词器： 实现支持 Unicode 的字节级 BPE 分词器，包括预分词、合并计算、编码（Encoding）和解码（Decoding）。
Transformer 模型架构： 构建一个 Pre-norm 结构的 Transformer。具体组件包括：
- Token 嵌入（Embeddings）。
- RMSNorm（均方根层归一化）。
- SwiGLU 激活函数的位置对应前馈网络（FFN）。
- 旋转位置嵌入（RoPE）。
- 因果多头自注意力机制（Causal Multi-Head Self-Attention）。
损失函数与优化器： 实现 交叉熵损失（Cross-entropy loss） 和 AdamW 优化器。
训练与基础设施： 编写支持序列化和加载模型/优化器状态（Checkpointing）的训练循环，并实现支持内存映射（mmap）的高效数据加载器。

具体作业要求

"从零开始"（From Scratch）原则： 严禁使用 torch.nn、torch.nn.functional 或 torch.optim 中的绝大部分定义。
- 允许使用： torch.nn.Parameter、容器类（如 Module, ModuleList, Sequential）以及 torch.optim.Optimizer 基类。
- 禁止使用： 例如 nn.Linear、nn.Embedding、nn.LayerNorm 或内置的 Adam 优化器等，必须手动实现其逻辑。
测试驱动： 必须通过 test_*.py 中提供的所有单元测试，学生需要通过 adapters.py 将自己的代码接入测试框架。
AI 工具使用限制： 允许使用 LLM 咨询低级编程或高级概念问题，但禁止直接用其求解作业问题；强烈建议禁用 IDE 的 AI 自动补全功能（如 Copilot）。
性能要求： 对于 BPE 训练，要求在处理 TinyStories 数据集时，在 30GB 内存下 30 分钟内完成（建议使用多进程并行化处理预分词）。 数据集
本次作业用到的是纯文本数据集： TinyStories 和 OpenWebText 子样本。
项目结构如下，我是在项目目录下创建 data 目录，将数据集下载到 data 目录下

tree assignment1-basics/
assignment1-basics/
├── CHANGELOG.md
├── LICENSE
├── README.md
├── cs336_basics
│ ├── Einops.ipynb
│ ├── init.py
│ ├── nn_utils.py
│ ├── pretokenization_example.py
│ └── transformer.py
├── cs336_spring2025_assignment1_basics.pdf
├── make_submission.sh
├── pyproject.toml
├── tests
│ ├── init.py
│ ├── _snapshots
│ │ ├── test_4d_scaled_dot_product_attention.npz
│ │ ├── test_adamw.npz
│ │ ├── test_embedding.npz
│ │ ├── test_linear.npz
│ │ ├── test_multihead_self_attention.npz
│ │ ├── test_multihead_self_attention_with_rope.npz
│ │ ├── test_positionwise_feedforward.npz
│ │ ├── test_rmsnorm.npz
│ │ ├── test_rope.npz
│ │ ├── test_scaled_dot_product_attention.npz
│ │ ├── test_swiglu.npz
│ │ ├── test_train_bpe_special_tokens.pkl
│ │ ├── test_transformer_block.npz
│ │ ├── test_transformer_lm.npz
│ │ └── test_transformer_lm_truncated_input.npz
│ ├── adapters.py
│ ├── common.py
│ ├── conftest.py
│ ├── fixtures
│ │ ├── address.txt
│ │ ├── corpus.en
│ │ ├── german.txt
│ │ ├── gpt2_merges.txt
│ │ ├── gpt2_vocab.json
│ │ ├── special_token_double_newlines_non_whitespace.txt
│ │ ├── special_token_trailing_newlines.txt
│ │ ├── tinystories_sample.txt
│ │ ├── tinystories_sample_5M.txt
│ │ ├── train-bpe-reference-merges.txt
│ │ ├── train-bpe-reference-vocab.json
│ │ └── ts_tests
│ │ ├── model.pt
│ │ └── model_config.json
│ ├── test_data.py
│ ├── test_model.py
│ ├── test_nn_utils.py
│ ├── test_optimizer.py
│ ├── test_serialization.py
│ ├── test_tokenizer.py
│ └── test_train_bpe.py
├── tmp.py
└── uv.lock

环境配置

项目使用 uv 管理，pyproject.toml 提供了需要的依赖，在项目目录下执行 uv sync即可配置本项目的虚拟Python 环境并安装所有需要的依赖，然后 source .venv/bin/activate激活并进入 cs336-basics 环境。

3. 涉及的专业知识

完成此作业需要掌握以下领域的深层知识：

文本处理： 深入理解 Unicode 标准 （Code points）、UTF-8 编码 以及子词分词（Subword Tokenization）的压缩逻辑。
线性代数与数值计算：
- Einsum 符号： 强烈建议使用 einsum 或 einops 进行高效的张量维度变换和矩阵运算。
- 数值稳定性： 在实现 Softmax 和交叉熵时，需应用"减去最大值"等技巧防止数值溢出。
深度学习架构： 理解 Pre-norm 与 Post-norm 的区别、因果掩码（Causal masking）的原理、以及现代模型（如 Llama/Qwen）中使用的 SwiGLU 和 RoPE 等改进组件。
资源核算（Resource Accounting）： 计算模型的参数量、浮点运算量（FLOPs）、内存占用（包括梯度和优化器状态），并理解模型 FLOPs 利用率（MFU）。
优化算法： 掌握带动量的优化器原理、权重衰减（Weight Decay）、学习率预热（Warmup）与余弦退火调度（Cosine Annealing），以及梯度裁剪（Gradient Clipping）。

4. 实验与评估任务

数据集： 在 TinyStories 和 OpenWebText 两个数据集上进行实验。
模型消融实验： 比较层归一化（有无 RMSNorm）、归一化位置（Pre-norm vs. Post-norm）、有无位置编码（RoPE vs. NoPE）以及激活函数（SwiGLU vs. SiLU）对模型性能的影响。
排行榜竞赛： 在规定的计算预算（1.5 H100 小时）内，通过调整架构和超参数优化 OpenWebText 上的验证集损失。

比喻理解： 这份作业就像是给你一堆原材料（基础张量运算），要求你不仅要亲手打磨出每一个零件（线性层、注意力机制），还要组装成一台复杂的精密机器（Transformer），最后还要在有限的油耗（计算时间）下，让这台机器在赛道（排行榜）上跑出最高的速度。