课程作业

Assignment 1: Basics

第一个作业是整个课程的基石，你将从零实现语言模型训练所需的核心组件。

任务概览

• BPE Tokenizer：实现 Byte Pair Encoding tokenizer 的完整训练和编码/解码流程。包括从 corpus 中学习 merge rules、处理 Unicode 文本、以及高效的 tokenization 算法。
• Transformer Language Model：实现完整的 decoder-only Transformer 架构，包括 multi-head self-attention、RoPE (Rotary Position Embedding)、RMSNorm、SwiGLU activation 等现代 LLM 使用的组件。
• Cross-Entropy Loss：实现适用于语言模型的 cross-entropy loss function，处理 padding token 和 label smoothing。
• AdamW Optimizer：从零实现 AdamW optimizer，包括 weight decay decoupling、bias correction 和 gradient clipping。
• Training Loop：将以上组件组合成完整的训练流程，包括 learning rate scheduling（warmup + cosine decay）、gradient accumulation 和 checkpoint saving。
• 模型训练：在 TinyStories 和 OpenWebText 数据集上训练你实现的语言模型，验证其生成文本的质量。

学习目标

通过这个作业，你将深入理解：

• Tokenization 如何将文本转化为模型可处理的 token 序列
• Transformer 中 attention mechanism 的计算细节和内存特性
• 现代 LLM 的架构选择（RoPE vs absolute position, RMSNorm vs LayerNorm, SwiGLU vs ReLU）
• AdamW 与 Adam 的关键区别以及 weight decay 的正确实现方式
• Training recipe 中各个超参数（learning rate, batch size, warmup steps）的作用

预计工作量

约 40-60 小时。这是工作量最大的一个作业，因为需要从零实现大量代码。

代码仓库：github.com/stanford-cs336/assignment1-basics

Assignment 2: Systems

第二个作业聚焦于 GPU 系统优化和分布式训练，让你理解 LLM 训练的系统层面挑战。

任务概览

• Flash Attention 2 (Triton)：使用 OpenAI Triton 编程框架实现 Flash Attention 2 算法。你需要理解 GPU 的内存层次结构（SRAM vs HBM），实现 tiling 策略来减少 HBM 访问次数，并处理 causal masking。这是整个课程中技术难度最高的部分之一。
• DDP (Distributed Data Parallel)：实现基本的数据并行训练，包括 gradient all-reduce、正确处理 batch normalization 统计量、以及 gradient bucketing 优化。
• Optimizer Sharding：实现 ZeRO Stage 1 风格的 optimizer state sharding，将 optimizer state 分布到多个 GPU 上以减少内存占用。
• Profiling：使用 PyTorch Profiler 和 Nsight Systems 等工具分析模型的性能瓶颈，包括 GPU 利用率、内存使用、通信开销等。

学习目标

• 理解 GPU 的 SRAM/HBM 内存层次结构及其对算法设计的影响
• 掌握 Triton 编程模型，能够编写高效的 GPU kernel
• 理解分布式训练中的通信原语（all-reduce, all-gather, reduce-scatter）
• 学会使用 profiling 工具诊断和优化性能问题

前置要求

需要完成 Assignment 1 中的 Transformer 实现。需要多 GPU 环境（至少 2 个 GPU）进行分布式训练实验。

预计工作量

约 30-50 小时。Flash Attention 的 Triton 实现是主要难点。

代码仓库：github.com/stanford-cs336/assignment2-systems

Assignment 3: Scaling

第三个作业探索 Scaling Laws，让你亲手拟合和验证 LLM 训练中最重要的经验规律。

任务概览

• IsoFLOP Scaling Laws 拟合：在不同规模的模型上进行 IsoFLOP 实验（固定总 FLOPs，改变模型大小和训练数据量），记录训练 loss，拟合 Chinchilla-style scaling laws。
• Scaling Laws 预测：使用拟合的 scaling laws 来预测更大规模训练的最优配置（最优模型大小、最优 token 数量），并通过实际训练验证预测的准确性。
• Compute-Optimal Analysis：给定一个 compute budget，使用 scaling laws 确定 compute-optimal 的训练配置，分析 Chinchilla 模型与之前工作（如 Kaplan et al.）的区别。

学习目标

• 理解 Scaling Laws 的经验公式及其适用范围
• 学会设计 IsoFLOP 实验并拟合 power law 函数
• 掌握 compute-optimal training 的概念和实践方法
• 理解为什么 Scaling Laws 对 LLM 研发决策如此重要

实验规模

这个作业需要训练多个不同大小的模型（从数百万到数亿参数），每个模型训练到不同的 token 数量。总计算量较大，但每个单独实验的规模可控。

预计工作量

约 20-30 小时（不含 GPU 训练等待时间）。主要工作在实验设计、数据收集和分析上。

代码仓库：github.com/stanford-cs336/assignment3-scaling

Assignment 4: Data

第四个作业让你亲手构建一个从 Common Crawl 原始数据到训练就绪数据的完整处理 pipeline。

任务概览

• HTML 到文本转换：处理 Common Crawl 的 WARC 文件，使用 trafilatura 或 resiliparse 从 HTML 中提取高质量的正文文本。你需要处理各种边界情况，如编码错误、JavaScript 渲染的页面、多语言内容等。
• Quality Filtering：实现多层质量过滤，包括语言检测（使用 FastText lid 模型）、perplexity 过滤（使用 KenLM n-gram 模型）、基于规则的启发式过滤（文档长度、特殊字符比例、重复行比例等）。
• Deduplication：实现完整的去重 pipeline，包括基于 hash 的精确去重和基于 MinHash + LSH 的近似去重。你需要处理去重在大规模数据上的可扩展性问题。

学习目标

• 理解 Common Crawl 数据的格式和特点（WARC, WET, WAT）
• 掌握 HTML 文本提取工具的使用和局限
• 学会设计和实现多层质量过滤 pipeline
• 深入理解 MinHash 和 LSH 算法的原理和实现
• 体会大规模数据处理的工程挑战

数据规模

作业使用 Common Crawl 的一个子集（约数 GB），足以体验完整的数据处理流程，同时保持在单机可处理的范围内。

预计工作量

约 25-40 小时。MinHash + LSH 的实现和调试是主要挑战。

代码仓库：github.com/stanford-cs336/assignment4-data

Assignment 5: Alignment

最后一个作业让你实现完整的 alignment pipeline，从 SFT 到 RL，将一个 pre-trained 模型转变为能有效遵循指令的 aligned 模型。

任务概览

• Supervised Fine-Tuning (SFT)：在指令数据上对 Qwen 2.5 Math 1.5B 模型进行 SFT。你需要实现 chat template 格式化、只计算 response 部分的 loss、以及合适的学习率策略。
• Expert Iteration：实现 rejection sampling + SFT 的 expert iteration 流程。对每个数学问题生成多个候选回答，用答案验证器选出正确的回答，然后用这些高质量回答进行 SFT。
• GRPO 及其变体：实现 GRPO (Group Relative Policy Optimization) 算法，包括 group-relative advantage 计算、clipped policy gradient、KL penalty 等核心组件。还需实现 GRPO 的一些变体进行对比实验。
• RL 训练实验：在 Qwen 2.5 Math 1.5B 模型上进行 RL 训练，使用数学答案正确性作为 verifiable reward。分析训练曲线、reward 变化、KL divergence 等指标，并评估 alignment 后模型在数学 benchmark 上的改进。

学习目标

• 理解 SFT 数据格式和训练细节（loss masking, chat template）
• 掌握 Expert Iteration / Rejection Sampling 的实现
• 深入理解 GRPO 算法及其与 PPO 的区别
• 体会 RL 训练中 KL 约束的重要性和超参数调优的挑战
• 观察 alignment 前后模型行为的变化

模型与硬件

使用 Qwen 2.5 Math 1.5B 作为 base model，这是一个针对数学推理优化的轻量级模型。RL 训练需要 GPU 资源（推荐至少 1 张 A100 40GB 或等效硬件），因为需要同时维护 policy model 和 reference model。

预计工作量

约 30-45 小时。GRPO 的实现和 RL 训练的调试是主要挑战。

代码仓库：github.com/stanford-cs336/assignment5-alignment