从啥也不会到DeepSpeed————一篇大模型分布式训练的学习过程总结

你猜怎么着？我写上一篇GPT教程的时候，后台就被读者轰炸了：“大佬，模型理论我啃懂了，但想训个7B的模型，单卡直接黑屏，分布式到底怎么玩啊？”

哎，说实话，一年前的我，也是这么蒙圈的。当时觉得什么分布式，不就是多插几张卡嘛，简单！结果真上手才发现——坑深不见底啊！

很多教程呢，要么上来就扔论文，看得人脑壳疼；要么只讲概念，实操一个字不提。

所以今天，咱们就来点实在的——我直接回答你最关心的四个问题。全是血泪经验，踩过的坑都给你标好了，能救一个是一个。

拷问一：大模型凭什么非要分布式？我用单卡不行吗？

两个字：不行。原因就两个——要么塞不下，要么跑不动。

先说显存。

拿GPT-2举例，1.5B参数（15亿），FP16存权重只要3GB。你觉得“这有什么了不起的？”对吧？

结果一训练，直接OOM。为什么？因为你忘了一个东西——优化器状态！

用Adam + 混合精度训练（现在基本都这么干），你需要的不仅仅是那3GB的FP16参数。看看账单：

• FP16模型参数：3GB
• FP16梯度：3GB
• FP32参数拷贝（Adam更新需要更高精度）：6GB
• Adam的momentum和variance（也都是FP32）：各6GB，总共12GB

加起来：3+3+6+12=24GB！

一张V100才32GB，再加上输入和激活值……你想想，是不是一下就满了？

我当时就是栽在这儿：以为3GB的模型能有什么大事，结果训练代码一跑，直接OOM，查了半天才发现，优化器状态才是真正的“显存黑洞”。

那GPT-3呢？175B参数，FP16就是350GB，仅模型状态（参数、梯度、优化器）就需要约2.8TB！就算今天H100 NVL有188GB显存（组合），单卡也塞不下。所以分布式训练不是什么“高级选项”，是唯一出路。

另一个场景：模型虽然能塞进单卡，但训练时间忍不了。比如1.5B的模型，单卡训一个epoch可能要一周，你用8卡数据并行，一天搞定。这时候分布式就是为了加速，多花点配置时间，完全值得。

拷问二：什么DP、DDP、ZeRO、模型并行……到底选哪个？

直接给答案：别再手动搞DP了！无脑用torchrun + DistributedDataParallel（DDP）。如果DDP显存不够，再上DeepSpeed的ZeRO。

我第一次踩坑就是用的DP。

PyTorch的DataParallel封装确实方便，两行代码就跑起来了。但模型一变大，慢得你想砸电脑。

为什么？因为DP在每个batch反向传播后，要把梯度从所有worker复制到主卡做allreduce，主卡变成瓶颈。卡越多，通信开销越爆炸。

后来我用两张A100训GPT-2，做了个对比实验：

• DP：`train_samples_per_second = 1.96`
• DDP（有NVLink）：1.96（DS已经有点吃紧了，但两张卡差距还不大）
• DDP（无NVLink，设置`NCCL_P2P_DISABLE=1`）：1.52，直接慢了30%！

这还是两张卡。要是四张八张，DP的通信模式更糟糕。所以现在PyTorch官方都推荐DDP，而且torchrun把进程启动都封装好了，用起来不比DP复杂。

那什么时候上ZeRO呢？

当模型大到DDP单卡也放不下的时候。ZeRO本质还是数据并行，但它把模型状态（参数、梯度、优化器状态）切分到各个GPU上，去掉了冗余。

核心区别：DDP每张卡都存完整一份；ZeRO把重复的东西砍了，显存自然省下来。后面我会细讲三阶段。

至于模型并行（张量并行、流水线并行），我建议你先别碰。除非模型真到百亿千亿级别（比如Megatron-LM那套），否则ZeRO+数据并行就够用了。张量并行通信量巨大，一般只在单机多卡内部用，跨机器成本高到怀疑人生。

拷问三：训练一个大模型，显存到底是怎么被吃掉的？

你以为8B模型用FP16只要16GB？太天真了！训练时显存需求是推理的8倍以上。

我就拿8B模型举个例。FP16参数占16GB，听着不多对吧？但训练时每层还会存激活值（前向传播的中间结果），后向传播要用的。激活值大小和batch size、序列长度强相关，这部分也有几十GB。

但更固定的是模型状态：权重、梯度、优化器状态。

以混合精度训练为例，Adam优化器下，每个参数需要：

• FP16参数：2 bytes
• FP16梯度：2 bytes
• FP32参数副本：4 bytes
• FP32 momentum：4 bytes
• FP32 variance：4 bytes

每参数共16 bytes。

8B模型 × 16 = 128GB！ 还没算激活值。

所以啊，要是你觉得16GB推理没问题，训练也能跑，那绝对会傻眼。记住这条公式：每个参数16 bytes，回头算一下你的模型，就知道分布式跑不掉了。

拷问四：DeepSpeed的ZeRO三阶段，到底省了什么？我该用哪个？

ZeRO全称“零冗余优化”，说白了就是把重复的砍掉。分布式训练每张卡都存一份完整参数，那是浪费。ZeRO通过分拆去冗余。

第一阶段：只分拆优化器状态。

参数和梯度每张卡还是完整的，但optimizer states（FP32参数拷贝、momentum、variance）分成N份，各卡只维护自己的那份。通信量和DDP一样（都是all-reduce梯度）。省多少？大约4倍。我一般建议先从Stage1起步，几乎不改模型代码，只需DeepSpeed配置文件。

第二阶段：把梯度也分拆。

每张卡只存自己负责的梯度（通过reduce-scatter分散），通信量仍然是数据并行级别。省更多，但反向传播后需要多一次gather操作，开销不算高。

第三阶段：模型参数也分拆。

每卡只存自己负责的参数。前向计算时需要从其他卡动态取参数（all-gather），后向计算梯度后再做reduce-scatter把梯度分散到各卡。通信量比Stage2高了约50%，网络带宽不够的话反而会变慢。

我的实战选择：

• 如果DDP能跑下，就用DDP（最简单）。
• DDP显存差一点，尤其batch size受限，上ZeRO Stage1就够了，几乎不增加通信。
• 如果模型大到DDP直接OOM，上Stage2，通常能解决大部分问题。
• Stage3慎用。我有一次因为网卡是25G以太网，训练比Stage2慢了一倍。只有单机多卡且NVLink连接时才考虑Stage3。

另外，ZeRO还可以配合Offload（把参数卸载到CPU内存）。但我个人建议：除非实在没办法，否则CPU-GPU来回搬数据，训练效率低到离谱。Offload适合微调或推理，但不适合从头训练大模型。

彩蛋：血泪教训大放送

分布式训练，除了上面那些，还有几个“坑中坑”，我一次说清楚。

1. 网络带宽比想象中关键一万倍。

没有NVLink，DDP默认走PCIe，跨机器还得走IB/RoCE。我调集群时，光nccl版本不匹配就搞了两天。所以训练卡能放同一台机器就别分开。

2. batch size变大，学习率要跟着调。

数据并行下总batch size是单卡batch size×卡数。翻倍卡数如果也翻倍batch size，学习率要按线性缩放调整。我第一次照搬单卡的lr，Loss直接发散，血压飙升。

3. 混合精度训练容易出现inf/nan。

记得开启loss scaling。PyTorch AMP自动化还行，但DeepSpeed需要自己配置。

4. torchrun启动命令别搞错。

我第一次跑DDP时用老式的python -m torch.distributed.launch，结果各种环境变量问题。后来直接用新版推荐的：

torchrun --nproc_per_node 8 train.py

一行搞定，自动管理环境变量。

5. Hugging Face Trainer可以直接集成DeepSpeed。

不用自己写Engine。在TrainingArguments里指定deepspeed配置文件路径，Trainer自动帮你搞定，省心省事。

6. 千万别一上来就啃论文！

我当初看ZeRO论文时觉得像天书。后来先跑通代码再回头看，才恍然大悟。建议学习顺序：

懂单卡训练 → 用DistributedDataParallel跑数据并行 → 用DeepSpeed跑ZeRO → 遇到问题翻论文。

最后，记得一件事：分布式训练和单卡训练最大的不同是——问题有时候出在代码，但更多时候出在环境、网络、驱动。

我调集群环境时，光nccl版本不匹配就搞了两天，心态差点崩了。

但别怕。

每一次OOM，都是一张通往大模型工程师的入场券。卡得多了，自然就会了。

有什么实战上的问题，评论区见，我尽量回。一起跑起来！🚀