复现 "Attention Is All You Need"

@hao CC BY 4.0
  • 这篇论文基本算是大模型的开山之作
  • 从零复现论文中的 Base Transformer 模型（德语→英语翻译）
  • 主要参考 The Annotated Transformer，逐步实现模型代码
  • 总耗时约 9 天（从精读论文到训练完成）
  • 最终 BLEU 达到 28.2，基本复现论文结果

过程

  • 先看了几个讲解论文的视频，精读论文原文
  • 跟着 Annotated Transformer 一步一步实现核心代码：模型架构、Multi-Head Attention、数据准备、训练函数
  • 先用小数据集 Multi30K（3 万条）在 Mac 上验证代码正确性
  • 确认无误后，上 AutoDL 用 H800 GPU 跑 WMT14 大数据集
  • 训练跑了4轮，中间做了吞吐量优化（bucket batching、LabelSmoothing 内存修复等）
  • 即便论文、代码、数据全部现成，从头做一遍仍花了 9 天——数据处理、训练参数调整、代码细节排查都需要不少时间

训练数据

  • 小数据集：Multi30K，3 万条德英翻译对，用于本地验证
  • 大数据集：WMT14 DE→EN，约 450 万条德英翻译对
  • 分词：SentencePiece，共享词表，37000 个 token
  • 平均序列长度约 30-50 个 token

训练环境

  • GPU：AutoDL 上的单卡 H800
  • 费用：8.8 元/小时，总训练费用约 40 元（约 4 小时训练）
  • 精度：bfloat16 autocast + TF32 matmul
  • 编译：torch.compile（max-autotune-no-cudagraphs 模式）
  • Micro batch size：128
  • Gradient accumulation：8 步

训练结果

  • 吞吐量：优化后单 GPU 约 80k tokens/second（优化前约 47k）
  • 单 epoch 耗时：优化后约 40 分钟（优化前约 53 分钟）
  • BLEU 进展：
    • 约 4 小时（7 个 epoch）后，BLEU 达到 27-28
    • 最佳 BLEU：28.2（step 30500，第 6 个 epoch）
    • 基本复现了论文中 Base Transformer 的结果（论文报告 27.3）
  • Validation Loss：从 epoch 0 的 2.22 逐步降至 epoch 6 的 1.51

学到的东西

  • 对 Transformer 架构有了更深的理解，结合之前复现 GPT-2 的经验，对整个训练流程更加熟悉
  • 虽然 NanoGPT-2 和 Transformer 模型架构有所不同，但整个流程（准备数据→tokenize→data loader→创建模型→训练迭代）是相似的
  • 即使所有资料都已公开，完整复现一个模型的细节仍然很多，并不是一两天就能跑出来的
  • 数据处理、训练参数、代码正确性等细节都需要仔细打磨
  • 使用 Claude Code 来辅助学习和排查问题。遇到模型架构、训练过程、优化方案中不理解的地方，直接向它提问，它能用简洁的文字和 ASCII 图解来解释清楚，节省了很多时间，非常有帮助

例子：Bucket Batching 的梯度累积偏差

在做吞吐量优化时，引入了 bucket batching（按序列长度分组），TPS 提升了 77%，但发现 validation loss 比之前更高。通过和 Claude Code 的讨论，用 ASCII 图快速定位了问题的根因：

问题：每个 micro-batch 内的序列长度一致，但不同 micro-batch 之间长度差异很大。loss 按每个 micro-batch 的 token 数归一化后再 backward，导致短序列 batch 对梯度的贡献被放大了约 3 倍：

短序列 micro-batch                  长序列 micro-batch
┌───────────────────────┐           ┌───────────────────────┐
│ █████░░░░░░░░░░░░░░░░ │           │ ████████████████████░ │
│ ██████░░░░░░░░░░░░░░░ │           │ ███████████████████░░ │
│ █████░░░░░░░░░░░░░░░░ │           │ ████████████████████░ │
│ ████░░░░░░░░░░░░░░░░░ │           │ ██████████████████░░░ │
└───────────────────────┘           └───────────────────────┘
  ntokens ≈ 2,500                     ntokens ≈ 7,500
  梯度权重 = 1/2500 = 4.0e-4          梯度权重 = 1/7500 = 1.3e-4
  ← 短序列的影响力是长序列的 3 倍！

修复：改为在整个梯度累积窗口（8 个 micro-batch）结束后，按总 token 数统一归一化，确保每个 token 对梯度的贡献相同：

修复前：∇ = Σᵢ ∂Loss_i/∂θ / ntokens_i      ← 短序列 batch 权重过大
修复后：∇ = G × Σᵢ ∂Loss_i/∂θ / Σᵢntokens_i ← 每个 token 权重相同

参考资料