# 更大、更快、更省 · DDP、编译、Flash 与 MFU

> 一次收拾 nanoGPT 让训练上规模的几个开关——DDP 多卡数据并行（只在最后一个 micro-step 同步梯度）、torch.compile 编译、Flash Attention、混合精度，以及那个衡量「GPU 跑得多满」的 MFU 指标。

- 作者：David（道雾轩）
- 专栏：nanoGPT 源码解读（https://daiw.org/manual/nanogpt.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/nanogpt/scaling-and-speed
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

# 更大、更快、更省 · DDP、编译、Flash 与 MFU

主干读完了。这最后一章，把散落在 `train.py`、`model.py` 里那几个「让训练又快又能上规模」的开关一次收拢。它们不改变模型学什么，只改变**学得多快、能不能铺到多卡上**。

## DDP：把训练铺到多张卡

单卡太慢，就上多卡。nanoGPT 用 PyTorch 的 **DDP（Distributed Data Parallel，分布式数据并行）**。核心思想：

> **每张 GPU 各跑一份**完全相同的模型副本，各喂**不同的数据**；每步反向后，把各卡的梯度做一次**全局平均(all-reduce)**，让所有副本以同样的梯度更新——于是它们始终保持一致，等效于用「N 张卡 × batch」的大 batch 在训练。

代码里（约 82–101 行）靠环境变量判断是不是 DDP 运行：

```python
ddp = int(os.environ.get('RANK', -1)) != -1  # torchrun 会设 RANK
if ddp:
    init_process_group(backend=backend)       # 建立进程组
    ddp_rank = int(os.environ['RANK'])
    ddp_local_rank = int(os.environ['LOCAL_RANK'])
    ddp_world_size = int(os.environ['WORLD_SIZE'])
    device = f'cuda:{ddp_local_rank}'          # 每个进程绑一张卡
    master_process = ddp_rank == 0             # 0 号进程负责打印/存档
    assert gradient_accumulation_steps % ddp_world_size == 0
    gradient_accumulation_steps //= ddp_world_size  # 梯度累积在多卡间摊分
```

用 `torchrun --nproc_per_node=4 train.py` 一启动，就有 4 个进程各占一卡。几个要点：

- **`master_process`（0 号）**：只有它打印日志、写 checkpoint——不然 4 个进程抢着写，乱套。
- **梯度累积摊分**：总的 `gradient_accumulation_steps` 除以卡数，每张卡少累积几轮（反正梯度会跨卡平均）。
- **只在最后一个 micro-step 同步梯度**（约 293–298 行）：

```python
if ddp:
    model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
```

这是个漂亮的优化：[梯度累积](https://daiw.org/manual/nanogpt/training-loop)中间那些 micro-step 的梯度反正要攒着，**没必要每步都跨卡同步**（同步很贵）；只在**最后一个 micro-step** 才触发一次 all-reduce。作者还特意吐槽：官方推荐用 `model.no_sync()` 上下文管理器做这事，但那样代码臃肿，他直接拨这个内部开关——**又一个「简单优先」的 nanoGPT 式选择**。

训练结束记得 `destroy_process_group()` 收尾。

## torch.compile：一句话换来大提速

```python
if compile:
    print("compiling the model... (takes a ~minute)")
    model = torch.compile(model) # requires PyTorch 2.0
```

**`torch.compile` 是 PyTorch 2.0 的招牌**：它把模型的一堆小算子**融合、编译成优化过的 GPU 内核**，通常能带来可观的提速，而你只需加这一行。代价是第一次要花约一分钟编译。

<Callout type="info">
  还记得[断点续训](https://daiw.org/manual/nanogpt/train-setup)时代码里那个 `unwanted_prefix = '_orig_mod.'` 的清理吗？就是 `torch.compile` 干的——它会给参数名加上 `_orig_mod.` 前缀，所以 resume 加载 checkpoint 时要把这前缀剥掉。**一个上游细节，在另一处留下的「疤」**，不读源码根本想不到。
</Callout>

## Flash Attention：免费的注意力提速

这个我们在[注意力那章](https://daiw.org/manual/nanogpt/attention)已经见过：

```python
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
...
y = torch.nn.functional.scaled_dot_product_attention(q, k, v, ..., is_causal=True)
```

有 PyTorch ≥ 2.0 就自动走 **Flash Attention**——一个 **IO 感知**的融合内核，**不在显存里实际造出那张 `T×T` 的大注意力矩阵**，而是分块流式计算。省显存、又快，数学结果和手写版完全一致。这也是[《大模型解构》里 Mistral 敢弃用滑动窗口](https://daiw.org/manual/llm-anatomy/sliding-window)的底气——标准注意力被 Flash 优化得太好了。

## 混合精度 + TF32：用低精度换速度

[训练循环那章](https://daiw.org/manual/nanogpt/training-loop)讲过的 `autocast` + `GradScaler`，也属于这一类：用 **bf16/fp16** 低精度算，更快更省显存，fp16 再靠 `GradScaler` 防下溢。外加开头设的 `allow_tf32 = True`，让矩阵乘法走 **TF32**（Ampere+ 上「精度略降、速度大涨」的模式）。这些和[《大模型解构》训练侧那章的 FP8](https://daiw.org/manual/llm-anatomy/training-tricks) 是同一条思路的不同档位——**在精度和速度之间选个平衡点**，DeepSeek 把它推到了 FP8，nanoGPT 默认停在 bf16。

## MFU：你的 GPU 到底跑了几成力

最后一个有意思的东西：`estimate_mfu`（`model.py` 约 289–303 行），算 **MFU（Model FLOPs Utilization，模型算力利用率）**：

```python
def estimate_mfu(self, fwdbwd_per_iter, dt):
    N = self.get_num_params()
    cfg = self.config
    L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
    flops_per_token = 6*N + 12*L*H*Q*T          # 每 token 的浮点运算量(PaLM 附录估法)
    flops_per_fwdbwd = flops_per_token * T
    flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
    flops_achieved = flops_per_iter * (1.0/dt)  # 实际每秒 FLOPs
    flops_promised = 312e12                     # A100 bf16 理论峰值 312 TFLOPS
    mfu = flops_achieved / flops_promised
    return mfu
```

- `flops_per_token = 6*N + 12*L*H*Q*T`：估算「跑一个 token 要多少次浮点运算」——`6*N` 是参数相关部分（前向 2、反向 4），`12*L*H*Q*T` 是注意力那部分，出处是 **PaLM 论文附录 B** 的估法。
- 拿「实际达到的 FLOPs/秒」除以「A100 的理论峰值 312 TFLOPS」，得到 **MFU**——**你实际用了这张卡几成的力**。

**MFU 是一个特别实用的体检指标。** 理论峰值谁也摸不到，但如果你的 MFU 只有 10%，说明 GPU 大半时间在空等（数据没喂上、内核没优化）；调到 40–50%，才算把卡用得比较满。它把「我的训练快不快」从玄学变成了一个**能对着优化的数字**——这也很符合 nanoGPT「一切都要看得见、可量化」的气质。

## 通关

至此，**nanoGPT 的每一行，你都读完了**：从 `prepare.py` 把莎士比亚压成整数，到 `model.py` 五个类搭起 GPT，到 `train.py` 的训练循环把它一步步教会，到 `sample.py` 让它开口，再到这一站的多卡、编译、Flash、MFU。六百行，没有一处是黑箱了。

最后一章，我们回望整趟旅程，并把这副「读懂了的骨架」接到真实世界——它怎么一路长成你在[《AI 大模型解构》](https://daiw.org/manual/llm-anatomy)里见到的那些庞然大物。👉 [收尾：从 bit 到会写莎士比亚，再往上是什么](https://daiw.org/manual/nanogpt/recap)

## 参考文献

- [1] Dao et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022. arXiv:2205.14135. —— Flash Attention。
- [2] Chowdhery et al. “PaLM: Scaling Language Modeling with Pathways.” 2022. arXiv:2204.02311. —— MFU 与 `6N` FLOPs 估法（附录 B）。
- [3] Li et al. “PyTorch Distributed: Experiences on Accelerating Data Parallel Training.” VLDB 2020. arXiv:2006.15704. —— DDP。
- [4] Micikevicius et al. “Mixed Precision Training.” ICLR 2018. arXiv:1710.03740. —— 混合精度。
- 源码：Andrej Karpathy, nanoGPT， `train.py` 与 `model.py`. https://github.com/karpathy/nanoGPT