大语言模型训练这件事真正上手之后你会发现最难的往往不是模型结构本身而是怎么把一堆参数塞进有限的显存里同时让多张卡协同干活还不互相拖后腿。MindSpore Transformers 这套框架我用了一段时间它在分布式并行和显存优化上给出的方案比较务实不像有些框架把配置藏得很深。这篇内容适合已经跑通过单卡小模型、想往多卡大模型方向走的同学也适合手上有几张卡但不知道怎么榨干算力的朋友。我会从并行策略的选择逻辑讲起把数据并行、模型并行、流水线并行这几条路各自的适用边界说清楚再落到显存优化的具体手段上包括重计算、优化器分片、混合精度这些能直接改配置就见效的操作。中间会穿插我自己踩过的坑比如切分维度选错导致通信量爆炸、重计算粒度没调好反而拖慢训练速度这类问题。1. 为什么大模型预训练绕不开分布式并行1.1 单卡放不下模型时的真实困境先算一笔账。一个 7B 参数的模型如果用 FP32 存储光权重就要 28GB 显存。训练的时候还要存梯度、优化器状态Adam 的话每个参数要存一阶矩和二阶矩加起来轻松超过 100GB。这还没算激活值。一张 80GB 的卡根本扛不住更别说消费级的 24GB 卡了。很多人第一反应是那我用混合精度不就行了。FP16 确实能把权重和激活砍一半但优化器状态通常还是 FP32 的省下来的空间有限。而且混合精度会带来数值稳定性问题loss 突然变成 NaN 是家常便饭。所以光靠精度压缩解决不了根本问题必须从并行维度想办法。分布式并行的本质是把一个大问题拆成多个小问题分给不同的设备去算再通过通信把结果拼起来。听起来简单但拆的方式不同通信开销和显存占用差别巨大。选错了策略可能卡越多越慢。1.2 三种并行策略的适用边界数据并行Data Parallel是最直观的每张卡存一份完整模型喂不同的数据算完梯度再 AllReduce 求平均。它的优点是实现简单、通信模式规整。缺点是每张卡都要存完整模型显存占用不降反升因为要额外存梯度通信的缓冲区。所以数据并行只适合模型能塞进单卡的情况一般 1B 以下的模型才考虑。模型并行Model Parallel是把模型本身切开不同的层或者同一层的不同部分放到不同卡上。这样每张卡只存一部分权重显存压力直接下降。但问题是切分之后前向和反向传播需要跨卡通信如果切得不好通信量会大到把计算时间全吃掉。张量并行Tensor Parallel是模型并行的一种把矩阵乘法按维度切开适合 Transformer 里的注意力层和 FFN 层。流水线并行Pipeline Parallel是按层切分把模型的不同层放到不同卡上数据像流水线一样依次流过。它的通信量比张量并行小但会引入气泡——也就是某些卡在等数据的时候闲着。气泡的大小取决于流水线的级数和微批次的数量。实际训练大模型时这三种策略通常是组合使用的。比如 8 卡训练一个 13B 模型可能用 2 路张量并行加 4 路流水线并行再叠加数据并行。MindSpore Transformers 把这套组合配置抽象成了几个参数改起来不算复杂但理解每个参数背后的含义很重要。2. MindSpore Transformers 的并行配置怎么落地2.1 并行维度参数的设置逻辑MindSpore Transformers 里控制并行的核心参数是parallel_config它包含data_parallel、model_parallel、pipeline_stage这几个字段。我一开始以为这些数字随便填就行后来发现它们之间有约束关系data_parallel × model_parallel × pipeline_stage必须等于总卡数。比如你有 8 张卡可以配成 2×2×2也可以配成 4×2×1但不能配成 3×3×1。选哪个组合取决于模型大小和卡间带宽。如果卡间是 NVLink 这种高带宽互联张量并行的通信开销可以接受可以多切一点。如果是普通的 PCIe张量并行的 AllReduce 会很慢这时候应该优先用流水线并行把通信量降下来。我自己的经验是先看模型能不能用流水线并行切开。Transformer 的层结构天然适合按层切所以pipeline_stage设成 2 或 4 通常没问题。然后再看单层能不能塞进一张卡如果单层太大比如 FFN 的隐藏维度特别大就需要张量并行来切。数据并行放在最后考虑因为它不解决显存问题。2.2 张量并行切分维度的选择张量并行最关键的决策是切哪个维度。以矩阵乘法 Y XW 为例W 的形状是 [in_features, out_features]。你可以按行切切 in_features也可以按列切切 out_features。这两种切法的通信模式完全不同。按列切的话每张卡算出部分输出最后需要 AllReduce 把结果加起来。按行切的话每张卡需要完整的输入输出直接拼接就行不需要 AllReduce。但按行切要求每张卡都存完整的输入显存占用会高一些。MindSpore Transformers 在 Transformer 层里默认用的是列切加行切的组合注意力层的 QKV 投影按列切输出投影按行切。这样前向传播只需要一次 AllReduce反向传播也只需要一次。如果你自己改切分方式很容易搞成需要多次通信性能直接掉一半。提示改张量并行切分方式之前先用小模型跑一遍用 profiler 看一下通信时间占比。如果通信时间超过计算时间的 30%说明切分方式有问题。2.3 流水线并行的微批次调优流水线并行的气泡问题很烦人。假设你有 4 个流水线阶段每个阶段处理一个批次需要时间 T那么第一个批次走完整个流水线需要 4T但后面三个阶段在第一个 T 时间里是闲着的。这就是气泡。减少气泡的办法是增加微批次数量。把一个大批次拆成多个微批次让它们像流水线一样重叠起来。微批次越多气泡占比越小。但微批次太多也有问题每个微批次的数据量变小计算效率会下降而且通信次数变多。MindSpore Transformers 里用micro_batch_num控制微批次数量。我的经验值是流水线级数的 2 到 4 倍。比如 4 级流水线微批次设成 8 或 16 比较合适。设成 4 的话气泡还是明显设成 32 的话每个微批次太小GPU 利用率上不去。这里有个容易忽略的点微批次数量必须能被全局批次大小整除。如果你全局批次是 32微批次设成 5那就对不上了。我踩过这个坑报错信息还不明显查了半天才发现是整除问题。3. 显存优化的几个实操手段3.1 重计算用时间换空间的经典操作重计算Recompute也叫梯度检查点Gradient Checkpointing思路是在前向传播时不保存中间激活值反向传播时重新算一遍。这样显存占用能降很多代价是计算量增加大约 30%。MindSpore Transformers 里开启重计算很简单在配置里设recomputeTrue就行。但粒度可以调可以只对注意力层重计算也可以对所有层重计算。只对注意力层的话省下来的显存有限但速度影响小对所有层重计算的话显存省得多但速度慢。我一般先试只对注意力层重计算如果显存还是不够再扩大到 FFN 层。实测下来7B 模型开注意力层重计算能省大约 20% 显存全开能省 40% 左右。但全开之后训练速度会掉 25% 到 35%这个取舍要看你的卡数和时间预算。还有一个细节重计算的粒度可以按层设置。比如前几层不开重计算因为它们激活值小后几层开。MindSpore Transformers 支持通过recompute_slice_activation参数来控制但这个参数文档里写得不太清楚我是看源码才搞明白的。3.2 优化器状态分片ZeRO 思路的落地优化器状态是显存大户。Adam 优化器每个参数要存一阶矩和二阶矩加起来是参数量的 2 倍。如果模型有 7B 参数优化器状态就要 56GBFP32。这还没算梯度和权重。ZeROZero Redundancy Optimizer的思路是把优化器状态、梯度、甚至权重切分到不同卡上每张卡只存一部分。MindSpore Transformers 通过optimizer_shard参数来开启这个功能。开启之后优化器状态会按数据并行的维度切分每张卡只存 1/N。但这里有个坑优化器状态分片之后更新参数时需要 AllGather 把状态收集起来通信量会增加。如果卡间带宽不够训练速度会明显下降。所以这个功能适合卡间带宽高、显存特别紧张的场景。我试过在 PCIe 互联的机器上开这个速度掉了 40%后来还是换回了普通的数据并行。3.3 混合精度与损失缩放的配合混合精度训练用 FP16 或 BF16 来存激活值和梯度用 FP32 存权重的主副本。这样显存能省一半左右计算速度也能提升因为很多 GPU 的 FP16 算力比 FP32 高。但 FP16 的数值范围窄梯度容易下溢变成 0。解决办法是损失缩放Loss Scaling把 loss 乘以一个大的系数反向传播得到的梯度也相应放大更新之前再除回来。MindSpore Transformers 里用loss_scale参数控制可以设成固定值也可以用动态缩放。动态缩放更省心它会根据梯度是否溢出自动调整缩放系数。我一般用动态缩放初始值设成 2^16。如果训练过程中频繁出现溢出日志里会有提示说明模型数值不稳定可能需要检查初始化方式或者降低学习率。BF16 的数值范围比 FP16 宽不容易溢出但精度低一些。如果硬件支持 BF16比如 Ascend 910 系列优先用 BF16可以省掉损失缩放的麻烦。4. 训练过程中的性能调优与排错4.1 通信瓶颈的定位方法多卡训练变慢十有八九是通信出了问题。定位方法是用 MindSpore 的 profiler 工具它能记录每个算子的执行时间和通信时间。跑几百步之后导出数据看 AllReduce、AllGather 这些通信算子的时间占比。如果通信时间占比超过 40%说明并行策略有问题。常见原因有几个张量并行切分方式不对导致通信次数过多流水线并行的微批次太少气泡太大或者卡间带宽本身就不够。我遇到过一次通信特别慢的情况查了半天发现是张量并行的切分维度搞反了。本来应该按列切的地方按行切了导致每次前向传播要多做一次 AllReduce。改过来之后训练速度直接翻倍。这个教训是改并行配置之后一定要跑个短训练验证一下别等跑了几小时才发现不对。4.2 显存溢出的排查链路显存溢出OOM是大模型训练的常客。排查的时候不要瞎猜按这个顺序来第一步看是哪个阶段溢出的。前向传播溢出通常是激活值太大反向传播溢出通常是梯度或优化器状态太大。MindSpore 的报错信息里会提示是哪个算子出的问题顺着算子找对应的层。第二步算一下理论显存占用。权重、梯度、优化器状态、激活值这四块分别占多少。如果理论值就超了那必须改并行策略或者开重计算。如果理论值没超但实际超了可能是显存碎片问题试试设置max_device_memory参数限制单卡显存。第三步逐步缩小范围。先把批次大小降到 1看还溢不溢出。如果不溢出了说明是批次太大需要调小或者增加梯度累积。如果还溢出说明是模型本身太大必须用模型并行。我踩过的一个坑是开了重计算但没开优化器分片结果反向传播的时候优化器状态还是把显存撑爆了。这两个要配合使用单开一个效果有限。4.3 训练不收敛的常见原因分布式训练不收敛原因往往比单卡训练更隐蔽。我总结了几种情况梯度同步问题。数据并行下如果 AllReduce 没正确执行不同卡上的梯度不一致模型就会乱跑。检查方法是打印每张卡的梯度范数看是否一致。如果不一致检查通信组配置。学习率没随批次大小调整。数据并行下全局批次大小等于单卡批次乘以卡数。如果学习率还按单卡批次设相当于学习率变小了收敛会变慢。一般规则是学习率随全局批次大小线性缩放但超过某个阈值后要改成平方根缩放。数值精度问题。混合精度下如果损失缩放系数设得不对梯度会下溢或者溢出。看日志里有没有overflow或underflow的提示。有的话调整损失缩放参数。数据加载问题。多卡下如果数据没正确分片每张卡拿到相同的数据相当于批次变小了。MindSpore 的 Dataset 会自动分片但如果你自己写了数据加载逻辑要确保每张卡拿到的数据不重叠。5. 从预训练到微调的衔接5.1 预训练权重的加载与适配预训练跑完之后权重保存成了 checkpoint。微调的时候加载这个 checkpoint但要注意并行策略可能不一样。预训练可能用了 8 卡微调只有 4 卡这时候需要做权重的重新切分。MindSpore Transformers 提供了load_checkpoint接口它能自动处理并行策略变化。但有个前提checkpoint 里要保存并行策略的元信息。如果没保存加载的时候会报维度不匹配。我建议在保存 checkpoint 的时候把parallel_config也存进去省得后面麻烦。另一个坑是预训练用的模型结构和微调用的结构必须一致。如果你在微调时改了层数或者隐藏维度权重就加载不进去了。这种情况只能加载部分权重或者重新训练。5.2 微调阶段的显存优化策略微调比预训练省显存因为批次大小通常小很多序列长度也可能短一些。但微调有自己的问题如果做全参数微调优化器状态还是很大如果做 LoRA 这类参数高效微调显存占用能降一个数量级。LoRA 的思路是冻结原模型权重只训练两个低秩矩阵。这样优化器状态只跟低秩矩阵有关显存占用大幅下降。MindSpore Transformers 支持 LoRA配置里设lora_config就行。我实测下来7B 模型做 LoRA 微调单张 24GB 卡就能跑起来批次大小还能设到 4。但 LoRA 的效果取决于秩的选择。秩太小模型学不到东西秩太大显存优势就没了。一般从 8 开始试如果效果不好再往上加。我试过秩设成 64效果和全参数微调差不多但显存占用还是比全参数微调低不少。5.3 微调后的模型合并与导出LoRA 微调完之后低秩矩阵需要和原权重合并才能导出成标准模型。MindSpore Transformers 提供了合并接口但要注意合并时的精度。如果原权重是 FP16低秩矩阵是 FP32合并的时候要统一精度否则会有数值误差。合并之后的模型可以导出成 MindIR 或者 ONNX 格式方便部署。导出的时候要注意算子兼容性有些自定义算子可能不支持导出。我遇到过一次导出失败原因是用了 MindSpore 特有的算子后来换成标准算子才成功。导出之后建议做一次推理验证确保合并后的模型输出和微调时的输出一致。我一般会跑几条测试数据对比一下 logits 的差异。如果差异很大说明合并过程出了问题。6. 一些容易被忽略的工程细节6.1 数据管道的吞吐量匹配多卡训练时数据管道的吞吐量必须跟得上计算速度。如果数据加载成了瓶颈GPU 会经常空转。MindSpore 的 Dataset 支持多线程加载和预取但默认参数不一定适合你的场景。我一般会把num_parallel_workers设成 CPU 核心数的一半左右prefetch_size设成批次大小的 2 到 3 倍。如果数据需要做复杂的预处理比如 tokenization可以考虑提前处理好存成二进制文件训练时直接读省掉预处理时间。还有一个细节数据分片要均匀。如果每张卡拿到的数据量不一样会出现有的卡算完了在等有的卡还在算。MindSpore 的 DistributedSampler 会自动做均匀分片但如果你自己写了采样逻辑要确保每个 epoch 的数据能被卡数整除。6.2 checkpoint 的保存频率与恢复大模型训练动辄几天几周checkpoint 保存策略很重要。保存太频繁I/O 会成为瓶颈保存太少一旦中断损失太大。我的经验是每几百步保存一次同时保留最近几个 checkpoint防止某个 checkpoint 损坏。MindSpore Transformers 支持异步保存也就是保存 checkpoint 的时候不阻塞训练。这个功能很实用但要注意磁盘写入速度。如果磁盘慢异步保存的队列会积压最后还是会影响训练。建议用 SSD 存 checkpoint或者存到网络存储但确保带宽足够。恢复训练的时候除了模型权重还要恢复优化器状态和学习率调度器的状态。如果只恢复权重优化器状态从零开始训练会有波动。MindSpore 的 checkpoint 默认会保存这些状态但如果你自定义了训练循环要确保把这些都存下来。6.3 日志与监控的配置分布式训练的日志很乱每张卡都往标准输出打混在一起看不清。建议把每张卡的日志写到单独的文件用 rank 编号区分。MindSpore 支持通过环境变量控制日志级别和输出位置。监控方面除了 loss 和准确率还要关注几个指标每步的耗时、通信时间占比、显存占用、GPU 利用率。这些指标能帮你快速定位性能问题。我一般用 TensorBoard 记录这些指标训练过程中随时能看到趋势。如果发现某张卡的显存占用明显高于其他卡说明并行策略可能不均衡。比如流水线并行下第一级和最后一级的显存占用通常比中间级高因为要存输入和输出。这种情况可以通过调整层分配来平衡把计算量大的层放到显存充裕的卡上。7. 我实际用下来的一些体会MindSpore Transformers 这套东西配置项确实多但每个配置项背后都有明确的工程考量。我刚开始用的时候喜欢抄别人的配置结果经常跑不起来。后来强迫自己把每个参数的含义搞清楚反而顺利多了。并行策略的选择没有标准答案得根据你的硬件条件来。卡多带宽高可以多用张量并行卡少带宽低优先流水线并行。显存优化也是重计算和优化器分片不是开得越多越好要权衡速度和显存。还有一个体会是小规模验证很重要。改完并行配置之后先跑几十步看看能不能跑通、速度怎么样别一上来就跑完整训练。我吃过这个亏配错了参数跑了一晚上第二天发现 loss 根本没降。最后说一个实用技巧MindSpore Transformers 的配置支持继承和覆盖。你可以写一个基础配置然后针对不同场景写子配置只覆盖需要改的字段。这样管理起来清晰也不容易出错。我现在都是这么干的预训练一套配置微调一套配置共用大部分参数。