大模型分布式训练:三大并行策略与显存优化实战指南
1. 为什么单卡训练大模型注定走不通——从显存账本说起我最早接触大模型训练时脑子里其实一直有个疑问网上那些动辄几十B参数的模型到底是怎么在GPU上塞进去的我手里一张24G显存的卡光是加载一个7B模型的权重就提示OOM更别说跑了。后来自己动手做了一次分布式训练的实验把几张卡拼在一起才恍然大悟——大模型训练从第一天起就不是单卡能干的事。先算一笔账。一个7B模型以FP16精度存储光权重就要占7 * 10^9 * 2 bytes 14GB。这还只是推理时的静态占用。训练时GPU显存里住着的东西可不止权重至少还有四类模型权重本身14GB梯度14GB和权重同尺寸优化器状态Adam优化器要保存一阶动量和二阶动量各一份FP32精度合计约14GB * 4 * 2 112GB这里需要重新算很多人第一次估算优化器状态时都会算错。Adam的m和v都是FP32一份权重对应两份FP32状态每份大小是7B * 4 28GB两份就是56GB。梯度如果用FP16算再占14GB。光这三样加起来就是14 56 14 84GB还没算激活值。激活值这东西才是真正吞显存的大户。序列长度一拉长、batch size一大激活值能轻松超过权重本身的几十倍。比如一个7B模型序列长度2048、batch size 16激活值随便就是上百GB。这也是为什么很多人一跑长文本训练就OOM根本不是网上说的显存不够装模型而是激活值把显存彻底撑爆了。单卡训练一个7B模型的真实显存需求保守估计在100GB以上。哪怕是消费级旗舰卡显存也就24GA100的80G版本勉强能跑H100也只有80G。所以单卡这条路从经济账和技术账上都走不通。分布式训练要解决的核心问题本质上只有一个把一个人扛不动的显存和算力负担拆给多张卡一起扛同时保证大家步调一致、结果一致。至于怎么拆、怎么同步、怎么通信这就是并行策略要回答的事。在往下拆并行策略之前我想先强调一个观念上的转变。很多人以为分布式训练就是把数据分给几个GPU各自训练最后把梯度加一下。这个想法只对了一半——它只是数据并行的一种朴素形态。真正工程级的分布式训练要在三个维度同时下手数据怎么分、模型怎么切、计算流程怎么排。这三个维度对应三种并行策略也是这篇文章的核心骨架。2. 分布式训练的三大基础范式——你可能同时用了它们而不自知分布式训练的并行策略业界一般分成三大类数据并行Data ParallelismDP、张量并行Tensor ParallelismTP、流水线并行Pipeline ParallelismPP。有些资料还会单列序列并行Sequence ParallelismSP和专家并行Expert ParallelismEP但从基础理论角度看先把三大范式吃透后面那些都是它们的变体或组合。2.1 数据并行DP把训练样本摊给每个节点最简单也最常用数据并行的思路最直观每张卡各持一份完整的模型副本把一个大batch的训练数据切成多份发给不同GPU处理。每个GPU用自己拿到的那份数据算前向和反向算出梯度后大家把梯度做个全局平均然后用平均梯度同步更新各自的模型副本。这样做的优点是模型在每张卡上都有一份完整副本实现简单、负载均衡。缺点是显存门槛被死死卡在单卡能不能装下一个完整模型上。我用一张24G卡跑13B模型数据并行连第一步都迈不出去。更麻烦的是通信开销。每步训练都要做一次全量梯度的AllReduce这个操作的通信量等于模型参数量乘以梯度字节数。13B模型在FP16下一次AllReduce就是26GB的通信量100张卡就是每轮2600GB的流量。通信时间甚至能超过计算时间这是数据并行在超大模型上失效的根本原因。后来出现的梯度累积、混合精度训练、梯度压缩这些技巧都是为了减轻数据并行的通信和显存压力但都没法突破显存必须装下完整模型这个天花板。2.2 张量并行TP把网络切块摊给多张卡类似多头并进张量并行是另一个维度的思路不把数据摊开把模型网络本身切开。它针对的是单独一层大矩阵运算比如Transformer里的自注意力计算、MLP层。具体做法是把权重矩阵按行或按列切分成多个子矩阵分配到不同GPU上每张卡只负责计算自己那一块算完后通过AllReduce合并结果。我举个例子你就明白了。一个MLP层里的线性变换Y XW假设W的形状是[4096, 4096]我把它按列切成两块[4096, 2048]分别放到GPU0和GPU1上。输入X两个GPU各持一份各自算X * W0和X * W1得到两个部分结果再通过AllReduce拼接成完整的Y。这个模式看起来数据并行也有点像关键区别在于数据并行的通信发生在整个batch的反向之后每步只通信一次张量并行在每个Transformer层的前向/反向过程中就要通信多次而且通信的是中间激活值而不是梯度单次通信量小但频率高。因此张量并行对互联带宽极其敏感一般只在单机多卡比如同一台机器里的8张A100里用跨机器用TP会导致通信延迟高到难以接受。张量并行的优势在于它实际上降低了单卡显存占用。每张卡只保存模型的1/N权重7B模型切成8份每张卡只需1.75G权重加上激活值和优化器状态也不至于爆显存。缺点是它要求所有GPU同步计算任何一个卡掉队整个训练都得等它所以对集群稳定性的要求极高。2.3 流水线并行PP把网络切成一段段像工厂流水线一样分工流水线并行是第三个思路。它不切数据也不切层的内部计算而是按层把网络切成多个阶段Stage每个阶段由一张或多张卡负责。比如一个24层的Transformer切成4段每段6层分别交给4张卡。输入数据先进入第一段算完前6层传给第二段依此类推到第四段最后得到loss。这种方式的通信量非常小只需要在相邻阶段之间传一次激活值所以理论上可以跨机器部署对网络带宽要求不高。但它有个著名的毛病——气泡Bubble。想象一下流水线刚启动的时候只有第一段在干活后面三段都闲着快结束的时候最后一段在干活前面三段都闲着。启动和收尾阶段的时间浪费就叫做气泡流水线切得越多气泡损失越大理论上限是(PP-1)/(PP-1MicroBatch)。现在的流水线并行已经不搞简单的一个batch走完全程了而是用micro-batch交错调度把一个大batch切成多个小batch流水线里同时跑好几个小batch让各个阶段尽量都保持忙碌。我实际测试过配合micro-batch的流水线并行气泡占比可以压到10%以下。2.4 三种并行范式怎么选我把三种范式整理成一张速查表方便你对照维度数据并行DP张量并行TP流水线并行PP切分对象训练数据单层内部权重网络层间阶段显存压力每卡完整模型不降权重占用权重等比降低权重等比降低通信频率每步一次全局梯度同步每层多次激活值同步仅相邻阶段间低频对带宽要求高极高较低典型应用中小模型扩batch单机多卡超大模型跨机多节点超大模型理解了三种基础范式的差异再看实践中常用的高级混合并行思路会清晰很多。实际上现代大模型训练几乎不会只用某一种范式而是把DP、TP、PP做个三维叠加比如8台机器、每台8卡就可以采用机间用PP、机内用TP、数据维度再用DP的混合布局。这个组合我在后面会详细展开。3. 显存与通信的定量估算——动手之前做算术题把账算明白再上资源做分布式训练最容易犯的一个错误就是一上来就跑模型、调参数结果各种OOM、通信超时、卡利用率低下搞了几天还不知道瓶颈在哪。我的习惯是动手之前先花半天把显存账和通信账算清楚。算清楚了选什么策略、上几张卡、怎么排拓扑答案基本自己就浮现出来了。3.1 用Menins公式估算显存需求业界最常用的显存估算方法是DeepSpeed文档里的Menins公式也常叫训练显存估计。假设模型参数量为M采用Adam优化器、FP16混合精度、ZeRO优化级别不同显存需求差异很大。以7B模型为例M 7e9我按常见的三种情况做计算情况一不用ZeRO原始DP每卡完整副本模型FP16权重:2M 14GB梯度FP16:2M 14GBAdam状态FP32:4M 4M 56GB合计:84GB不含激活值这已经超过了单张A100 80G的极限所以7B模型用原始DP跑单卡必然OOM。情况二ZeRO Stage 1优化器状态切分ZeRO Stage 1把Adam状态按DP rank数N切分每卡只保留56GB / N。假设用8卡权重:14GB梯度:14GBAdam状态:7GB合计:35GB8卡A100 80G可以轻松跑。情况三ZeRO Stage 3权重梯度优化器状态全切分ZeRO Stage 3在每张卡上只保留模型权重的1/N权重和梯度也都切分权重:14GB / N 1.75GB梯度:14GB / N 1.75GBAdam状态:56GB / N 7GB合计:10.5GB这才是ZeRO Stage 3真正的意义——让7B模型跑在24G甚至更小的显存上也成为可能。这也是为什么很多本地部署场景比如用消费级显卡微调大模型选用DeepSpeed或FSDP这类ZeRO方案而不是原始DP的原因。3.2 通信量估算AllReduce到底要传多少数据显存算完了还要算通信。假设你用N张卡做纯数据并行DP每步训练都需要做一次全梯度AllReduce通信量为通信量 2 * (N - 1) / N * 模型参数量 * 梯度字节数用8卡、13B模型FP16梯度26GB代入通信量 2 * 7/8 * 26GB 45.5GB也就是说每训练一步这张卡要向集群里发送45.5GB的数据。如果用的是PCIe 4.0 x16理论带宽32GB/s这一步光通信就要花1.4秒以上。如果单步计算只要0.3秒通信就是4倍于计算的开销——这时候你怎么调学习率都没用瓶颈在通信上。这也是为什么数据并行在跨机部署时需要特别注意机内用NVLink600GB/s级别机间用InfiniBand或RoCE200~400Gbps两者带宽差了一个数量级以上。如果机间的AllReduce频繁且量大跨机数据并行会慢到怀疑人生。3.3 训练时间粗估公式训练时间估算也是账本里的厚账本。已知模型参数M、训练token数T、GPU总算力F每秒浮点运算次数、GPU数量N、MFU模型浮点利用率即实际算力与理论峰值之比训练时间约等于训练时间 ≈ 6 * M * T / (F * N * MFU)这里的6倍是对应前向和反向各2倍FLOPs以及反向重算时梯度约2倍FLOPs实际推导是权重一次乘法和一次乘法转置大致为2倍前向总约3倍前向加上优化器开销后工程上常取6倍或更高。这个公式最大的价值是帮你判断花多少亿卡时才能训完以及反过来给定截止日期倒推需要多少卡。我用这个公式算过一个7B模型的微调微调token 5亿8张A100单卡FP16约312TFLOPSMFU按40%算耗时约6 * 7e9 * 5e8 / (312e12 * 8 * 0.4)约等于21小时。这个估算和我实际跑出来的结果很接近误差在10%以内。有了这个底至少不会出现租了10台机器结果跑两周才发现预算不够的尴尬。3.4 账算完之后的策略判断显存账算完你会立刻知道该上ZeRO还是TP通信账算完你会知道该不该跨机、该用什么互联时间账算完你会知道该租多少卡、该不该换更大batch。这三张账是我做任何分布式训练实验前必做的功课。4. 从DeepSpeed到FSDP——现代分布式训练的工程化形态理论讲完了接下来聊工程。基础范式是思想但工程落地时我们用的几乎都是封装好的框架。目前主流的选择有两个DeepSpeed和PyTorch的FSDP。搞懂它们的底层机制才能避免在配置上照着抄都抄歪。4.1 DeepSpeed的ZeRO优化器到底做了什么DeepSpeed打出的旗号就是ZeROZero Redundancy Optimizer核心思想是消除显存的冗余占用。前面提过原始数据并行下每张卡都持了一份完整模型副本但同一时刻每张卡算的是不同batch的数据权重完全相同这种复制就是显存浪费的根源。ZeRO三个阶段递进ZeRO Stage 1: 切分优化器状态。Adam状态从每卡一份完整变成N张卡各持1/N通信时优化器状态只在小范围内AllGather。ZeRO Stage 2: 在Stage 1基础上梯度也切分。每张卡只保留自己负责的那部分梯度反向算出全量梯度后只AllReduce自己那部分。这样可以进一步降低梯度累积时的显存峰值。ZeRO Stage 3: 权重也切分。每个参数只在某一张卡上存在前向计算时用到哪个参数临时从持有它的卡上AllGather拉过来算完就释放。这使得单卡显存占用随卡数线性下降理论上你甚至可以在单卡上加载一个根本装不下的超大模型。但注意ZeRO Stage 3的通信量其实不低。因为权重要前向、反向各AllGather一次梯度要Reduce-Scatter一次总通信量约为2倍权重大小的数据量乘以训练步数。虽然比TP的通信频率低但依然需要一个相对健康的网络环境。所以ZeRO Stage 3跑在单机8卡上很顺跨机时如果机间带宽不足性能会明显下滑。4.2 FSDP与DeepSpeed的取舍FSDPFully Sharded Data Parallel是PyTorch官方实现的ZeRO Stage 3思路大同小异但有几个差异点值得选型时注意易用性FSDP作为PyTorch原生API和torch生态无缝集成。你用torch写的模型改几行代码就能套上FSDP调试起来路径短。DeepSpeed则需要额外引入引擎配置和HuggingFace Trainer集成度虽高但排错链路会多一环。灵活性DeepSpeed支持ZeRO Stage 1/2/3、Offload把优化器状态或参数卸载到CPU/SSD、甚至支持混合ZeRO TP PP的部署编排。FSDP在最新版本里也支持了CPU Offload和混合精度的组合但在复杂混合并行编排上不如DeepSpeed灵活。通信优化DeepSpeed在一系列通信算子例如用不同NCCL算法处理AllReduce上有较多积累官方在大规模集群上验证过。FSDP的通信模式与大模型训练社区的调优经验也相当丰富。我的经验是如果你的场景是HuggingFace模型 单机多卡 微调FSDP上手最快如果你的场景是从头预训练 跨机大规模集群 需要混合并行编排DeepSpeed更成熟。顺带一提现在很多框架比如Megatron-LM已经内置了DeepSpeed的集成不用自己抠配置。4.3 配置一个简单的FSDP训练循环我贴一段最小可跑的FSDP训练循环伪代码让你感受一下工程门槛import torch import torch.nn as nn from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload, ShardingStrategy # 初始化进程组 torch.distributed.init_process_group(backendnccl) # 模型实例化后直接用FSDP包装 model MyTransformer(num_layers24, hidden_dim4096) model FSDP( model, sharding_strategyShardingStrategy.FULL_SHARD, # 对应ZeRO Stage 3 cpu_offloadCPUOffload(offload_paramsTrue), # 参数可以卸载到CPU mixed_precisiontorch.distributed.fsdp.MixedPrecision( param_dtypetorch.float16, reduce_dtypetorch.float16, buffer_dtypetorch.float16 ) ) optimizer torch.optim.AdamW(model.parameters(), lr1e-5) for batch in dataloader: optimizer.zero_grad() loss model(batch[input_ids], labelsbatch[labels]).loss loss.backward() optimizer.step()实际操作时你还需要处理分布式采样器、checkpoint的保存策略、以及梯度累积的逻辑。这些细节官方文档都有但我必须在下面单独开一节把最容易踩的工程坑都摆出来。5. 分布式训练最容易踩的五个工程坑——我和团队的实战排雷记录理论倒背如流一到真机训练就各种翻车这是分布式训练最常见的剧本。我把自己踩过、以及帮别人排查过的高频坑整理成五类每一类背后都是真金白银的算力时间买来的教训。5.1 卡间通信拓扑引发的掉速NVLink没连上训练速度减半有一次我们在一台8卡机器上做训练发现MFU只有20%怎么调超参都上不去。后来用nvidia-smi topo -m一看8张卡只挂了两条NVLink其余全靠PCIe切换器通信。这样一来TP通信变成走PCIe通信瓶颈直接把算力利用率掐死。排查这类问题我常用的命令就三条# 查看GPU间NVLink拓扑 nvidia-smi topo -m # 查看当前通信速率 nvidia-smi -q -d PCI # 跑一个NCCL测试验证实际通信带宽 all_reduce_perf -b 8 -e 256M -f 2 -g 8NCCL的all_reduce_perf非常重要它能告诉你当前集群的实际AllReduce带宽。如果理论应该到600GB/s的NVLink实测只有几十GB/s那几乎肯定是拓扑或驱动配置问题。我后来养成的习惯是任何新集群上训练前先跑半小时NCCL test把通信基准测出来再谈训练。这一步能省下后面几十小时的痛苦排错。5.2 梯度累积和loss spikes一个微型batch尺寸引发的连锁反应分布式训练里batch size不是由你单卡设的batch size决定的而是由单卡batch size x 卡数 x 梯度累积步数共同决定。很多人以为调大batch size就能提高训练稳定性结果一上来就是64K的global batch训练直接不收敛或者loss突然飙到NaN。原因在于超大global batch会改变优化轨迹。小batch随机性大相当于在损失曲面上的布朗运动能帮模型跳出较差的局部极小值一旦batch大到一定程度轨迹过于平滑更新方向和梯度几乎固定极易陷入尖锐极小值泛化反而变差。业界常用的做法是用学习率线性缩放global batch扩大K倍学习率大致也扩大K倍同时配合warmup前几千步从很小的学习率慢慢爬升让模型适应新的优化动态。我实际执行过的一次loss spike排查最后定位到的问题很朴素梯度累积的optimizer.step()频率没对上导致本该累积64个micro-batch的梯度才更新时间变成了一步一更新。这种情况下loss曲线会出现周期性震荡。修好步数对应关系之后loss曲线立刻回归平稳。所以排查loss spike先别怀疑模型结构先看一眼梯度累积和优化器step的节奏对不对。5.3 混合精度训练下的溢出与舍入错误分布式训练几乎必开混合精度FP16算前反向、FP32存主权重。但FP16的表示范围很有限权重和梯度的数量级相差悬殊时很容易出现梯度下溢underflow——小于FP16最小正常数的梯度被直接变零模型不再更新。我在一次LoRA微调里遇到loss纹丝不动查了整整一个下午最后发现是fp16下梯度过小全部被截断成0。换用BF16Brain Floating Point指数位比FP16多虽有不足但很少出现下溢问题后loss立刻正常下降。现在的卡基本都支持BF16能用BF16就别用FP16这是血泪教训。另一个坑是loss scaling。FP16训练需要动态loss scaling机制框架默认开着的。如果关掉或者scale值设置太激进loss会变成NaN。这个细节在分布式环境下尤其容易被忽视。做分布式训练时不建议手动干预loss scaling把这个活交给框架自动处理最稳妥。5.4 checkpoint保存策略单卡保存不等于全集群保存分布式训练的checkpoint比单卡复杂得多。如果用FSDP模型权重是切分在各卡上的单卡保存的state_dict只是1/N的权重。如果不做特别处理你最后拿到的checkpoint可能根本无法单卡加载。常见做法有两种保存前聚合每张卡把权重AllGather拼成完整模型再保存。缺点是大模型聚合会占用大量显存和带宽可能触发OOM实际中很少直接用。保存分片元信息每卡只保存自己那1/N同时记录分片映射表。加载时按分片还原。DeepSpeed和FSDP都推荐这种方案。我在训练7B模型时长期用的是FSDP的FullStateDictConfig做周期性完整聚合保存高频率用ShardedStateDictConfig做分片保存。聚合保存每N步做一次分片保存每个micro-batch收敛前做。这样既保证恢复时能快速restart又不至于因为聚合太频繁拖垮训练节奏。5.5 卡掉线导致的训练中断远比你想的频繁一个超过24小时的训练任务遇到node 3 rank 5 watchdog timeout这类报错的概率非常大。原因可能是网络闪断、GPU ECC错误、甚至散热问题。分布式训练几乎必须配备容错恢复机制。业界最标准的方案是周期性保存checkpoint 自动重启脚本。我现在的做法是每500步保存一个分片checkpoint训练框架挂掉后由调度器自动拉起新任务恢复到最近checkpoint。为此我还写了一个简单的bash循环检测训练进程退出码非0就自动重启并加载最新checkpoint。如果你用Kubernetes或Slurm这个逻辑可以托管给平台来做。还有一个隐性坑数据加载的随机种子。恢复训练时如果dataloader的shuffle种子没有和工作目录绑定恢复的步数会反复看到相同的数据导致过拟合和训练不稳定。正确做法是把epoch数和步数一并写进seed确保恢复后能断点续数据。6. 混合并行策略的实战编排——三张卡不够三十张卡该怎么排前面聊了三大基础范式和工程实现。现在的问题是当你有几十张卡时到底怎么把DP、TP、PP、ZeRO组合起来这个调度方案直接决定你的训练吞吐量上限。6.1 从2D并行到3D并行层次化编排最常见的集群拓扑是8卡机多台机互联。这时采用机间PP、机内TP、数据维度DP的3D并行方案几乎是黄金组合。为什么这样排机内8张卡通过NVLink高速互联恰好适合TP的高频通信需求。机间带宽相对较低适合PP这种低频通信只在阶段间传激活值。数据并行作为最顶层维度用于扩大整体batch它需要的AllReduce通信量虽然大但可以通过ZeRO的梯度切分缓解。举个例子我配过的一个48卡训练任务6台8卡机每台机内8卡做TP8把一层Transformer的权重切成8份。机器之间PP6把模型24层切成6个阶段每个阶段负责4层。数据维度DP1因为模型已经很大且单阶段显存占用颇高。这个48卡的拓扑用语言描述就是一张卡负责模型的1/48数据不再横向复制每张卡只算自己那1/48的层和1/8的权重。通信路径被压到最低机内走NVLink做高频TP机间走以太网或InfiniBand做低频PP。6.2 显存与通信的权衡菜谱在做编排决策时我习惯画一张权衡菜谱场景推荐配置理由7B微调单机8卡消费级卡FSDP/DeepSpeed ZeRO Stage 2/3TP1模型不大纯度DPZeRO最省事通信开销可接受13B~30B预训练单机8卡TP8 DPN或者TP4PP2单卡装不下全模型TP把权重摊开PP进一步降显存65B预训练多机多卡机内TP8 机间PPN DPK最大化NVLink利用率机间用PP规避低带宽超长序列比如32K在3D并行基础上叠加序列并行SP序列维度的激活值切分可以避免激活值OOM这里尤其要强调序列并行。近年训练长文本大模型逐渐成为刚需激活值内存占用会随序列长度线性暴涨。序列并行的思想是把LayerNorm、Dropout这类沿序列维度操作的部分在序列维度上切分和TP叠加后能显著降低激活值占用。现在主流的Megatron-LM实现里已经内置了序列并行支持用起来不复杂但理解它解决什么问题、和TP怎么嵌套才能调好配置。6.3 我踩过的编排反例TP跨机部署有一次图省事我没有把TP限定在单机8卡内直接把TP纬度跨了两台机器TP16跨机互联。结果训练收敛是能收敛但吞吐量直接掉了40%以上。原因很简单TP的AllReduce在单层内频繁发生任何一次都要经过机间低速互联通信等待时间摊到了每层计算的路径上。后来把TP收回单机改用跨机PP吞吐立刻回升。这件事让我明白TP放在机内是为了让最高的通信频率走最高的带宽通道PP放机间是用时间换空间用低频率通信换跨机扩展能力。这俩的取舍不是随意的背后是通信频率和带宽的匹配逻辑。6.4 并行规模的理论上限Amdahl定律的影子不要无限加并行度。每个并行维度都有它的串行部分。PP的气泡成本随PP深度增加TP的通信占比随TP规模提升DP的梯度AllReduce通信量随卡数线性增长。把这些串行占比套进Amdahl定律你会发现PP深度超过8以后气泡增长开始明显侵蚀加速比。TP超过8以后单层内通信占比过高加速比趋于饱和。DP卡数你受限于总batch sizebatch太大收敛变差。所以行业里大模型预训练的标配组合几乎都是机内TP8 机间PP4~8 数据并行适度扩展很少看到某个维度被拉得极端长。合理配置的目标不是每个维度都最大而是让每个维度都在边际收益为正的区间内。7. 从理论到实验——一个完整的7B模型分布式微调复盘理论聊了不少最后我拿一个真实做过的实验做复盘把前面所有内容串起来。任务是对一个7B基座模型做指令微调训练数据大约5亿token目标是跑进24小时以内。7.1 选型判断过程数据量不算大是微调任务不是预训练所以不需要大集群8卡A100单卡80G就够。7B模型完整权重14GB拿80G单卡去跑原始DP没问题但我不想浪费算力而且未来可能往13B、30B迁所以直接上FSDP的ZeRO Stage 3方案。微调任务要快不需要做大规模TPFSDP单机8卡足够。混合精度用BF16避免FP16下溢出问题。优化器用AdamW8卡ZeRO Stage 3下Adam状态每卡约7GB加上权重分片、梯度分片显存压力很小。主要瓶颈反而是激活值所以把batch per GPU控制在合理范围配合梯度累积把global batch撑到合理值。7.2 关键的训练参数配置model: name: 7B-base train: global_batch_size: 512 micro_batch_size: 4 gradient_accumulation_steps: 16 learning_rate: 2e-5 lr_warmup_steps: 1000 precision: bf16 sequence_length: 2048 optimizer: name: AdamW beta1: 0.9 beta2: 0.95 weight_decay: 0.1 parallelism: fsdp: sharding_strategy: FULL_SHARD cpu_offload: false tensor_parallel_size: 1 pipeline_parallel_size: 1 data_parallel_size: 8 checkpoint: save_interval: 500 sharded: true full_aggregate_interval: 5000这里global batch size 4 * 8 * 16 512。按照前面的学习率缩放经验这个规模下2e-5是个安全的起点。warmup 1000步是为了让模型在超大batch下平稳起步避免前期震荡。7.3 训练过程中的实际观察训练的前500步loss从3.2降到2.1基本符合预期。第1000步左右出现过一次小的loss spike排查后确认是数据加载的sample顺序在恢复训练时没对齐所致修正seed后平稳。全程MFU实测在38%左右和前面用Menins公式估算时假设的40%很接近。实际吞吐约为每步1.2秒5亿token按2048序列长度折算约25万步这里步数偏大是因为token数里多数是填充实际有效训练步更少算下来不到24小时就完成了训练。显存峰值约45GB/卡离80G还有不小余量说明这个配置下还有空间可以加大batch或者序列长度。7.4 这个实验带给我的三个判断第一FSDP的易用性确实名不虚传从原始PyTorch代码迁移到FSDP训练循环前后不到半天。第二而真正花时间的不是框架切换而是把显存账、通信账、batch size缩放规则、checkpoint恢复策略这些基建账算明白。这些决定了你的实验能不能稳定跑完而不是能跑起来。第三对着理论做完备的预计算再上手实操效率远高于先跑起来再说。这大概就是我坚持基础理论笔记一定要动手验证的原因。理论笔记不是纸面功夫是把每次实验里的隐形知识固化下来的地方。