ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

MindSpore Transformers大模型预训练与微调实战:并行策略与显存优化

MindSpore Transformers大模型预训练与微调实战:并行策略与显存优化 1. 项目全貌为什么需要重新思考大模型训练最近一直在折腾 MindSpore Transformers 下的大语言模型预训练和微调从单卡实验一路扩到多卡并行中间踩了不少坑也把分布式并行和显存优化这两块从“知道概念”练到了“能落地”。这篇文章不打算讲空话直接把我的配置、脚本、报错和排查思路整段甩出来给准备上手大语言模型训练的人做个参考。先说清楚这个项目要解决什么问题。大语言模型的参数规模动辄几十亿甚至上千亿单张显卡根本装不下更别说训练了。就算只做微调一张24G显存的卡跑7B模型开个全量微调也是分分钟爆显存。所以“高效预训练与微调”本质上就是两件事把模型拆开放到多张卡上跑把显存里的冗余计算和存储省下来。MindSpore Transformers 在这个场景里是一个很合适的载体它兼容了不少 Hugging Face Transformers 的接口迁移成本比想象中低而且并行策略的原生支持和显存优化手段都比较齐全。这套方案适合谁如果你已经跑过一些小模型的训练比如BERT、GPT-2现在想往更大规模走或者你手头有多张卡但不知道怎么组织并行这篇文章能给你一条可复现的路径。如果你还在犹豫“生成语言模型和大语言模型是一个东西吗”我稍微解释一下生成语言模型强调的是任务形式即按序列生成文本大语言模型强调的是模型规模和能力范式通常是基于海量文本预训练出来的大规模Transformer。两者有交集但不能直接画等号。MindSpore Transformers 里目前主流的LLM模型比如GPT系列、LLaMA系列都是既能做预训练、也能做微调同时还支持后续部署本地推理的。需要提醒一句大模型训练不是一个“装上就能跑”的活儿硬件规模、数据质量、超参设置、并行策略相互牵制。你可能会觉得 7B 模型挺大了但放到几千亿参数的场景里这套思路依然适用只是需要把每个环节抠得更细。这也是我写这篇实战记录的原因——把关键环节拆开逐个讲清楚为什么这么做、参数怎么填、坑在哪儿。2. 分布式并行策略从单卡到多卡的关键路径2.1 三种并行模式怎么选数据并行、张量并行、流水线并行并行策略是一个大模型训练项目最先要决定的事。很多新手一上来就问“我该用几路并行”其实不对你应该先问“我的模型有多大、卡有多少、卡间带宽多少”。模型参数在10B以下一般用数据并行加ZeRO优化就够了模型超过10B单卡哪怕是80G也放不下完整参数这时候才需要张量并行或流水线并行把模型本身切开。数据并行是最直观的每张卡持有一份完整的模型副本喂不同的batch然后通过梯度同步来保持一致。它的通信开销主要是梯度all-reduce卡数越多通信占比越高所以通常配合梯度累积来降低同步频率。MindSpore Transformers 里设置数据并行很简单一般通过parallel_config里的data_parallel维度指定。但数据并行有个硬伤模型参数和优化器状态仍然是每卡一份显存没有本质减少。所以它适合参数不大、但单卡显存刚好能装下模型的情况。张量并行是把一个层内的矩阵运算按维度切到多张卡上比如把hidden_size切成两块每张卡算一半算完再通过通信拼接。这个策略能真正降低单卡显存但会引入大量all-reduce通信对卡间带宽要求很高。NVLink、RoCE 这种高速互联跑张量并行才划算普通千兆以太网做张量并行会慢到怀疑人生。MindSpore Transformers 里通过tensor_parallel设置切分份数一般建议2、4、8这种二的幂而且切分维度必须能被整除否则会直接报shape错误。流水线并行则是按层切分把不同的层放到不同的卡上数据像流水线一样一段段流过各卡。它的通信量最小但存在流水线气泡问题也就是某些卡在等上游数据时的空闲时间。MindSpore 提供了pipeline_stages参数来设置切分点配合微批次micro-batch可以缓解气泡。实际项目里很少只用一种并行基本都是数据并行 张量并行 流水线并行组合使用。我常用的一组配置是 4 机 32 卡把 13B 模型切成tensor_parallel4, pipeline_stages2, data_parallel4也就是每4卡一组做张量并行两组之间做流水线四组之间做数据并行。这样既能压单卡显存又能保持整体吞吐。2.2 并行配置实操以MindSpore Transformers为例MindSpore Transformers 的并行配置并不是散落在各个脚本里的而是集中在run_pretrain.py这类入口脚本里通过参数传入。我以一次7B模型预训练为例给出一份实际可用的配置片段。python run_pretrain.py \ --model_name_or_path llama2_7b \ --data_path ./data/train.jsonl \ --output_dir ./output \ --use_parallel True \ --parallel_mode data_parallel \ --data_parallel 8 \ --tensor_parallel 1 \ --pipeline_stages 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 16 \ --learning_rate 3e-4 \ --weight_decay 0.1 \ --num_train_epochs 3 \ --max_seq_length 4096 \ --save_steps 1000这份配置是典型的数据并行8张卡各持一份完整模型每张卡batch size为4经过16步梯度累积后等效全局batch size是4 * 8 * 16 512。为什么全局batch要512因为大模型预训练对batch size比较敏感太小了梯度噪声大loss波动明显太大了收敛速度不会同比例提升反而浪费算力。512是我在7B模型上实测下来性价比比较高的值。如果模型升级到13B单卡可能放不下我会改成--parallel_mode semi_auto \ --data_parallel 4 \ --tensor_parallel 4 \ --pipeline_stages 2 \这里parallel_mode从data_parallel切到semi_auto意思是让框架自动处理张量和流水线的切分。很多人以为设置完就完事了其实还需要在模型定义里检查每一层的shard配置。MindSpore 的自动并行并不等于“什么都不用管”它只知道哪些算子可以切分但切分的维度是否合理需要你确认。比如nn.Embedding默认只在词表维度做张量并行切分如果你的词表大小不能被tensor_parallel整除会报错。我通常会把词表做个对齐处理或者让框架开启vocab_emb_dp选项来改成数据并行切分。2.3 通信开销与卡间拓扑的取舍并行策略定完还要考虑通信拓扑。张量并行通信量最大必须把对应卡放到同一个节点内走NVLink。流水线并行通信量小一些可以跨节点但节点间的带宽也不能太差。数据并行是梯度级别的all-reduce通信量中等。我在实际测试中发现一个现象同样是4机32卡如果张量并行卡分散到不同机器训练速度会直接掉一半以上。原因很简单跨机的网络延迟和带宽远不如机内NVLink。所以做张量并行时第一件事是确认卡分组逻辑。MindSpore 在启动分布式时会根据rank_table文件划分卡组你可以通过npu-smi info查看卡所在节点然后把[0,1,2,3]这类连续四卡分给同一张量并行组。如果用的是昇腾环境还会涉及HCCL通信库的配置一定要和组网方式匹配否则可能出现通信初始化超时。另一个经常被忽略的点是梯度累积与通信频率。数据并行每步都要做梯度同步但是可以等累积到若干个micro-batch后再同步降低总通信次数。MindSpore 里的gradient_accumulation_steps就是干这个的。有人会问梯度累积等于增大了batch为什么不直接调per_device_train_batch_size因为单卡显存往往是瓶颈你没办法一次塞进512条样本只能通过累积来绕过显存限制。这个思路本质上是用时间换空间在算力约束下提升模型训练能力的一种常见资源配置方式。3. 显存优化实战把每一兆显存都榨干3.1 显存都去哪了训练态内存组成很多人对显存优化的理解停留在“开混合精度、开梯度检查点”这种操作层面但不知道这些手段分别作用于哪一块显存。搞清楚显存花在哪你才能对症下药。一次大语言模型训练中显存主要由四部分构成模型参数、梯度、优化器状态、激活值。模型参数就是权重本身7B模型用FP16存储大概是14GB。梯度同样大小也是14GB。优化器状态最夸张如果用AdamW每个参数要维护一阶动量、二阶动量和参数副本在混合精度下这又是好几倍的开销。光这三样7B模型的单卡需求已经超过40GB。激活值则是前向传播时每层产生的中间结果它和序列长度、batch size、层数强相关序列越长越恐怖。之前跑一个4K序列长度的训练激活值一度占了总显存的三分之一。所以显存优化不是一个技巧的事而是一整套组合拳。目标很明确把模型参数、梯度、优化器状态、激活值这四块的冗余都压下去。3.2 ZeRO分级、重计算和混合精度的组合拳ZeRO 是解决参数、梯度、优化器状态冗余的主流手段。它把这三类状态按一定策略分片到多张卡上每卡只存一部分用的时候再通过通信聚合。ZeRO-1只切优化器状态ZeRO-2把梯度和优化器状态都切了ZeRO-3进一步把模型参数也切了也就是每卡只持有模型参数的一个分片。MindSpore 对应的接口是zero_level我实测下来数据并行场景用zero_level2性价比最高因为只改优化器状态和梯度的分布通信增量可控zero_level3虽然能把参数也省掉但每次前向反向都要做参数收集通信开销明显上升小规模集群不划算。如果你的模型大到单卡确实放不下ZeRO-3 配合 offload 才能勉强训起来。重计算是另一大杀器。它把前向传播的激活值丢掉反向传播时用一版临时参数重新算一遍激活。典型的重计算配置是选择性地对部分Transformer层开启而不是全开。全开虽然最省显存但训练时间可能增加30%以上。我习惯的做法是每隔两层开启一层重计算或者只对序列比较长的模块开。MindSpore 里通过recompute配置控制你可以指定recompute_layers的编号列表。混合精度是默认操作FP16或BF16把显存直接减半。要注意的是用FP16时容易出现loss异常因为梯度值太小会被截断。解决方案是开dynamic_loss_scale让loss缩放值自动调整。BF16虽然没有精度截断问题但对硬件有要求昇腾和部分新GPU都支持。我在MindSpore Transformers 上跑起来后基本固定用BF16 AdamW 动态loss scale的组合稳定性和收敛速度都更好。3.3 Offload到底什么时候开Offload 是把优化器状态或参数从显存挪到CPU内存甚至磁盘上常见的有optimizer_offload和param_offload。很多人一听到能省显存就无脑开结果训练速度惨不忍睹。原因是CPU与GPU/昇腾设备之间的传输带宽远低于显存带宽每步都要把状态搬来搬去总耗时自然暴涨。我的经验是只有当显存缺口在20%以内时offload值得开如果缺太多说明并行策略本身就不合理应该先把tensor_parallel加上而不是靠offload硬撑。比如单卡24G想跑7B全参训练就算开offload也很勉强因为模型参数14G、梯度14G、优化器状态超大这些开销不是offload能救的。但如果你已经开了ZeRO-2只差一点点显存激活值装不下那开optimizer_offload把优化器状态挪到CPU通常能解决问题性能损失也还能接受。显存优化的另一个偏门技巧是micro_batch_interleaving把一个小batch再切成更小的micro-batch顺序执行减少单次持有的激活值。这和梯度累积不同它不改变batch size只是把一次前向/反向切碎。效果上有点像“显存版的梯度检查点”但副作用少一些。MindSpore 里在流水线并行场景下经常使用能显著降低激活峰值。4. 预训练与微调的完整实操流程4.1 环境搭建与数据准备先从环境说起。我用的是昇腾环境MindSpore 版本 2.2MindSpore Transformers 直接pip install mindspore-transformers安装依赖会自动带上mindspore核心库。为了在VSCode里调试模型代码我还装了 MindSpore 内核扩展可以在Jupyter里直接跑训练脚本断点调试体验比纯命令行好不少。首次启动时如果报缺少CANN包记得按设备类型安装对应的CANN工具链这个属于环境基础跳过的话连设备都识别不到。数据准备阶段需要把文本序列化成语料。我用的是JSONL格式每行一个样本包含text字段然后通过分词器转成input_ids。数据清洗特别重要大语言模型预训练很吃数据质量重复度高、噪音多的语料会让loss下降变慢甚至训出一个“复读机”。我一般会做三步去重、过滤超短文本、按照长度排序然后分桶保证每个batch内的样本长度接近减少padding浪费。4.2 预训练脚本参数逐项说明预训练脚本的参数很多但核心就几个模型路径、数据路径、并行配置、学习率、序列长度、保存间隔。前面已经给过一份完整配置这里重点解释容易被坑的参数。max_seq_length决定序列长度直接影响激活值显存。很多人想一步到位用8192结果OOM。建议先跑一个短序列验证数据流和模型逻辑再逐步拉长。save_steps别设太频繁7B模型每次保存都要写几十GB权重频繁保存会严重拖慢训练。我通常每1000步保存一次同时保留最近两版权重避免磁盘爆满。学习率调度也不能忽视。预训练一般用warmup cosine decaywarmup步数设为总步数的1%~2%。MindSpore 支持在训练脚本里传lr_schedule_type和warmup_ratio。注意learning_rate指的是峰值学习率和batch size强相关。如果全局batch翻倍学习率也应适当上调不然收敛速度会变慢。4.3 微调阶段的技术要点微调和预训练虽然代码相似但策略完全不同。预训练追求的是在巨量数据上稳定收敛微调则是在特定任务上快速适配所以学习率要小得多一般1e-5到2e-5之间。全参微调对显存压力极大所以很多场景用LoRA这类参数高效微调方法。MindSpore Transformers 里也有LoRA支持通过peft配置指定lora_rank16、lora_alpha32只训练注入的低秩矩阵可训练参数量少到可以忽略但效果往往出乎意料地好。如果你要微调的模型已经通过预训练产生了权重注意加载路径要对应正确。MindSpore 的权重文件可能是.ckpt格式而Hugging Face上公开的权重是.bin或.safetensors格式。MindSpore Transformers 提供了一些转换脚本但转换时经常遇到不同的key命名。这个环节比较耗时好在社区有预转换好的权重可以直接下载省去很多麻烦。微调数据也要按任务调整。指令微调需要构造“指令输入输出”的结构一条样本通常包含多个轮次用特殊token分隔。我在构造Alpaca格式时会在系统提示、用户指令、回答之间加上### Instruction、### Response这类分隔符。分隔符的选取会影响模型输出格式不能随意。4.4 配置继承与命名冲突的坑这里要专门说一个坑在迁移Hugging Face模型配置到MindSpore Transformers时如果直接把原来的config.json扔过去偶尔会报一个aimv2 is already used by a transformers config, pick another name.之类的错误。意思是某个模型名或配置名已经被已有的Transformers配置占用了导致加载时命名冲突。我遇到的情况是自定义了一个模型结构但没有修改model_type字段结果和内置的模型注册表冲突。解决方案很简单在config.json里换一个不冲突的model_type同时确保模型注册时的类名一致。还有一个更隐蔽的版本兼容问题MindSpore Transformers 更新很快旧版本保存的config.json里某些字段可能在当前版本中已经被移除或改名加载时会抛警告甚至直接报错。我的习惯是每次升级库之后先跑一个加载模型的冒烟测试确认权重能正常读进去再启动正式训练。不要嫌这一步多余它能帮你省下半天排查时间。5. 踩坑实录与排查清单5.1 常见错误一览先整理一份常见错误速查表都是我实际跑训练时撞过的场景。错误现象可能原因处理方式启动时卡在初始化HCCL通信组网失败或rank_table不匹配检查卡分组逻辑和rank_table文件训练几小时后OOM激活值增长的内存碎片问题开启重计算或减小batch sizeloss为NaN或持续不降学习率过大、FP16丢失精度换BF16、开动态loss scale、调低学习率模型加载时命名冲突config中的model_type与内置配置重叠更换model_type名称微调评测效果和预期差很多数据格式不正确或LoRA超参不合理检查指令格式、调整lora_rank训练速度低于预期大量使用跨机通信或offload过多调整并行策略把张量并行收敛到机内这其中的前三项是新手最容易碰到的。遇到问题不要慌按“环境、数据、模型、超参”的四个维度逐层排查通常能很快定位。5.2 显存溢出OOM排查OOM是最常见的故障但“显存不够”不一定是同一个原因。我先看日志是在哪一步爆的如果在前向传播阶段爆多半是激活值太多优先考虑重计算或减小batch如果在优化器更新阶段爆多半是优化器状态太大优先考虑ZeRO或offload如果是加载权重时爆那就是模型参数本身太大需要上张量并行。还有一个容易忽视的情况显存碎片。训练过程中每次前向/反向都会分配和释放不同大小的临时块长时间运行后显存碎片越积越多明明总显存够用但这块新申请的大块内存分配不到最终OOM。这种问题在序列长度不固定时尤其明显。解决方法是尽量固定序列长度或者定时重启训练从checkpoint续跑把显存布局重置一下。5.3 分布式训练卡死的定位思路多卡训练一旦卡死先看卡状态。用npu-smi info查看每张卡的利用率如果某张卡一直是0而其他卡在动大概率是数据加载不均或者同步逻辑出了问题。另一个常用手段是打开通信日志MindSpore 有HCCL的详细日志级别开启后能看到每个通信原语的收发状态。但我建议只在排障时开开着日志跑大规模训练会拖垮性能。还有一种“静默卡死”很烦人程序不报错但loss不动、时间不再推进。多半是异步数据加载把CPU占满了主线程拿不到数据训练循环空转。检查数据集的num_parallel_workers和预读取队列长度必要时降低并行度给训练主线程留出足够的CPU资源。5.4 效果相关loss不降、评测波动如果loss不降先看是不是数据问题。文本质量差、标签错位、特殊token处理错误都会导致模型学不到东西。其次是学习率太大会让loss震荡太小会让loss下降极慢。评测波动则是另一个头疼的事。同一个checkpoint在不同评测集上表现不同并不一定是模型退化了可能是解码参数不一致比如temperature、top_p或者max_new_tokens设置不同。我每次评测都固定一套生成参数并把随机种子固定住这样才能在不同checkpoint之间做公平比较。6. 从实验到上线的经验总结6.1 合理规划实验的迭代节奏大模型训练最贵的是时间最怕的是方向错了还在猛跑。我在实际项目中养成了一个习惯任何大规模训练之前先拿一个小数据集、小模型配置跑通全流程。比如用7B模型但只训500步确认loss在下降再启动正式训练。这一步看起来浪费时间实际上能拦住大量配置错误。尤其是并行切分配置如果tensor_parallel设置不合适小规模试跑时就能暴露shape不匹配的问题不用等到几十个小时之后才报错。如果要以较小的成本去验证不同并行策略和显存优化的组合效果我的建议是固定batch size只调整并行度和重计算参数看单位时间吞吐是多少。记录下每个组合下的训练吞吐再做选择。尤其当你准备换更大的模型时参考这些数据会很有帮助因为它们直接对应了卡间通信和显存压力的变化趋势。6.2 后续扩展建议这套方案目前主要覆盖文本类大语言模型但同样的并行和显存优化思路可以直接迁移到视觉大语言模型上。视觉大模型的前半部分是视觉编码器后半部分是LLM唯一的变化是输入变成图像特征序列显存占用会更依赖分辨率。如果你已经跑通了文本LLM的分布式训练再上手视觉模型时只需要调整数据管线和序列构造部分。另外如果你的目标不只是训练还想在本地部署大语言模型做推理训练阶段采用的并行策略会直接影响推理时的部署方式。训练时用张量并行切开的模型权重在推理时也必须按同样的切分方式加载否则会出现权重不匹配。好在 MindSpore Transformers 可以保存带有并行信息的分片权重部署时直接用分片权重启动服务即可。最后再分享一个个人心得显存优化和分布式并行永远不是一次到位的事情每次换模型、换数据、换硬件都可能需要重新调参。但核心原则是不变的——先搞清楚瓶颈在哪再决定用什么手段。把所有技巧都开一遍看起来很厉害实际上可能只是在浪费算力。保留一份最简配置能让你在遇到突发问题时光速回到基线剩下的空间再一点一点挤出来。
返回列表