更大、更快、更省 · DDP、编译、Flash 与 MFU · nanoGPT 源码解读

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

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

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

DDP：把训练铺到多张卡

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

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

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

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 行）：

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

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

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

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

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

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

还记得断点续训时代码里那个 unwanted_prefix = '_orig_mod.' 的清理吗？就是 torch.compile 干的——它会给参数名加上 _orig_mod. 前缀，所以 resume 加载 checkpoint 时要把这前缀剥掉。一个上游细节，在另一处留下的「疤」，不读源码根本想不到。

Flash Attention：免费的注意力提速

这个我们在注意力那章已经见过：

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 敢弃用滑动窗口的底气——标准注意力被 Flash 优化得太好了。

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

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

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

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 大模型解构》里见到的那些庞然大物。👉 收尾：从 bit 到会写莎士比亚，再往上是什么

参考文献

[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

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