今年我们接了个有点棘手的活儿:把线上一个日调用量千万级的NLP推理服务,从国际主流的GPU平台平滑迁到国产AI加速卡上。任务听起来不复杂——换个硬件,模型重新导出一下不就行了?实际上从立项到灰度上线,我们整整折腾了两个月,中途数次想放弃,问题一个接一个:图编译直接崩、算子不支持、推理精度飘了三个多点、性能上不去只能跑 GPU 的一个零头。这篇就把我们踩过的坑和最终沉淀下来的迁移思路完整复盘一遍,希望那些同样被“算力迁移”和“推理适配”夹在中间的团队能少走一些弯路。先说结论:从 GPU 迁到国产算力,本质上不是“改几行代码”的活,而是一次完整的软硬件适配工程。模型结构、算子映射、精度策略、显存管理、部署方式,每一个环节都可能埋雷。但好消息是,这个过程的坑是结构性的、可预测的,只要按对顺序一步步排查,绝大多数问题都能提前暴露并解决。1. 迁移前的算力摸底:先别急着改代码我们最初犯的第一个错误,就是拿到新卡之后立刻把原来的 PyTorch 模型部署脚本跑了一遍,结果当然是一堆看不懂的报错。后来跟某国产算力厂商的 FAE(现场技术支持工程师)聊了两次,才意识到迁移工作的第一步根本不应该发生在跑代码的环境里,而应该在会议室里完成一次完整的算力实测和算子梳理。1.1 为什么算力选型直接决定后面要踩多少坑不同国产 AI 加速卡之间的差异,远比不同型号的 NVIDIA GPU 之间大得多。NVIDIA 的卡从 1080 到 A100 再到 H100,至少 CUDA 生态是统一的,算子库、推理引擎、量化工具链都是同一套体系,迁移成本主要在显存和计算精度上。但国产加速卡这边,每一家的指令集、片上存储结构、算子实现方式都不同,有的卡对 PyTorch 的兼容性做得比较好,有的卡则要求模型先转成特定中间格式才能运行。这里的核心矛盾不是“跑不跑得动”,而是生态成熟度和算子覆盖率的差异。我们在选型时列了一个简单的评估维度:评估维度具体问题评估方式算子覆盖模型里的算子(Conv、LayerNorm、Softmax、MatMul等)新卡支持几个用模型结构解析脚本统计算子类型分布,再对照新卡的算子支持列表精度支持是否支持 FP16 / BF16 / INT8,以及量化的激活函数支持情况看厂商提供的精度支持文档,实测小模型跑一遍显存容量峰值显存占用是否超过卡上内存在 GPU 上先测模型峰值显存,留出 30% 余量框架对接PyTorch 是直接适配还是要转 ONNX用最简模型跑通测试工具链成熟度是否提供模型转换工具、性能分析器、调试工具询问 FAE,翻开发者社区我们最终选了一款显存和算力规模接近现有 GPU 七成水平的国产 NPU 加速卡,理由很简单:它的 PyTorch 适配层做得比较成熟,我们的模型不需要额外改网络结构就能完成图编译。1.2 模型算子分布:一张图看出迁移工作量在决定适配方案之前,我用一个简单的 ONNX 解析脚本把模型的算子分布统计了一遍。这里不是让大家必须转 ONNX,而是建议用这种方式快速了解自己的模型里到底有哪些“稀有算子”——比如一些自定义的 Attention 实现、奇怪的激活函数、或者是用了torch.einsum这种会被编译器展开成复杂子图的写法。import onnx from collections import Counter model = onnx.load("model.onnx") node_types = Counter(node.op_type for node in model.graph.node) total = sum(node_types.values()) for op, cnt in node_types.most_common(): print(f"{op:20s} {cnt:5d} {cnt / total * 100:.1f}%")我们模型里当时的算子分布大概是:MatMul 占 35%、LayerNormalization 占 18%、Softmax 占 12%、Gelu 占 15%,剩下的是 Reshape、Transpose、Concat 这类数据搬运算子。表面上看都是标准算子,应该没什么问题。但这恰恰是最迷惑人的地方。标准算子名不代表标准实现,比如 LayerNorm 在一张卡上可能是用 CUDA kernel 直接算的,在另一张卡上则可能被编译器展开成一堆乘加和开方操作。一旦展开,精度特性就不一样了,性能也完全不是一回事。这个领悟是我们后续所有排查思路的出发点。2. 部署与适配的第一个坎:算子编译报错与替换策略环境搭建阶段比预想中耗时更久,光是“跑通一个示例模型”就花了近一周。这个过程中遇到的坑,大致可以分为环境配置类、图编译类和算子缺失类三种。环境配置类的坑网上一搜一大堆,这里只说两个我们反复踩的。2.1 环境依赖的版本锁死:换个卡连 Python 版本都得重新对齐国产算力的软件栈一般分为三层:驱动、运行时、框架适配插件。我们最开始按照厂商文档在服务器上装了最新版驱动,结果死活无法加载推理运行时,反复报driver mismatch。后来发现是依赖的某个系统库版本跟预装的 CUDA 兼容层冲突——即使你完全不使用 GPU,某些国产加速卡的驱动安装包仍然会尝试跟 CUDA 的运行时组件做版本检查。解决办法是彻底清空服务器上的 NVIDIA 驱动和 CUDA,再装国产卡的运行时。这里的关键教训是:迁移到国产算力的机器,最好做一次系统级的环境初始化,不要指望跟原有的 GPU 环境共存。如果确实有混跑需求,用容器隔离是最稳妥的方案,而且容器里也要严格指定基础镜像版本。2.2 图编译失败:动态 Shape 是最大的敌人图编译环节,我们遇到的核心报错是shape inference failed。原因是模型里有两处torch.where混着动态 mask 输入,导致张量形状在运行期才能确定。NVIDIA 的 TensorRT 通过显式设置 optimization profile 也能处理一部分动态形状,但国产加速卡的编译器对动态 Shape 的容忍度低得多,很多直接要求模型输入是静态的。这一块没有捷径,唯一的办法是在模型导出阶段把动态维度固定。# 导出时固定序列长度,避免动态 shape dummy_input = ( torch.randn(1, 128, hidden_size, dtype=torch.float32).to(device), torch.tensor([[1] * 128], dtype=torch.int64).to(device), ) torch.onnx.export( model, dummy_input, "model.onnx", input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes=None, # 关键:关闭动态轴 opset_version=17, )我们当时的做法是统计线上服务的实际请求长度分布,发现绝大多数请求集中在 64~128 token 之间,于是统一填充到 128,超出部分走一个回退逻辑。这样做虽然牺牲了一点响应速度,但换来的是模型在国产卡上的稳定编译。另外,如果模型里有if分支控制流,比如if seq_len > 64: use_flash_attention,图编译器会很难处理。建议在导出前把这种运行期条件分支全部去掉,统一走一条计算路径。2.3 算子替换:不支持的算子怎么处理做完上一步,图编译能跑通了,但日志里出现警告:有两个算子类型回退到了 CPU 执行。一个是torch.fft,我们用来做某种频域特征增强的;另一个是torch.linalg.eigh,用在后处理阶段。CPU 回退意味着这部分计算量会从加速卡上搬到 CPU,对于推理性能是毁灭性的。我们尝试了两种替换策略:数学等价替代:torch.fft.rfft用在特征提取,跟一个固定卷积核的 Conv1d 在效果上近似,于是直接把 FFT 换成几个 Conv1d,模型效果几乎不变,但算子完全落在 NPU 上。拆分替代:torch.linalg.eigh用于对 2x2 矩阵求特征值,这种小矩阵直接用解析公式手写。手写之后不仅避免了 CPU 回退,速度还快了 3 倍。这里有个经验:大多数不支持的算子都可以通过数学变换绕过去,关键是你要对自己的模型逻辑足够清楚。如果实在绕不过,就写自定义算子融合进厂商的算子库。写自定义算子的学习成本很高,建议只在算子确实无法替代且直接影响正确性时才考虑。3. 精度漂移排查:那些“差不多”的数字害死人模型适配完、在国产卡上能正常出结果之后,我们进入精度验证阶段。先在测试集上跑了一遍,发现整体准确率比 GPU 低了 3.5 个百分点。这个差异足够大,绝对不能直接上线,但问题是:它到底从哪一层开始错的?3.1 逐层输出比对:定位误差起始层我的排查方法非常简单原始:在模型的关键中间层都插上print,同时对 GPU 和国产卡跑同一个输入,逐层比对输出张量的均值、方差和最大绝对误差。这个方法虽然笨,但绝对有效。第一次排查发现,从第 4 个 Transformer 层开始输出差异突然放大,前 3 层几乎完全一致。这很奇怪:如果第 4 层之前都一样,那第 4 层的输入肯定也一样,为什么偏偏这一层输出误差变大?于是我把第 4 层内部拆开:Q、K、V 投影输出单独比对,发现 Q 和 K 的投影没有差异,但 V 的投影输出从第 1024 个维度开始出现明显偏差。再往下追,发现是 V 投影之后的reshape操作后,国产卡编译器做了一次不必要的重排,导致内存布局和 GPU 上的不同,最终影响到了数值累加顺序。3.2 三个最常见的精度偏差来源这个案例其实只是冰山一角。后面前前后后还排查出了另外几个不同的根因,归纳下来,国产卡上的精度偏差主要来自三个方面:偏差来源机理典型表现浮点数累加顺序差异编译器把求和改写成了树状归约,跟 GPU 上的线性扫描不一样误差在几十个维度上缓慢积累,到深层才爆发归一化实现差异LayerNorm 的均值/方差计算可能先转 FP32 再做除法,也可能全程在 FP16 下对比精度时 LayerNorm 层误差单独增大特殊函数实现Gelu、Softmax 等函数用近似多项式或查表实现逐点误差在 1e-2 量级,模型最终输出累积偏差明显排查到这一步,你不能只是把代码拿回来改,而是要针对不同的根因做对应处理。累加顺序差异,可以尝试强制让模型以 FP32 做中间累加:# 在模型计算中显式提升累加精度 class LayerNormFP32(nn.Module): def forward(self, x): return F.layer_norm(x.float(), x.shape[-1:], self.weight.float(), self.bias.float()).to(x.dtype)归一化实现差异,可以尝试修改模型,先把输入转成 FP32 做完整计算再转回 FP16。特殊函数差异,则可以考虑用更保守的函数实现来替代,或者接受差异进行少量微调。3.3 一次性调平所有层:先对齐全局再回到局部逐层反复调很花时间。后来我意识到一个更高效的策略:先让模型跑完整的 FP32 精度,确认结果与 GPU 完全一致,然后再逐个模块换成 FP16,看误差从哪一步开始冒出来。这个“白盒二分法”把排查时间从一周压缩到了两天。最终的解决方案其实是混合的:绝大多数层保持 FP16 不变,只把 LayerNorm、Softmax 这几个对精度敏感的模块强制用 FP32 计算,整体精度损失从 3.5% 压到了 0.2% 以内,完全满足业务要求。4. 性能优化:从“能跑”到“跑满算力”的四个杠杆精度对齐之后,模型的推理是正确了,但性能非常难看——同样的 batch size 下,单卡吞吐只有原来 GPU 的三成。这个时候才进入真正的性能调优阶段,这一步更考验对硬件底层的理解。4.1 显存管理与内存碎片国产 NPU 的显存管理和 GPU 很不一样。GPU 上 PyTorch 的缓存分配器把显存块复用得很好,但国产卡的运行时可能会出现显存碎片化严重的情况。同样是跑一个大 batch,提前分配好显存可以比动态申请快 20% 以上。解决方法是:把推理服务里的显存分配尽量提前到进程启动阶段完成,不要在每次请求时反复申请释放。我们用一个显存池把输入、输出、中间结果全部提前分配好,后面每次请求只做数据拷贝,几乎没有动态申请。4.2 图融合与算子融合的取舍图融合在 GPU 上通常是纯赚的,但在国产 NPU 上需要具体情况具体分析。我们实测对比下来,把 8 个小算子融合成一个大算子之后,某些情况下性能反而下降了,因为融合后的算子内部并行度不够,调度器没法塞满计算单元。这里我踩了一次很深的坑:听厂商 FAE 建议把 MHA 里的 QKV 投影全部融合成一个大的 GEMM,结果性能掉了一半。后来分析发现,我们模型的 hidden_size 是 768,融合后的 GEMM 的 M 维度是 1(单条样本),计算单元利用率极低;拆分成三个独立的 256 维 GEMM 反而能并行执行。所以做算子融合时一定要结合自己的 shape 特点做实测,不要盲信工具默认配置。4.3 Batch 大小与动态 Padding多 batch 推理时,动态 padding 是提升性能的关键。GPU 上大家习惯把所有请求 pad 到固定长度,再一次性跑完。这个策略在国产 NPU 上问题很大——因为静态 Shape 要求固定长度,一旦 pad 到 128,而实际平均请求长度只有 40,算力全浪费在填充 token 上了。我们用了一个粗粒度的长度分桶策略:把请求按长度分成 32、64、96、128 四个桶,每个桶分别走对应的静态图,这样就能兼顾动态 Shape 的灵活性和静态 Shape 的性能。实测整体吞吐提升了 60%。优化项优化前优化后提升幅度显存预分配8 req/s11 req/s+37%算子融合调整11 req/s12 req/s+9%分桶动态 Padding12 req/s19 req/s+60%4.4 CPU 端的预处理卸载一个很容易忽略的性能瓶颈在 CPU 端。我们原先的 tokenizer 和后处理都在主进程里,结果 CPU 占用率直接打满,NPU 一直在等 CPU 喂数据。后来把 tokenizer 和后处理丢到独立进程,用队列传输数据,NPU 的利用率从 50% 提升到 85% 以上。这一步可以说是性价比最高的优化,代码改动少,收益立竿见影。5. 迁移步调与上线验证:如何把风险降到最低技术上的坑基本填完,接下来是执行层面的策略。我们最后花了一周时间做灰度上线,中途还回滚了一次,这里面也有不少值得说的。5.1 迁移推进的四个阶段迁移不能一上来就把生产流量切过去,最好是按下面这个节奏推进:离线验证:模型精度对比 + 性能压测。这个阶段要产出“差异清单”,明确哪些模块有精度损失、是多少。影子模式:生产流量同时打到 GPU 和国产卡上,对比两边结果。我们在这个阶段发现了几个只在真实流量分布下才会触发的长尾 Shape,之前测试集里完全没暴露。小流量灰度:先切 5% 的流量到国产算力上,观察一周的延迟、准确性、报错率。此时如果出了问题,影响面是可控的。逐步放量:确认稳定后逐步从 5% 到 20%、50%、100%,每一步至少观察 24 小时,有任何指标波动都要回退。5.2 影子模式里发现的坑:不只是模型本身影子模式阶段,我们发现国产卡在长尾请求上偶尔超时。后来定位到是模型里有一个极长的文本分类分支,句子长度超过 512 后触发了某种低效的内存访问模式,导致一遍推理耗时飙升到 300ms,而 GPU 上同样请求只需要 40ms。解决办法是给这个分支增加长度阈值,超过阈值的请求走简化逻辑。这种业务级的适配,测试集里完全不会覆盖到,只有影子模式能暴露。5.3 用监控指标把“正确”定义清楚上线前,我们和运维一起定了三个监控门槛:单次推理延迟 P99 不得高于原有 GPU 的 1.2 倍、坏样本率(跟 GPU 结果比对不一致的比例)不得高于 0.5%、服务可用性不低于 99.9%。这三个指标不仅是上线门槛,也是后续所有优化的目标。后来的回滚事件就是坏样本率超标触发的,当时有一个语料库的预处理逻辑在国产卡的 tokenizer 上略有不同,导致部分输入被截断,模型输出偏差巨大。没有明确的监控指标,你会发现迁移项目很容易陷入“到底能不能上线”的无限拉扯中。定义好量化门槛,所有争论都能被一锤定音。6. 写在迁移完成之后:我个人的几点体会复盘整个迁移过程,我最深的体会是:国产算力迁移与其说是一次技术替换,不如说是一次对自身模型的完整“体检”。在 GPU 上跑得好好的模型,你从没仔细看过它的算子分布,也从来没关注过哪个算子被 CPU 回退、哪个模块精度其实很脆弱。这次迁移逼着我们把模型从上到下翻了个遍,很多潜在问题反而提前暴露了。如果你的团队也准备做类似迁移,我给三个具体建议:尽早接触目标算力环境的厂商技术支持,他们的文档可能不够完善,但 FAE 手里往往有一批前人踩坑的案例,提前问一圈能省下好几天的排查时间。把精度对比工具和算子统计流程提前固化到 CI 里,不要靠手工跑。后续只要模型一更新,就能立刻知道新的算子是否会被 CPU 回退、精度是否会漂移。做好“一些算子永远无法完美对齐”的心理准备,业务上能接受的误差范围才是关键,不需要追求逐比特一致。最后一个实用小技巧:排查精度问题时,可以用np.save把 GPU 和国产卡的中间层输出都存下来,然后写一个脚本自动对比每层的余弦相似度。这样几万个维度一眼就能看出异常层,比肉眼盯几个数值高效得多。这次迁移最终顺利给了我们信心——国产算力完全能够支撑生产级的推理业务,前提是做好适配工作。过程虽然曲折,但走通之后再回头看,很多坑都是有迹可循的。希望这篇复盘能帮后面的人省掉我们当初交的那些学费。
今年我们接了个有点棘手的活儿:把线上一个日调用量千万级的NLP推理服务,从国际主流的GPU平台平滑迁到国产AI加速卡上。任务听起来不复杂——换个硬件,模型重新导出一下不就行了?实际上从立项到灰度上线,我们整整折腾了两个月,中途数次想放弃,问题一个接一个:图编译直接崩、算子不支持、推理精度飘了三个多点、性能上不去只能跑 GPU 的一个零头。这篇就把我们踩过的坑和最终沉淀下来的迁移思路完整复盘一遍,希望那些同样被“算力迁移”和“推理适配”夹在中间的团队能少走一些弯路。先说结论:从 GPU 迁到国产算力,本质上不是“改几行代码”的活,而是一次完整的软硬件适配工程。模型结构、算子映射、精度策略、显存管理、部署方式,每一个环节都可能埋雷。但好消息是,这个过程的坑是结构性的、可预测的,只要按对顺序一步步排查,绝大多数问题都能提前暴露并解决。1. 迁移前的算力摸底:先别急着改代码我们最初犯的第一个错误,就是拿到新卡之后立刻把原来的 PyTorch 模型部署脚本跑了一遍,结果当然是一堆看不懂的报错。后来跟某国产算力厂商的 FAE(现场技术支持工程师)聊了两次,才意识到迁移工作的第一步根本不应该发生在跑代码的环境里,而应该在会议室里完成一次完整的算力实测和算子梳理。1.1 为什么算力选型直接决定后面要踩多少坑不同国产 AI 加速卡之间的差异,远比不同型号的 NVIDIA GPU 之间大得多。NVIDIA 的卡从 1080 到 A100 再到 H100,至少 CUDA 生态是统一的,算子库、推理引擎、量化工具链都是同一套体系,迁移成本主要在显存和计算精度上。但国产加速卡这边,每一家的指令集、片上存储结构、算子实现方式都不同,有的卡对 PyTorch 的兼容性做得比较好,有的卡则要求模型先转成特定中间格式才能运行。这里的核心矛盾不是“跑不跑得动”,而是生态成熟度和算子覆盖率的差异。我们在选型时列了一个简单的评估维度:评估维度具体问题评估方式算子覆盖模型里的算子(Conv、LayerNorm、Softmax、MatMul等)新卡支持几个用模型结构解析脚本统计算子类型分布,再对照新卡的算子支持列表精度支持是否支持 FP16 / BF16 / INT8,以及量化的激活函数支持情况看厂商提供的精度支持文档,实测小模型跑一遍显存容量峰值显存占用是否超过卡上内存在 GPU 上先测模型峰值显存,留出 30% 余量框架对接PyTorch 是直接适配还是要转 ONNX用最简模型跑通测试工具链成熟度是否提供模型转换工具、性能分析器、调试工具询问 FAE,翻开发者社区我们最终选了一款显存和算力规模接近现有 GPU 七成水平的国产 NPU 加速卡,理由很简单:它的 PyTorch 适配层做得比较成熟,我们的模型不需要额外改网络结构就能完成图编译。1.2 模型算子分布:一张图看出迁移工作量在决定适配方案之前,我用一个简单的 ONNX 解析脚本把模型的算子分布统计了一遍。这里不是让大家必须转 ONNX,而是建议用这种方式快速了解自己的模型里到底有哪些“稀有算子”——比如一些自定义的 Attention 实现、奇怪的激活函数、或者是用了torch.einsum这种会被编译器展开成复杂子图的写法。import onnx from collections import Counter model = onnx.load("model.onnx") node_types = Counter(node.op_type for node in model.graph.node) total = sum(node_types.values()) for op, cnt in node_types.most_common(): print(f"{op:20s} {cnt:5d} {cnt / total * 100:.1f}%")我们模型里当时的算子分布大概是:MatMul 占 35%、LayerNormalization 占 18%、Softmax 占 12%、Gelu 占 15%,剩下的是 Reshape、Transpose、Concat 这类数据搬运算子。表面上看都是标准算子,应该没什么问题。但这恰恰是最迷惑人的地方。标准算子名不代表标准实现,比如 LayerNorm 在一张卡上可能是用 CUDA kernel 直接算的,在另一张卡上则可能被编译器展开成一堆乘加和开方操作。一旦展开,精度特性就不一样了,性能也完全不是一回事。这个领悟是我们后续所有排查思路的出发点。2. 部署与适配的第一个坎:算子编译报错与替换策略环境搭建阶段比预想中耗时更久,光是“跑通一个示例模型”就花了近一周。这个过程中遇到的坑,大致可以分为环境配置类、图编译类和算子缺失类三种。环境配置类的坑网上一搜一大堆,这里只说两个我们反复踩的。2.1 环境依赖的版本锁死:换个卡连 Python 版本都得重新对齐国产算力的软件栈一般分为三层:驱动、运行时、框架适配插件。我们最开始按照厂商文档在服务器上装了最新版驱动,结果死活无法加载推理运行时,反复报driver mismatch。后来发现是依赖的某个系统库版本跟预装的 CUDA 兼容层冲突——即使你完全不使用 GPU,某些国产加速卡的驱动安装包仍然会尝试跟 CUDA 的运行时组件做版本检查。解决办法是彻底清空服务器上的 NVIDIA 驱动和 CUDA,再装国产卡的运行时。这里的关键教训是:迁移到国产算力的机器,最好做一次系统级的环境初始化,不要指望跟原有的 GPU 环境共存。如果确实有混跑需求,用容器隔离是最稳妥的方案,而且容器里也要严格指定基础镜像版本。2.2 图编译失败:动态 Shape 是最大的敌人图编译环节,我们遇到的核心报错是shape inference failed。原因是模型里有两处torch.where混着动态 mask 输入,导致张量形状在运行期才能确定。NVIDIA 的 TensorRT 通过显式设置 optimization profile 也能处理一部分动态形状,但国产加速卡的编译器对动态 Shape 的容忍度低得多,很多直接要求模型输入是静态的。这一块没有捷径,唯一的办法是在模型导出阶段把动态维度固定。# 导出时固定序列长度,避免动态 shape dummy_input = ( torch.randn(1, 128, hidden_size, dtype=torch.float32).to(device), torch.tensor([[1] * 128], dtype=torch.int64).to(device), ) torch.onnx.export( model, dummy_input, "model.onnx", input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes=None, # 关键:关闭动态轴 opset_version=17, )我们当时的做法是统计线上服务的实际请求长度分布,发现绝大多数请求集中在 64~128 token 之间,于是统一填充到 128,超出部分走一个回退逻辑。这样做虽然牺牲了一点响应速度,但换来的是模型在国产卡上的稳定编译。另外,如果模型里有if分支控制流,比如if seq_len > 64: use_flash_attention,图编译器会很难处理。建议在导出前把这种运行期条件分支全部去掉,统一走一条计算路径。2.3 算子替换:不支持的算子怎么处理做完上一步,图编译能跑通了,但日志里出现警告:有两个算子类型回退到了 CPU 执行。一个是torch.fft,我们用来做某种频域特征增强的;另一个是torch.linalg.eigh,用在后处理阶段。CPU 回退意味着这部分计算量会从加速卡上搬到 CPU,对于推理性能是毁灭性的。我们尝试了两种替换策略:数学等价替代:torch.fft.rfft用在特征提取,跟一个固定卷积核的 Conv1d 在效果上近似,于是直接把 FFT 换成几个 Conv1d,模型效果几乎不变,但算子完全落在 NPU 上。拆分替代:torch.linalg.eigh用于对 2x2 矩阵求特征值,这种小矩阵直接用解析公式手写。手写之后不仅避免了 CPU 回退,速度还快了 3 倍。这里有个经验:大多数不支持的算子都可以通过数学变换绕过去,关键是你要对自己的模型逻辑足够清楚。如果实在绕不过,就写自定义算子融合进厂商的算子库。写自定义算子的学习成本很高,建议只在算子确实无法替代且直接影响正确性时才考虑。3. 精度漂移排查:那些“差不多”的数字害死人模型适配完、在国产卡上能正常出结果之后,我们进入精度验证阶段。先在测试集上跑了一遍,发现整体准确率比 GPU 低了 3.5 个百分点。这个差异足够大,绝对不能直接上线,但问题是:它到底从哪一层开始错的?3.1 逐层输出比对:定位误差起始层我的排查方法非常简单原始:在模型的关键中间层都插上print,同时对 GPU 和国产卡跑同一个输入,逐层比对输出张量的均值、方差和最大绝对误差。这个方法虽然笨,但绝对有效。第一次排查发现,从第 4 个 Transformer 层开始输出差异突然放大,前 3 层几乎完全一致。这很奇怪:如果第 4 层之前都一样,那第 4 层的输入肯定也一样,为什么偏偏这一层输出误差变大?于是我把第 4 层内部拆开:Q、K、V 投影输出单独比对,发现 Q 和 K 的投影没有差异,但 V 的投影输出从第 1024 个维度开始出现明显偏差。再往下追,发现是 V 投影之后的reshape操作后,国产卡编译器做了一次不必要的重排,导致内存布局和 GPU 上的不同,最终影响到了数值累加顺序。3.2 三个最常见的精度偏差来源这个案例其实只是冰山一角。后面前前后后还排查出了另外几个不同的根因,归纳下来,国产卡上的精度偏差主要来自三个方面:偏差来源机理典型表现浮点数累加顺序差异编译器把求和改写成了树状归约,跟 GPU 上的线性扫描不一样误差在几十个维度上缓慢积累,到深层才爆发归一化实现差异LayerNorm 的均值/方差计算可能先转 FP32 再做除法,也可能全程在 FP16 下对比精度时 LayerNorm 层误差单独增大特殊函数实现Gelu、Softmax 等函数用近似多项式或查表实现逐点误差在 1e-2 量级,模型最终输出累积偏差明显排查到这一步,你不能只是把代码拿回来改,而是要针对不同的根因做对应处理。累加顺序差异,可以尝试强制让模型以 FP32 做中间累加:# 在模型计算中显式提升累加精度 class LayerNormFP32(nn.Module): def forward(self, x): return F.layer_norm(x.float(), x.shape[-1:], self.weight.float(), self.bias.float()).to(x.dtype)归一化实现差异,可以尝试修改模型,先把输入转成 FP32 做完整计算再转回 FP16。特殊函数差异,则可以考虑用更保守的函数实现来替代,或者接受差异进行少量微调。3.3 一次性调平所有层:先对齐全局再回到局部逐层反复调很花时间。后来我意识到一个更高效的策略:先让模型跑完整的 FP32 精度,确认结果与 GPU 完全一致,然后再逐个模块换成 FP16,看误差从哪一步开始冒出来。这个“白盒二分法”把排查时间从一周压缩到了两天。最终的解决方案其实是混合的:绝大多数层保持 FP16 不变,只把 LayerNorm、Softmax 这几个对精度敏感的模块强制用 FP32 计算,整体精度损失从 3.5% 压到了 0.2% 以内,完全满足业务要求。4. 性能优化:从“能跑”到“跑满算力”的四个杠杆精度对齐之后,模型的推理是正确了,但性能非常难看——同样的 batch size 下,单卡吞吐只有原来 GPU 的三成。这个时候才进入真正的性能调优阶段,这一步更考验对硬件底层的理解。4.1 显存管理与内存碎片国产 NPU 的显存管理和 GPU 很不一样。GPU 上 PyTorch 的缓存分配器把显存块复用得很好,但国产卡的运行时可能会出现显存碎片化严重的情况。同样是跑一个大 batch,提前分配好显存可以比动态申请快 20% 以上。解决方法是:把推理服务里的显存分配尽量提前到进程启动阶段完成,不要在每次请求时反复申请释放。我们用一个显存池把输入、输出、中间结果全部提前分配好,后面每次请求只做数据拷贝,几乎没有动态申请。4.2 图融合与算子融合的取舍图融合在 GPU 上通常是纯赚的,但在国产 NPU 上需要具体情况具体分析。我们实测对比下来,把 8 个小算子融合成一个大算子之后,某些情况下性能反而下降了,因为融合后的算子内部并行度不够,调度器没法塞满计算单元。这里我踩了一次很深的坑:听厂商 FAE 建议把 MHA 里的 QKV 投影全部融合成一个大的 GEMM,结果性能掉了一半。后来分析发现,我们模型的 hidden_size 是 768,融合后的 GEMM 的 M 维度是 1(单条样本),计算单元利用率极低;拆分成三个独立的 256 维 GEMM 反而能并行执行。所以做算子融合时一定要结合自己的 shape 特点做实测,不要盲信工具默认配置。4.3 Batch 大小与动态 Padding多 batch 推理时,动态 padding 是提升性能的关键。GPU 上大家习惯把所有请求 pad 到固定长度,再一次性跑完。这个策略在国产 NPU 上问题很大——因为静态 Shape 要求固定长度,一旦 pad 到 128,而实际平均请求长度只有 40,算力全浪费在填充 token 上了。我们用了一个粗粒度的长度分桶策略:把请求按长度分成 32、64、96、128 四个桶,每个桶分别走对应的静态图,这样就能兼顾动态 Shape 的灵活性和静态 Shape 的性能。实测整体吞吐提升了 60%。优化项优化前优化后提升幅度显存预分配8 req/s11 req/s+37%算子融合调整11 req/s12 req/s+9%分桶动态 Padding12 req/s19 req/s+60%4.4 CPU 端的预处理卸载一个很容易忽略的性能瓶颈在 CPU 端。我们原先的 tokenizer 和后处理都在主进程里,结果 CPU 占用率直接打满,NPU 一直在等 CPU 喂数据。后来把 tokenizer 和后处理丢到独立进程,用队列传输数据,NPU 的利用率从 50% 提升到 85% 以上。这一步可以说是性价比最高的优化,代码改动少,收益立竿见影。5. 迁移步调与上线验证:如何把风险降到最低技术上的坑基本填完,接下来是执行层面的策略。我们最后花了一周时间做灰度上线,中途还回滚了一次,这里面也有不少值得说的。5.1 迁移推进的四个阶段迁移不能一上来就把生产流量切过去,最好是按下面这个节奏推进:离线验证:模型精度对比 + 性能压测。这个阶段要产出“差异清单”,明确哪些模块有精度损失、是多少。影子模式:生产流量同时打到 GPU 和国产卡上,对比两边结果。我们在这个阶段发现了几个只在真实流量分布下才会触发的长尾 Shape,之前测试集里完全没暴露。小流量灰度:先切 5% 的流量到国产算力上,观察一周的延迟、准确性、报错率。此时如果出了问题,影响面是可控的。逐步放量:确认稳定后逐步从 5% 到 20%、50%、100%,每一步至少观察 24 小时,有任何指标波动都要回退。5.2 影子模式里发现的坑:不只是模型本身影子模式阶段,我们发现国产卡在长尾请求上偶尔超时。后来定位到是模型里有一个极长的文本分类分支,句子长度超过 512 后触发了某种低效的内存访问模式,导致一遍推理耗时飙升到 300ms,而 GPU 上同样请求只需要 40ms。解决办法是给这个分支增加长度阈值,超过阈值的请求走简化逻辑。这种业务级的适配,测试集里完全不会覆盖到,只有影子模式能暴露。5.3 用监控指标把“正确”定义清楚上线前,我们和运维一起定了三个监控门槛:单次推理延迟 P99 不得高于原有 GPU 的 1.2 倍、坏样本率(跟 GPU 结果比对不一致的比例)不得高于 0.5%、服务可用性不低于 99.9%。这三个指标不仅是上线门槛,也是后续所有优化的目标。后来的回滚事件就是坏样本率超标触发的,当时有一个语料库的预处理逻辑在国产卡的 tokenizer 上略有不同,导致部分输入被截断,模型输出偏差巨大。没有明确的监控指标,你会发现迁移项目很容易陷入“到底能不能上线”的无限拉扯中。定义好量化门槛,所有争论都能被一锤定音。6. 写在迁移完成之后:我个人的几点体会复盘整个迁移过程,我最深的体会是:国产算力迁移与其说是一次技术替换,不如说是一次对自身模型的完整“体检”。在 GPU 上跑得好好的模型,你从没仔细看过它的算子分布,也从来没关注过哪个算子被 CPU 回退、哪个模块精度其实很脆弱。这次迁移逼着我们把模型从上到下翻了个遍,很多潜在问题反而提前暴露了。如果你的团队也准备做类似迁移,我给三个具体建议:尽早接触目标算力环境的厂商技术支持,他们的文档可能不够完善,但 FAE 手里往往有一批前人踩坑的案例,提前问一圈能省下好几天的排查时间。把精度对比工具和算子统计流程提前固化到 CI 里,不要靠手工跑。后续只要模型一更新,就能立刻知道新的算子是否会被 CPU 回退、精度是否会漂移。做好“一些算子永远无法完美对齐”的心理准备,业务上能接受的误差范围才是关键,不需要追求逐比特一致。最后一个实用小技巧:排查精度问题时,可以用np.save把 GPU 和国产卡的中间层输出都存下来,然后写一个脚本自动对比每层的余弦相似度。这样几万个维度一眼就能看出异常层,比肉眼盯几个数值高效得多。这次迁移最终顺利给了我们信心——国产算力完全能够支撑生产级的推理业务,前提是做好适配工作。过程虽然曲折,但走通之后再回头看,很多坑都是有迹可循的。希望这篇复盘能帮后面的人省掉我们当初交的那些学费。