# 复现 "Attention Is All You Need"

- 这篇论文基本算是大模型的开山之作
- 从零复现论文中的 Base Transformer 模型（德语→英语翻译）
- 主要参考 [The Annotated Transformer](https://nlp.seas.harvard.edu/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 权重相同
```

## 参考资料

- [论文：Attention Is All You Need](https://arxiv.org/abs/1706.03762)
- [The Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/)
- [视频：Transformer 论文逐段精读（李沐）](https://www.youtube.com/watch?v=nzqlFIcCSWQ)
- [视频：Introduction to Transformers w/ Andrej Karpathy](https://www.youtube.com/watch?v=XfpMkf4rD6E)
- [The Illustrated Transformer](https://jalammar.github.io/illustrated-transformer/)