昇腾大模型训练迁移实践：从GPU到国产算力的完整攻略

去年我们组接到一个任务把一套在通用GPU集群上跑通的千亿参数模型训练方案完整迁移到国产昇腾平台上并且训练吞吐不能掉得太多。这是典型的信创场景业务侧的诉求一句话就能说清——模型必须能在国产算力底座上完成训练和推理。但真正动手才发现迁移的工作量远不止“换一台机器重新装环境”那么简单。这个领域有一个关键认识昇腾不是一个“长得像GPU的加速卡”而是一整套以昇腾处理器为核心的计算平台。硬件有训练服务器和加速卡软件有底层驱动、图编译框架、算子库、集合通信库再往上是MindSpore框架和PyTorch的适配层。整套体系的组织方式和CUDA生态接近但实现细节差异很大。这篇内容我就以亲历的迁移过程为线索把昇腾大模型训练的软件栈、迁移路径、分布式配置、排障经验和调优思路完整梳理一遍给正准备在国产算力平台上做大模型训练的同学一条可落地的路线。1. 从GPU到昇腾迁移动机、平台组成与真实难度1.1 信创场景里最常见的三种驱动因素我在实际接触过的项目中“为什么要做昇腾适配”通常有三个层面的原因。第一是硬件采购侧的国产化替代要求新建项目明确要求算力底座选用国产方案第二是运行环境的私有化交付需求有些客户的数据不能出域训练和推理都要跑在客户指定的国产服务器上第三是成本与供应侧的考量在特定阶段国产加速卡的可获得性更稳定等待周期更短。不管出于哪种原因落到工程层面只有一个结果原本在CUDA生态里写成的那套训练代码要在昇腾上重新跑起来还要跑得好。很多人初听这件事会觉得“不就是换个后端吗”这是最大的误解。PyTorch的代码确实可以做到一套代码多后端运行但前提是算子必须都有对应实现。大模型训练涉及的算子数量动辄几百个其中任何一个关键算子没有适配整个训练就会卡住或异常。更麻烦的是通信策略、显存管理、算子融合这些底层机制CUDA和昇腾完全是两套实现不是改两行配置就能平移过来的。1.2 昇腾平台到底由什么组成先把平台结构讲清楚。昇腾训练平台从底往上大体分为四层。最底层的硬件是昇腾AI处理器为核心的加速卡和训练服务器。单卡集成了AI计算单元、向量计算单元、矩阵计算单元等异构核同时自带高带宽的片上存储。训练服务器通常是一台机器插多张加速卡通过高速总线组成一个节点多个节点之间再通过高速网络组成训练集群。硬件之上是CANNCompute Architecture for Neural Networks这是昇腾的计算架构层地位相当于CUDA。CANN包含驱动、运行时、图编译引擎、算子库、集合通信库等组件。模型代码最终会通过CANN编译成可以在昇腾处理器上执行的指令序列。再往上是深度学习框架适配层。现在主流的两条路线是原生路线用MindSpore兼容路线用PyTorch加Ascend Extension for PyTorch通常叫torch_npu。torch_npu做的事情相当于一个翻译层把PyTorch的算子调用翻译成昇腾能执行的算子。最顶层才是算法工程师熟悉的模型代码和训练脚本。到这一层迁移工作才开始。1.3 评估迁移工作量先算算你的模型踩中多少算子我在评估一个模型能不能快速迁移时通常先做一次粗糙的算子覆盖扫描。方法是把模型跑一遍推理打开CANN的算子落盘日志统计哪些算子命中了昇腾的高性能算子库哪些走了自定义实现路径哪些根本没有对应的NPU实现。这个扫描结果基本决定了迁移难度。如果模型是Transformer类架构用的都是标准算子如矩阵乘、LayerNorm、Softmax、注意力计算那迁移路径非常顺畅。如果模型里有大量自定义CUDA算子、自定义反向传播、特殊的索引逻辑那就要做好自己写算子或者改造模型结构的准备。我们迁移的这个模型主体是标准的稠密Transformer架构但内部有几个自定义的稀疏注意力算子。就是这几个算子在后续的适配过程中消耗了最多时间。2. 昇腾训练环境的硬件组网与软件栈先搞清楚任务跑在哪一层2.1 从单机八卡到跨机集群训练组网方式昇腾训练集群的基本单元是一台训练服务器常见配置是单机八卡。八张卡在节点内部通过高速互联总线组成一个完整拓扑卡与卡之间可以直接通信不需要经过CPU内存转发。这个拓扑结构对并行策略设计影响很大通常建议先通过工具查看卡间的物理连接拓扑再做张量并行或者数据并行的切分位置。跨机训练时节点之间通过高速网络互联比如RoCE或InfiniBand方案。每台服务器上会有若干网口用于跨机通信同时还有管理网口用于系统运维。组网规划里最容易犯的错是把训练网络和管理网络混在一起导致跨机通信跑不满。真实训练时通信数据量极大一个千亿模型的梯度同步动辄几十GB级别网络带宽不足会直接让集群吞吐崩掉。单机八卡适合百亿规模以下模型的预训练和大部分微调场景千亿级模型的预训练基本都要跨机。我们可以用一张简表说明模型规模常见并行方式组网需求十亿级单机单卡到单机四卡可不组网百亿级单机八卡数据并行为主单节点内部互联即可千亿级多机多卡张量并行加流水并行跨机高速网络必须到位这里要提醒一下跨机训练时除了计算卡本身你还要规划好存储系统和数据加载路径。训练数据通常放在共享存储上每个节点拉取自己的数据分片。如果存储的并发吞吐不够训练过程会频繁出现“等数据”的状态加速卡的利用率会掉下来。2.2 CANN软件栈编译、算子、运行时三层CANN把模型从图到指令经历了几大步理解这几步对调优很有帮助。第一步是图编译。PyTorch动态图或MindSpore静态图传入后CANN会先把模型表达成计算图然后做算子融合、冗余消除、内存复用等优化。算子融合是最重要的一环典型例子是把LayerNorm里的多个小算子合并成一个融合算子减少访存和算子调度开销。大模型的训练性能很大程度上取决于这些融合做得怎么样。第二步是算子实现。昇腾的算子库已经覆盖了绝大多数常见算子矩阵乘、卷积、归一化、激活函数这些都有深度优化版本。没有覆盖到的算子需要走自定义算子开发路径一种方式是用Ascend C语言手写算子。我们有同事前期写了一个稀疏注意力算子的原型前后花了两周才把性能调到可接受水平。第三步是运行时调度。CANN负责把编译好的任务下发到板上管理显存分配、多流调度、事件同步。PyTorch的显存管理策略和CANN不同迁移时偶尔会遇到显存碎片化或者申请失败的问题后面排坑章节细说。2.3 框架适配层PyTorch如何被“翻译”成昇腾算子指令PyTorch在昇腾上跑靠的是torch_npu这个适配层。它的工作原理很直接通过PyTorch的私有的后端扩展机制把tensor上的算子操作转发到CANN的算子接口上。训练脚本里只要增加一行import torch_npu并且把模型和计算数据通过.npu()方式放到昇腾设备上就能实现设备端的切换。但这里有个陷阱适配层不是全算子的。如果你在模型里用了某个相对生僻的PyTorch算子而这个算子在torch_npu的适配表里不存在程序不会直接编译失败而是可能出现两种行为——要么抛异常要么该算子在CPU上执行数据反复在CPU与板卡间拷贝。后一种情况最坑人因为训练看起来在继续但速度会慢到让人怀疑人生而且这种慢通常很难一眼发现。排查方法也简单跑一小步训练看日志里有没有警告信息或者用性能分析工具看算子耗时凡是耗时异常高的算子多半就落在CPU上了。我们当时有个模型里用了某个自定义采样函数就是被静默放到了CPU执行每步训练多出好几秒开销定位这个过程花了不少时间。MindSpore那边的逻辑不太一样。MindSpore是昇腾的原生框架图编译和算子调度天然整合得更紧密不需要经过适配层。但MindSpore的生态和社区资源相比PyTorch还是有差距部分开源模型的官方实现只提供PyTorch版本。所以现实中的大多数项目还是会优先走torch_npu路线这也是本文重点讲的方向。3. 模型迁移实操从算子盘点到大模型跑通3.1 迁移前的准备算子清单、依赖清单、环境确认实操第一步不是改代码而是清点。我们把迁移准备工作分成三份清单。第一份是算子清单。把模型的forward和backward涉及的算子全部列出来重点标注自定义算子、动态shape相关的算子、依赖第三方实现了的算子。现在大模型代码普遍基于Transformer架构多数场景扫一遍就能确定适配等级。第二份是依赖清单。项目里用了哪些Python库、哪些版本比如datasets、tokenizers、safetensors这些常见依赖以及是否用到了一些与设备绑定的库。一个典型的坑有些可视化监控库会隐式调用CUDA属性查询换成昇腾环境后在导入阶段就报错需要在生成的依赖文件里提前剔除。第三份是环境清单。确认CANN版本、PyTorch版本、torch_npu版本三者之间的兼容关系。昇腾的版本矩阵比CUDA那边更严格版本对不齐会出现各种诡异问题例如算子编译不过、分布式初始化卡死。我个人的建议是先固定一套官方文档里明确验证过的组合再去动模型代码不要上来就把所有东西都升到最新版。3.2 两种适配方式的对比与选择根据项目情况迁移路径大体可以分成两种典型方式。方法A是“PyTorch原版代码加torch_npu适配”核心工作是改脚本里跟设备相关的部分然后逐个处理不兼容的算子。这种方式优点是对代码改动小迭代速度快能借住PyTorch已有的生态组件。缺点是需要依赖适配层的算子覆盖度遇到性能瓶颈时排查链条比较长。方法B是“模型结构迁移到MindSpore”配合MindFormers这类大模型套件。优点是原生的算子融合和并行策略实现更完整训练性能的上限通常更高而且从编译到通信整个链路问题定位更直接。缺点是代码改动量大模型重新用MindSpore实现一遍工期长。我们项目选择的是方法A核心原因是时间紧且模型的逻辑复杂度高重写一遍风险太大。但这里要说明白方法A并不是简单跑通就完事。跑通只是开始后续的性能调优、算子替换、融合优化工作量不比重写小。3.3 权重转换与数据流替换大模型预训练或微调的初始化往往要加载已有的模型权重而这些权重大多是在GPU上用PyTorch的torch.save存的格式是.pt或.bin。昇腾侧加载这些权重理论上直接torch.load就能读取因为权重文件的格式和硬件无关。真正需要注意的是加载后的存储布局和设备放置。例如一些旧权重文件里把张量默认放在CPU上加载后需要显式执行.npu()迁移。如果有使用accelerate这类库要把它的设备映射配置改成适配昇腾设备名的写法。很多代码会写死cuda字符串这也是迁移中最高频的报错点建议全局搜索代码里的cuda、device相关逻辑统一改掉。数据流这边重点看数据加载器。如果原项目用的是PyTorch原生的DataLoader通常改动不大。但若用了一些专用加速库做数据预处理和缓存就要确认库是否在昇腾环境可用。一个真实例子我们项目的数据管线里用了一个把样本预编码成二进制文件的缓存组件该组件依赖特定的GPU张量运算来做数据增强迁移后这部分被整体替换成了普通Python预处理逻辑。3.4 迁移后的第一轮冒烟验证代码迁移完成之后不要直接上多卡训练。必须在单卡上做一次小规模冒烟验证三步走。第一步跑一个极小的batch比如两三条数据确认前向和反向能完整执行完loss能正常算出来。第二步对比同一份数据的loss数值比如原GPU环境和昇腾环境各跑几个step观察loss曲线是否符合预期。第三步才是开多卡跑分布式先从单机两卡开始再单机八卡最后跨机。为什么要这么保守因为分布式训练里一旦通信、拓扑、并行切分有问题报错信息往往千奇百怪很难直接就定位到根因。先用单机小规模把基础链路跑稳后面排查问题的范围会小很多。我们当时跨过单机直跨两机实验结果一个通信超时问题折腾了整整两天后来退回单机才发现其实就是设备序号没有对齐。4. 大模型分布式训练并行策略、显存优化与超参配置4.1 并行策略的搭配DP、TP、PP怎么组合大模型训练最核心的问题是模型参数、梯度、优化器状态、中间激活这四类数据加起来的占用远超单卡显存。分布式训练本质上解决两件事放得下算得快。围绕这两个目标形成了三套基础并行策略。数据并行DP最简单每张卡都放一份完整模型训练数据切成多份喂给不同卡每轮迭代后梯度做一次AllReduce同步。适合模型的参数总量不大、单卡能放下的场景。张量并行TP是把模型某一层的大矩阵按行或按列切分到多张卡上每张卡只算自己负责的矩阵块前后通过通信算子拼接结果。这种并行策略能有效降低单卡显存占用但每层的矩阵运算都有额外的通信开销通信量和切分粒度密切相关。一般建议把TP的度数限制在单机或同一交换机下的少数几张卡内跨机TP通信延迟会急剧上升。流水线并行PP是把模型按层切成若干段每张卡负责其中一段。数据按微批次流水式地在段与段之间传递。好处是通信量远小于TP坏处是存在流水线气泡即部分卡在等上一段算完利用率有损失。PP的段数、微批次大小对吞吐影响很大需要实验调优。实际项目里几乎不会只用一种并行策略。经典组合是DPTPPP混合比如一个规模为64卡集群可以组织成8路数据并行、4路张量并行、2路流水并行的组合总并行度就是8乘4乘2等于64卡。配比的选择没有万能公式主要参考模型结构和集群拓扑。一般来说矩阵运算密集的层多TP收益大模型很深PP收益大数据量大DP收益大。我们的实践是从一个偏保守的组合开始比如TP优先选择4PP选择2然后通过profiling观察各阶段时间占比再调整。4.2 显存优化的几件套重计算、梯度累积、Offload在很多场景下并行策略还不足以让一个巨大的模型完整放进显存。于是有了显存优化的经典三板斧。第一是重计算也叫激活检查点。把模型划分为多个段训练前向时不保存每层的中间激活只保存分段边界上的检查点。反向传播需要中间结果时再重新执行一次前向去计算。效果是显存占用大幅下降代价是训练时间多出约20%到40%的额外计算量。这是在无法扩容时最常用的手段我们千亿模型迁移后的第一版配置就开了重计算。第二是梯度累积。当单卡能处理的batch太小导致梯度噪声过大时可先累积若干个微批次的梯度再把累加后的梯度统一更新参数。它不直接减少显存占用但能稳定大batch训练效果同时减少频繁同步梯度的次数。第三是Offload机制。把优化器状态或部分参数临时搬到CPU内存计算时再搬回板卡。昇腾和GPU一样板卡显存和CPU内存是独立的offload能大幅降低显存需求但CPU和板卡之间的拷贝带宽有限。Offload调度不好训练会卡在数据搬运上。所以一般只在显存捉襟见肘的极端场景用或者配合梯度累积一起用。4.3 混合精度与Loss Scaling大模型训练现在都是混合精度即用FP16或BF16表示大多数计算数据。BF16相比FP16有更大的指数范围对大模型训练更稳定。昇腾对BF16的支持是完整的这点没遇到阻碍。混合精度训练的关键是loss scale。前向传播后计算出的loss需要放大再反向传播防止梯度在反向过程中下溢变成0。动态loss scale会周期性检查梯度是否有数值溢出自动调整缩放倍数。迁移过程中如果发现loss突然变成NaN或者训练loss维持在某个值不动首先要怀疑loss scale的设置。另一个精度相关的坑是随机种子。不同框架底层随机数生成算法不同即使设置相同的seed初始权重和训练过程的采样结果也不会完全一致。所以做完迁移后不要拿原GPU环境的结果做精确对比只能观察收敛趋势和量级是否一致。4.4 一份千亿参数模型的关键训练配置样例下面给出一份在我们实际项目中验证过的配置模板大致的参数含义做了说明。这不是标准答案只是一个可供参考的起始点。配置项建议值说明模型规模1000亿参数稠密Transformer需要跨机训练集群规模4节点每节点8卡共32卡单机拓扑优先TP并行策略数据并行4张量并行4流水并行232卡总量对应32卡微批次大小1梯度累积16步结合重计算将单卡显存压到可用范围精度策略BF16混合精度动态loss scale稳定性优先重计算开启检查点间隔4层显存与算力的折中优化器AdamWoffload关闭显存够用时不开启学习率3e-4余弦衰减大模型常见起点这个配置在实测中吞吐不是最优的但它是我们排完所有坑之后第一个能稳定跑起来的起点。性能调优是在这个基础上逐步进行的。5. 实测排坑十六次“看起来挂了”和三次真挂了5.1 设备ID错位与多进程初始化卡死分布式训练启动第一轮我们撞上的问题就是卡死。多进程任务在每张卡上启动一个进程每个进程按local_rank分配设备。torch_npu环境下的设备序号排列方式和CUDA环境略有不同尤其在一台机器插了两类加速设备时映射关系很容易错位。症状很典型每个进程都正常启动日志输出到了一半然后整个任务僵住不报错也不退出。用系统工具查看时发现有一部分进程始终停在初始化状态。排查过程先确认每个进程拿到的设备ID和实际物理卡ID是不是一一对应打印进程rank和设备ID然后手工在单机上只启动两个进程做最小复现。最后定位到问题出在启动脚本里用了一个默认的设备映射函数这个函数在CUDA环境下能正确返回设备号在torch_npu下返回的却是空值。换成官方推荐的环境API之后问题解决。这类问题让我养成了一个习惯所有分布式启动脚本先在单机双卡上验证再上多机。5.2 精度对齐的连环坑随机种子、Loss Scale与权重加载顺序模型能跑起来之后更大的挑战是“loss不下降”。我们曾遇到过连续十轮迭代loss纹丝不动的情况。一开始怀疑是学习率设置问题把学习率调大后loss直接变NaN才意识到跟学习率无关。排查链路先查权重加载确认预训练权重有没有正确加载再查数据加载顺序确认每个rank拿到的数据分片不重复也不缺失最后查了混合精度相关配置发现是loss scale策略需要被显式指定否则某些算子内核路径上梯度被截断成了0。另一个容易忽略的细节是权重加载顺序。多卡场景下不同进程如果各自独立加载了初始化权重时用的随机种子不一致会导致同一模型在不同卡上初始化状态不一致。即使后面同步了梯度模型的某些内部状态比如BatchNorm的统计量也可能会出现偏差。大模型一般不用BatchNorm但LayerNorm内部计算涉及缩放因子如果权重加载顺序里某个缓冲区没有被正确同步loss曲线也会异常。5.3 通信超时跨机带宽和拓扑的双重影响训练规模从单机八卡扩到双机十六卡时我们遇到了通信超时的问题。现象是训练跑了十几个step之后随机某个进程报通信超时整个训练任务退出。一开始怀疑是网络不稳定但排查硬件链路、交换机端口、光模块状态都正常。后来用集合通信测试工具实测跨机全互联带宽发现实际带宽远低于网络标称值。进一步排查发现多机通信的最佳路径依赖卡与网口的物理拓扑配置不当会走很长的绕路路径。解法很吃亏但很有效调整网卡中断绑定的CPU核心让负责通信的进程占用的CPU插核与相应网卡本地核心一致这样跨机时少了一次跨NUMA访问。另一个改进是把AllReduce分组把跨机的大通信拆成多次小的分组通信减少单次通信对网络的压力。这一轮调完之后跨机训练才真正稳定下来。5.4 判断“真挂了”与“假挂了”进程级日志与Profiling的使用方法训练跑着跑着有时候你怀疑它挂了但它其实只是在做一个特别耗时的算子。早期我们有个阶段单步训练偶尔要十几分钟从日志看什么都没有输出团队成员开始怀疑程序卡死。后来我们用profiling工具记录了一整个step的所有算子执行时间才发现问题不在卡死而是某个重计算模块触发了重复前向而这个重复前向内部还有一个没有被融合的稀疏算子单次执行要数秒钟多次串联起来就把整个step拖到十几分钟。这里给一个经验判断训练是否健康别只看日志有没有输出要看每个算子阶段的耗时是否在预期范围内。一个可靠的办法是定期保存performance记录观察两个相邻checkpoint之间的平均step时间曲线。如果曲线平滑即使某个step慢一点也不用担心如果曲线出现断崖或持续爬升才是真正要排查的信号。6. 性能调优从“能跑”到“跑得快”6.1 先看Profile而不是猜原因跑通不是终点信创项目验收时对性能是有要求的。我们在拿到一份可稳定运行的版本后开始做性能调优。第一步永远是做性能分析确认时间到底花在哪里。昇腾提供了多种性能分析工具可以输出算子的执行时间、通信时间、空闲时间占比。我拿到一份完整报告后通常先看三个数字计算密集算子矩阵乘、卷积总耗时占比是否够高通信时间占总时间比例是否超过预期是否存在明显的空闲等待如果通信占比偏高优先优化并行配置和通信算子如果计算占比偏低大概率存在过多的未融合算子和CPU回退。这两个方向的优化手段完全不同不先看数据就乱调事倍功半。6.2 关键性能指标与评估方法衡量昇腾大模型训练性能我常用的指标有三个第一个是真实吞吐即每秒处理的样本数或token数。这是最直观的指标。第二个是单卡算力利用率可以理解为MFU模型浮点运算利用率越高说明计算资源越接近上限。大模型训练的MFU受并行策略、通信占比、算子融合度共同影响业界优秀水平多数在40%到60%这一档。第三个是有效训练时间占比用纯计算时间除以总时间把等待数据、等待通信、等待梯度同步的时间全排除掉。我们会建立一个看板记录不同版本配置下这三个数字的变化。调优本质上是持续改进这三个数字的过程而不是简单地把某个超参数调大调小。6.3 几个能直接用的调优手顺基于实际经验我总结出一套适合大模型场景的调优手顺顺序很重要别跳。第一步优化算子融合。用性能分析工具找出耗时最高的非融合算子通过开启自动算子调优功能或用融合算子替代。这几个改动往往收益最大我们实测有接近两成的性能提升来自这一项。第二步调并行配置。在固定集群拓扑下尝试不同的TP和PP组合。通信和计算的交叉点不同会有明显性能差异。我们最后锁定的组合比初始方案快了约一成。第三步调数据加载与预处理线程数。把数据加载时间压到step总时间的5%以内。这一步很容易被忽略但它直接影响整机利用率。第四步开启通信与计算重叠。利用多流机制让梯度通信和下一轮前向计算同时进行隐藏通信延迟。第五步调超参。比如微批次大小、梯度累积步数、重计算粒度。这些参数要在前四步做完之后再调否则调了也不准。6.4 最后一件事把“稳定复现”变成工程能力性能调到目标范围之后还有一个最容易忽视的问题训练的确定性。大模型训练动辄几周中间任何一次中断都可能造成之前全部工作无效。我们后来把整个训练链路做成了可复现的工程包括固定的启动脚本、每次训练前自动检查算子兼容性、定期保存可恢复的训练快照、以及一套监控告警脚本。根据个人经验给准备做昇腾大模型训练的同学几条实在的建议。第一第一次做迁移时预留2到3倍的预期工期算子兼容问题和通信调优永远比你想的费时间。第二尽量先用小规模数据把完整链路跑通再上大规模。第三时刻关注平台版本之间的关系环境升级要在独立环境验证完再推到训练集群。这件事没有什么巧劲就是靠一个个细节堆出来的稳定。