CS336 Assignment 1: Building a Transformer LM

本项目是 CS336 课程作业一的编程任务实现，涵盖了从零构建一个完整的 Transformer 语言模型，以及配套的字节级 BPE (Byte Pair Encoding) 分词器。

实现细节

1. BPE Tokenizer (字节对编码分词器)

• 底层支持：基于 UTF-8 字节序列进行训练，有效避免了内存溢出（OOM）和未知字符问题。
• 性能优化：在合并（Merge）逻辑中，引入了有序队列（最大堆）数据结构来替代传统的蛮力扫描，大幅提升了寻找最高频字节配对的效率。
• 自适应并行：针对预分词（Pre-tokenization）阶段，实现了基于文件大小（10MB 阈值）的单进程与多进程自适应切换，减少了小文件处理时的进程调度开销。

2. 模型架构 (Transformer)

本项目采用 Pre-norm 架构的现代 Transformer 语言模型：

• RMSNorm：替代传统 LayerNorm，提升训练稳定性。公式实现为 $RMSNorm(a_{i})=\frac{a_{i}}{RMS(a)}g_{i}$。
• RoPE (旋转位置编码)：通过对词向量施加与位置相关的旋转变换，利用两个词向量旋转角度的差值将相对位置信息注入模型。
• Causal Multi-Head Attention：带有下三角因果掩码的自注意力机制，防止当前位置关注未来 Tokens，保证自回归生成的正确性。
• SwiGLU FFN：前馈网络使用 SiLU 激活函数与门控线性单元（GLU）结合。核心计算为 $FFN(x)=W_{2}(SiLU(W_{1}x)\odot W_{3}x)$。

3. 训练基础设施

• 优化器：实现了带有权重衰减解耦的 AdamW 优化器。
• 学习率调度：采用带有预热（Warmup）阶段的余弦退火（Cosine Annealing）策略，帮助稳定训练初期的梯度。
• 稳定性保障：引入全局梯度截断，通过限制全局梯度向量的 L2 范数防止梯度爆炸。
• 高效数据加载：针对大规模语料读取，使用 numpy.memmap 实现内存映射，极大地优化了内存占用。

实验与性能

• 学习率调优：在 TinyStories 数据集上进行了超参数搜索，确定了最优策略为 max_lr = 1e-3, min_lr = 1e-4。
• 最终表现：在 A800 GPU 上训练约 327M Tokens 后，验证集 Loss 成功降至 1.39（低于 1.45 的目标阈值）。
• 文本生成：支持带温度缩放的 top-p 采样机制。经过全量数据训练后的模型，生成的长篇幅故事展现出了良好的因果逻辑和上下文连贯性。