大模型分布式训练核心策略：从数据并行到ZeRO显存优化实战

聊大模型的时候，大家最津津乐道的是参数量、上下文窗口、各种评测分数。但真正把这些模型从论文变成可用系统的，是背后的分布式训练这套基础设施。我接触分布式训练也有几年了，从早期的单机多卡，到现在需要跨节点协调的大规模集群，踩过的坑不少，也积累了一些心得。这篇笔记不打算罗列概念定义，而是想把分布式训练的底层逻辑、关键策略和实际落地时会遇到的问题串一遍，讲清楚每个选择背后的原因，以及在真实环境中怎么权衡取舍。

在展开之前，先明确一个核心事实：大模型之所以需要分布式训练，不是因为它“大”，而是因为“单设备装不下”和“单设备算不完”这两个硬约束同时存在。装不下指的是显存容量有限，一个上百B参数的模型可能就需要数百GB的显存，单卡根本无法承载；算不完指的是算力时间，即便能装下，按单卡算力训练也需要数年甚至更久，根本无法接受。分布式训练的目标，就是用多台设备的时间换单台设备的空间，再把算力池化，通过并行调度来缩短训练周期。

这套技术栈从底子上分两个维度：数据并行是“复制模型、分数据”，模型并行是“切分模型、聚合计算”。二者不是互斥的，实际落地时往往是多层并行策略的混合组合，而如何组合、如何分层，就是整个分布式训练设计的核心命题。

在后面的章节里，我会逐一拆解这些并行策略的适用场景、核心实现逻辑、参数设定方法，以及实际项目中容易踩的坑。如果你正打算从单卡实验转向大规模训练，或者只是想搞清楚集群训练底层的设计思路，这些内容应该能帮你少走不少弯路。

1. 分布式训练的整体设计思路拆解

1.1 数据并行与模型并行的基座逻辑

数据并行（Data Parallelism，DP）的逻辑最容易理解：每个设备都保存一份完整的模型副本，训练数据被切分成多份分发给各个设备，每个设备独立做前向和反向计算，然后在梯度更新之前，对所有设备上的梯度做一次全局同步，得到的平均梯度再用于更新每份模型副本。

这种设计的好处是逻辑简单、扩展性好。但它有一个隐藏瓶颈：单设备必须能装下整个模型，否则无从谈“每份模型副本”。一旦模型规模超出单卡显存，数据并行就会立刻失效。

模型并行（Model Parallelism，MP）解决的就是数据并行装不下的问题：把模型本身切碎，分散放到多个设备上。模型并行又有两种主流切法。

一种是按层切分，也就是把几十上百层的Transformer网络按层分配到不同设备上，前向和反向过程按顺序在各设备间依次执行。这种方案在实现层面比较简单，但存在明显的通信串行瓶颈——任何时刻只有一部分设备在计算，其他设备都在等待前向或反向的中间结果，设备利用率不高。

另一种是算子内并行，业内更常叫张量并行（Tensor Parallelism，TP），把单层内部的矩阵乘法切块，分布到多张卡上同时计算。这种方式设备间的通信非常频繁，但对单层的大矩阵运算能起到立竿见影的加速效果。

在实际的大模型训练架构里，几乎不会只用一种并行策略。典型的做法是：在单机内部用张量并行把大层切开，在机器之间用数据并行扩大整体吞吐，在更深层的结构上穿插流水线并行来缓解跨机通信压力。越接近底层硬件，越依赖通信高速的并行方式；越接近全局调度，越依赖逻辑简单的复制式扩展。

1.2 从单卡到集群的演进路径

刚开始做单人实验时，单卡训练就足够了，最多在PyTorch里套一个DataParallel来用用多卡。但DataParallel在每次前向计算时都会复制模型、收集梯度输出，通信效率很低，遇到稍大的模型或者稍多的卡，性能就会断崖式下跌。后来大家普遍改用DistributedDataParallel（DDP），它只在反向传播结束后做一次梯度同步，通信量小得多，而且能直接对齐多进程多机的模型。

当模型大到单卡放不下时，就进入了多卡模型并行阶段。最开始可以从张量并行入手，把注意力层和前馈层的矩阵按列或按行切块，分到不同卡上计算。但张量并行会产生大量的all-reduce通信，所以只适合用在高速互联的节点内部。如果模型进一步变大，单节点都放不下了，就必须跨节点训练，这时流水线并行和混合并行的组合就派上了用场。

这个演进路径本质上是一个决策链：先看模型能不能塞进单卡，不能就考虑张量并行；再看单节点够不够装，不够就引入流水线并行和跨节点数据并行；最后才是在全局层面配置ZeRO（零冗余优化器）这类显存优化技术，进一步压榨设备利用率。

1.3 通信拓扑与集群架构的基础认知

分布式训练的瓶颈往往不是计算而是通信，理解这一点后，你就需要关心集群的拓扑结构了。

单机多卡场景下，卡间通信走NVLink或PCle，延迟极低、带宽很高，适合做高频次的小数据通信，比如张量并行里面的逐层all-reduce。跨节点场景下，通信需要经过网络接口和交换机，带宽和延迟比机内通信差一到两个数量级，所以跨节点的通信要尽量稀疏化、大块化——把多个小张量的梯度打包成一个大的通信包，减少通信次数。

我在实践中的做法是：先画一张节点拓扑图，标注哪些卡在同一个交换域内，哪些跨交换域，然后根据这张图来做并行策略分配。把通信量最大的并行组放在同一个节点内部，把通信量小的复制式并行放到跨节点层面，是整个集群架构设计的不变原则。

2. 并行策略的底层原理与选型实战

2.1 数据并行DDP的通信机制与梯度同步细节

DDP 的核心机制是反向传播完成后的梯度全局同步。每个进程计算得到本地梯度，然后通过AllReduce操作在所有进程间求平均，平均梯度再用于优化器更新。

常见的AllReduce实现有两种树形聚合和ring-allreduce。树形聚合像一棵二叉树，数据从叶子节点逐层向上聚合再广播下来，通信量随着节点数对数增长，节点少时效率不错。ring-allreduce则是把节点排成一个环，每个节点只跟相邻的两个节点通信，把整个梯度向量切块后绕环一圈做聚合和分发。它的通信量不随节点数增加而增加（总通信量固定），所以在节点数较多时扩展性更好，这也是NCCL库实际使用的核心算法。

实际使用DDP时，bucket_size参数值得注意。PyTorch会把小的梯度张量打包成桶（bucket），一个桶满了以后再做一次通信，而不是每个梯度单独通信。桶越大，通信次数越少，但每次通信的延迟越高。我实测下来，对于一般规模的模型，bucket_cap_mb设置为25左右是个平衡点；如果模型特别大，适当调大这个值能有效减少通信次数，提升吞吐。

另外要注意，DDP 的后端通常有gloo和nccl两种。GPU集群一定用nccl后端，它深度优化了NVIDIA GPU之间的通信路径；CPU训练或者调试阶段可以考虑gloo，更稳定但性能差一些。用torchrun启动时，确保每个进程的local_rank正确映射到对应的物理GPU上，否则会出现多个进程共享同一张卡导致显存溢出的问题。

2.2 张量并行TP的切分方式与通信开销

张量并行的核心是把一个矩阵乘法拆到多张卡上并发执行。以Transformer中的线性层为例，常见的切法有按行切和按列切。

对于一个前馈层Y = XW，其中X是输入激活，W是权重矩阵。按列切分就是把W按列切成两块W1和W2，分别放到两张卡上，每张卡计算Y1 = XW1和Y2 = XW2，然后再把结果拼接成完整的Y。这个过程在计算时不需要卡间通信，但在拼接结果时需要一个AllGather操作。反向传播时，输入的梯度需要跨卡AllReduce聚合，通信开销就体现出来了。

按行切分则是把输入X按列切块，每张卡持有完整权重的一部分，各自计算部分输出，之后需要一次跨卡的AllReduce来合并输出。这种方式的反向传播不需要额外的通信，因为每张卡只需要更新自己那一份权重切片。

具体选哪种切法，取决于模型的哪一层以及相邻层的切分方式。在多头注意力里面，典型的做法是把QKV三个矩阵按列切分，把注意力头的维度切开，每个设备负责一部分头——这样切完以后注意力计算本身完全不需要设备间通信，只有最后拼接输出时通信一次，非常高效。这也是为什么张量并行在注意力层特别好用的原因。

但TP有个显著的痛点：通信频率太高，每个Transformer层的前向和反向都要做多次AllReduce/AllGather。所以TP的规模一般控制在单节点内部，通常每台机器最多8卡，TP=8是常见的上限。一旦跨节点用TP，网络延迟会拖垮整体性能。

2.3 流水线并行PP的切层调度与气泡问题

流水线并行是按Transformer层切分，把连续的一段层放到一张卡或一组卡上。数据以micro-batch的形式依次通过各层流水段。理想情况下，前一个micro-batch算完第1层，马上可以算第2层，同时后一个micro-batch进入第1层，设备就能像工厂流水线一样持续工作。

但实际上，流水线并行必然面临“气泡”问题——启动阶段和排空阶段会有部分设备空闲。气泡占比 = 流水线段数减一并乘以一个系数再除以micro-batch数量，段数越多、micro-batch数量越少，气泡占比越高。所以在平衡气泡和显存的前提下，要尽量增大micro-batch的数量。实践中，梯度累积步数就可以看作增加micro-batch数量的手段。累积步数越大，流水线填得越满，设备利用率越高。

另一个核心参数是切分均匀性。流水线各段的计算量如果不均衡，比如第一段的Embedding和最后一段的LM Head计算量明显小于中间的Transformer层，那么算得快的段会出现等待，算得慢的段会成为瓶颈。在做层分配时，要根据实际profile数据来调整切分边界，而不是简单地把层数均分。

2.4 ZeRO与显存优化的核心构想

ZeRO（零冗余优化器）可以看作数据并行的增强版。传统数据并行要求每张卡都保存完整的模型状态（参数、梯度、优化器状态），显存浪费非常严重。ZeRO的核心洞察是：这些状态在训练过程中并不需要每张卡都保存完整副本，完全可以分片存储，在需要用到的时候再通过通信收集回来。

ZeRO有三个阶段的递进。ZeRO-1对优化器状态分片，比如Adam里的动量项和方差项，原本每张卡各存一份，分片后每张卡只存1/N。ZeRO-2进一步把梯度也分片，ZeRO-3则把模型参数本身也分片，每张卡只在用到某个参数的时候才临时拉取完整副本。阶段越高，显存节省越多，但通信开销也随之增加。

实际工程里，第一阶段和第二阶段用的比较多，显存节省明显而通信开销增幅可以接受。第三阶段在超大模型场景下很有效，但需要非常仔细地调优通信调度，否则通信延迟会吃掉显存节省带来的收益。

此外，还可以配合激活重计算技术。前向传播时不再保存中间激活值，而是在反向传播需要时重新计算一遍。这等于用算力换显存，通常能让激活值显存占用减少70%左右，是单卡装下大模型的另一把钥匙。激活重计算和ZeRO并不互斥，两者搭配使用效果更佳。

3. 实操过程与核心环节实现

3.1 从零开始设计一个混合并行方案

假设现在要在4台8卡机器上训练一个参数量约65B的模型，每张卡显存80GB。拿到这个需求，我会按下面的路径来做方案设计。

第一步估算参数显存。65B参数，以FP16存储，参数量乘2字节，大概是130GB。再加上Adam优化器状态，每个参数需要额外存储两倍参数的FP32主副本、动量和方差，算下来优化器状态约65B42=520GB。光参数和优化器状态加一起就已经超过650GB，单卡80GB乘以32卡的总显存2560GB看起来装得下，但绝不能全塞进单卡。

第二步确定并行策略。每台机器8卡，先在单机内部用张量并行TP=8，这样每一层的权重被切成8份，单卡上模型参数降为65B/8，约等于8GB。然后把4台机器用数据并行方式复制模型，DP=4，这样整体模型占用显存在可接受范围内。

第三步处理超长序列和激活显存。如果序列长度很长，激活值会非常惊人，这时开启激活重计算，把激活内存压下来。再用流水线并行把层切成4段分布在4个DP组上，这样模型参数进一步降为原来的1/4。整个并行配置就是TP=8，PP=4，DP=1。如果数据并行维度也想保留，可以在多个TP=8×PP=4的副本之间做DP，形成三维混合并行。

混合并行维度的设置不是拍脑袋决定的，而是根据节点拓扑、模型规模、显存容量三层约束反推出来的。实际做的时候我会先跑一个最小的profile实验，量出各层的耗时和通信量，再决定最优切分方式。

3.2 关键参数与Batch Size的计算方法

确定了并行策略后，还要精确计算单卡Batch Size和全局Batch Size的关系。全局Batch Size = 单卡Batch Size × 数据并行度 × 梯度累积步数。注意这里的数据并行度指的是模型副本的数量，而不是所有卡的总数。比如TP=8，PP=4，DP=2的结构下，模型副本数为2，32张卡实际只相当于2个数据并行副本，每条数据要过两个模型副本，所以全局Batch size要按2倍来算。

梯度累积步数也需要匹配。流水线训练中梯度不是每个micro-batch后立即更新，而是累积多个micro-batch的梯度后再做一次优化器更新，这个累积步数直接决定流水线能被填得多满。步数太少则气泡增大，步数太多则会延后梯度更新、影响收敛曲线。

实际设定中，我习惯先从目标全局Batch Size倒推：先设定好总吞吐量，再分配到各维度上。一个常见的经验法则是全局Batch Size不超过训练数据量的1‰，而单卡Batch Size首先受显存限制，一般从1、2、4开始试，能跑到多大跑多大。

另外还需要特别留意学习率与Batch Size的联动关系。增大Batch Size通常要同步线性或平方根地调整学习率，否则收敛速度会明显变慢。在从单卡切换到分布式训练时，这是一个容易被忽略但影响非常直接的坑。

3.3 混合精度训练与损失缩放细节

大模型训练几乎都会用混合精度训练，即FP16或BF16计算，FP32做参数更新主副本。BF16相对于FP16有更大的数值范围，但精度更低，好处是在大部分场景下不需要做损失缩放。FP16则必须处理梯度下溢问题，需要使用动态损失缩放来调整梯度幅度。

在工程落地时，梯度裁剪的顺序要先于损失缩放。先算梯度范数，在FP32的梯度主副本上做裁剪，再把裁剪后的梯度转为低精度格式。如果顺序搞反了，先缩放了低精度梯度再裁剪，梯度范数已经被截断或溢出，裁剪就失去意义了。

选择FP16还是BF16，要看硬件对两者的支持情况。新的高性能GPU普遍原生支持BF16，计算速度更快，数值稳定性也更好，所以新卡集群上我更推荐默认BF16。老一代显卡如果原生不支持BF16，只能用FP16加动态损失缩放的方式。混合精度的损失缩放系数一般从2的24次方附近起步，连续若干步梯度没溢出就适当调大，溢出就跳步并调小，这个逻辑框架要理解清楚，才能在出现异常时快速定位问题。

3.4 训练启动与日志监控的实操记录

用PyTorch生态启动分布式训练，最推荐的方式就是torchrun，它负责管理进程组初始化、分配RANK和环境变量，不要手动去init_process_group里写死地址。启动前需要检查三件事：节点间网络连通性、显卡驱动和通信库的版本匹配、训练脚本中是否有需要置乱随机种子的逻辑。

训练启动之后，最重要的就是日志监控。我一般会关注四个核心指标：吞吐量（每秒处理的样本数）、计算利用率（GPU利用率）、通信占比（通信耗时占总耗时的比例）、显存峰值。其中通信占比这个指标最容易被忽视，如果超过30%，说明并行策略大概率配置不当，需要检查是不是跨节点通信太频繁或者TP规模超出节点边界。

日志监控工具上，单机阶段我用简单的日志加可视化面板就够了；跨集群训练时一定要把训练日志落到分布式文件系统，集中收集到统一的可视化面板上。否则训练跑到一半节点挂了，日志随容器一起蒸发，排障会非常痛苦。

4. 常见问题与排查技巧实录

4.1 显存溢出：边界情况与动态显存管理

显存溢出是最常见的问题，但它不一定是因为模型太大。有一种隐蔽的溢出场景是多进程启动时local_rank映射错误，多个进程被分配到了同一张卡上，导致该卡显存直接爆掉，其余卡空闲。排查方式很直接：训练启动后检查每张卡的显存占用，如果某张卡占用异常高于其他卡，八成就是rank映射问题。

另一种是动态显存峰值问题。显存不是训练全程恒定不变的，在反向传播计算梯度的瞬间，或者在梯度累积和通信的某些节点，显存会出现短暂的峰值。如果静态分析觉得显存够，但训练中频繁OOM，就需要用显存profiler抓峰值时刻，看是哪一步操作触发的。常见解法包括开启激活重计算、减小Batch Size、调整通信桶大小以及分片模型状态。

显存问题还要注意和序列长度的关系。序列长度变长，激活值显存呈线性甚至平方级增长。用FP16训练时，如果输入长度不均，最后一个Batch可能远大于平均长度，这种“长尾样本”往往是压垮骆驼的最后一根稻草。处理办法是按长度分桶并做动态padding，把能耗拉平。

4.2 训练速度不达标：通信瓶颈与数据加载瓶颈

训练速度上不去，很多人第一反应是加卡，但有时候加卡反而变慢。原因很简单：通信开销增长超过了算力增长。要先量化出计算和通信各自的时间占比。如果通信占比高，下一步要判断通信到底发生在哪个环节——是梯度同步的AllReduce，还是张量并行的逐层通信，还是数据加载的IO阻塞。

梯度同步通信占比较高时，先检查梯度累积步数和Batch Size是否过小，再看通信后端是不是被换成了gloo（性能差很多），最后检查网络协议和网络配置，比如把默认的TCP改为RDMA类的高速协议，会有非常明显的改善。

数据加载瓶颈是我见过最容易出现的问题。训练脚本里没有使用DataLoader的多进程读取，或者num_workers设得太小，数据加载成了隐藏瓶颈。分布式训练还有更隐蔽的一个坑：多个进程同时访问同一个数据源时，如果没有做好分片，每个epoch都可能出现数据重复或缺失，影响训练正确性。好的做法是每个进程分配一个独立的数据分片索引，配合合理的num_workers和预取机制，确保数据管道不是短板。

4.3 训练不稳定：Loss突刺与梯度异常

分布式训练里出现Loss突然暴涨或变成NaN，原因往往和单机训练不同。常见的一个原因是梯度累积把多个micro-batch的梯度叠在一起时，并没有做梯度归一化。如果累积步数设了8步，但梯度累积代码里没除以8，梯度范数就会偏大8倍，优化器很容易被推飞。排查时先用小学习率跑通，再逐项排查梯度缩放和梯度累积逻辑。

另一个典型问题是数据并行副本之间的权重不一致。虽然DDP理论上保证所有副本梯度一致同步，但如果用了某些自定义的优化器或者做梯度裁剪时没有在所有进程上同步执行同样的操作，副本间权重就会出现细微偏移。这种偏移不会一上来就爆，但会在长时间训练后累积成明显的精度下降。检测办法是在固定步数打印所有副本的参数均值，看是否一致。

序列长度不等导致的计算图动态性问题也值得记录。当模型支持动态序列长度时，不同micro-batch的计算负载差异很大，在流水线并行里会直接导致各段负载不均衡，某些段长期等待，某些段长期过载。这种场景下尽量固定序列长度，或者按长度分桶调度。

4.4 常见问题速查表

4.5 几个值得养成的实战习惯

第一，训练脚本里从第一步就显式设置随机种子，并且在每个进程上设置不同的偏移。这能保证数据加载顺序在各副本之间错开，同时保证实验可复现，排查问题时也能缩小范围。

第二，养成每个固定步数打印一次“全局Batch序号、Loss、学习率、吞吐量、显存峰值、通信耗时”的习惯。分布式训练的一次迭代可能包含多个micro-batch，只看单卡日志根本不知道全局进度，必须建立全局视角的日志规范。

第三，配置好告警。分布式环境下，任何一台机器掉线都可能导致整个训练任务失败。不要指望训练结束后再人工检查，要设置显存、温度、网络错误、进程心跳的自动化告警，尽早发现问题，避免浪费大量算力。

5. 一个完整的端到端训练方案示例

为了让前面的理论落到实处，我在这里给出一个模拟项目的完整训练方案设计流程，项目代号就叫“模拟项目X”，它是一个基于Transformer解码器结构的大语言模型，参数量约13B，训练数据规模约200B tokens，要在4台8卡节点共32张80GB显卡上进行训练。

节点拓扑是每台机器内有高速卡间互联，跨节点走标准网络。基于这个拓扑，并行方案定为单机内张量并行TP=8，跨节点流水线并行PP=4。这样32张卡被组织成一条4段流水线，每段内部是8卡张量并行。参数显存估算：13B参数在FP16下约26GB，每卡分摊到TP=8后约3.25GB，再配合Adam的优化器状态，每卡总参数相关显存约20GB出头，为激活值和通信缓冲留出了充足空间。

全局Batch Size的计算：单卡Batch Size设为4，数据并行副本数为1，因为整个模型只被复制了一份。如果走梯度累积的话，累积步数设为16，则全局Batch Size为4×16=64。这里没有用数据并行，是因为该模型以TP+PP为主，跨节点复制模型副本的显存代价太高，不如直接拉大流水线的micro-batch数来填满流水线。梯度累积16步下，流水线有足够多的micro-batch来掩盖气泡，预计设备利用率能稳定在85%以上。

混合精度选择BF16，不做动态损失缩放。学习率采用预热加余弦衰减，前2000步从0线性升到3e-4，之后按余弦曲线衰减。为了配合全局Batch Size的变化，在分布式训练中最忌讳固定学习率不动，必须根据有效Batch Size做缩放。

训练中的监控指标以三个为主：全局吞吐（目标超过2000 tokens/s）、单节点通信占比（目标低于15%）、显存峰值（目标不超过75%）。任何一个指标连续若干步不达标，就触发告警。另外每1000步计算一次验证集Loss，一旦验证Loss和训练Loss的差值开始拉大，提前停止训练或者回滚到前一个检查点重新调整学习率调度。

这套方案在实际运行中遇到的问题也挺典型。第一次启动速度只有预期的一半，profiler一看，通信占比逼近40%，原因是流水线切层不均匀——前三段各分到9层Transformer，第四段只有Embedding加1层Transformer，整体等待明显。把层重新切成7、7、7、4加尾部的结构之后，通信占比降到15%，吞吐直接翻了一倍。另一次遇到Loss周期性尖峰，排查了整整一天，最后发现是某个节点的时间同步出了问题，导致梯度累积的跨进程归一化出现偏差。从此之后我所有集群训练脚本都会加一个启动时的时间同步检查。

这类案例再往下展开，涉及的细节就是实践中的“手感”了。分布式训练的框架知识和系统设计可以从文档里学，但如何在具体硬件拓扑、具体模型结构、具体数据分布下做取舍，只能靠一次一次调试和记录积累出来。

6. 几个值得反复回看的经验

分布式训练这个领域，实际上是由大量权衡构成的。没有“绝对最优”的并行策略，只有“当前资源约束下最合适”的并行策略组合。随着模型结构、硬件拓扑、数据规模的变化，同样的配置可能从最优变成劣化，所以训练过程中的持续观察和动态调整非常重要。

有一点我体会很深：分布式训练排障时，第一步永远不是改代码，而是先确认数据流和处理逻辑有没有问题。很多疑难杂症，比如模型Loss不收敛、梯度同步异常、训练吞吐莫名下降，根源往往在数据分片、数据加载和日志收集这些“边缘”环节，而不是核心的并行逻辑本身。把这些基础设施打牢，后面能省下大量排查时间。

后续如果条件允许，这个方向可以继续往深处扩展：一方面可以研究自动并行搜索，让系统根据模型结构和硬件拓扑自动寻找最优的并行切分方案，省去人工试错；另一方面可以关注最新的异构训练框架，把计算、通信、显存三者做更细粒度的统一调度。分布式训练这个领域的演进速度非常快，今天写在笔记里的方案，过一两年可能又有更优解替代，但这种基于第一性原理的分析方法——先看约束、再看策略、最后选方案——始终是核心的决策路径。