企业官网建设流程全解析

HSTU模型在GPU上跑得好好的，为什么要迁到昇腾NPU？这是我们组接到任务后的第一反应。但现实是，生成式推荐这类对算力和成本同时敏感的业务里，国产算力的供给稳定性和单卡性价比确实已经具备足够的吸引力。我们花了大几周时间把HSTU（Hierarchical Sequential Transduction Unit）这套生成式推荐模型从A100集群迁移到昇腾910B，中间踩了不少坑，也沉淀出一套可复用的适配方法论。这篇文章不聊论文复现，只讲迁移过程中真正卡过我们的算子适配、精度对齐、分布式通信改造和性能调优，给后面要碰昇腾NPU适配的同学一份能直接参考的路线图。

1. 为什么要搬HSTU：生成式推荐上国产芯片的真实动机

1.1 先搞明白HSTU是什么，才能知道适配难在哪

HSTU是Meta提出的一套推荐系统架构，和传统DLRM这类深度推荐模型最大的区别在于：它把推荐任务建模成一个序列生成问题，用类似Transformer的堆叠层去直接建模用户行为序列，然后预测下一个要交互的商品ID。它不是一个单一模型，而是一整套包含Embedding层、多层HSTU Block、输出头在内的生成式推荐范式。

HSTU Block本身有三个核心子模块：Multi-head Latent Attention（MHLA）、Mixer Layer、Fusion Layer。MHLA负责在稀疏特征序列上做多头注意力计算，Mixer Layer是Position-wise的MLP加上非线性激活（HSTU论文里用的是SwiGLU），Fusion Layer则通过门控机制把前面两者的输出融合在一起。另外它还有一个很关键的设计，就是相对位置编码（Relative Positional Encoding）和多种Attention Mask的组合——这部分是后面在NPU上算子适配最大的难点。

理解这个结构的意义在于：适配工作的难度，基本取决于模型里有哪些算子和张量操作，而不是这个模型叫什么名字。HSTU里大量用到Einsum、Masked Softmax、相对位置偏置计算、动态Shape的Padding和Gather操作，再加上Embedding表又大又稀疏，这些特征决定了它在昇腾NPU上不会像普通CV模型那样"装上torch_npu就能跑"。

1.2 迁移前必须做的三层评估：算子覆盖、生态依赖、业务收益

我强烈建议任何团队在动手之前先做三层评估，别直接上来就改代码。

第一层是算子覆盖评估。把HSTU模型结构里用到的所有PyTorch算子梳理一遍，对照昇腾CANN算子库的覆盖清单，重点看那些不常见的高阶算子是否支持。我们当时的结论是：70%的算子可以直接映射，20%需要替换或改写，10%需要深入算子级适配。这个比例决定了工作量上限。

第二层是生态依赖评估。你的训练脚本里除了PyTorch之外还依赖哪些库？比如我们用了DeepSpeed做ZeRO优化，但DeepSpeed对昇腾的支持当时并不完整，这就导致我们需要评估是否降级到DDP+手动梯度累积，或者接受某些功能在NPU上不可用。

第三层才是业务收益评估。昇腾单卡和GPU单卡的性能差距、单位算力成本、供货稳定性，这些数字要放在一起算总账。当时我们测下来，昇腾910B单卡的训练吞吐大约是A100的0.7倍左右，但同预算下能拿到的昇腾卡数更多，算总账反而是划算的。如果你们跑的是小规模模型，迁移成本摊不平，那就没必要折腾。

2. 环境搭建第一课：版本矩阵比代码改造更早卡住你

2.1 CANN、torch_npu、PyTorch三者版本必须严格咬合

昇腾NPU上跑PyTorch，不是pip install torch_npu就行。它依赖一整套CANN工具链，而CANN、torch_npu、PyTorch三者之间有着严格的版本对应关系。我们在最开始就吃过这个亏：装了一套CANN 8.0的驱动，配了某个旧版本的torch_npu，结果一跑训练就报莫名其妙的内存错误，查了两天才发现是版本不匹配。

我的经验是：先确定你的PyTorch版本，然后去昇腾社区找到与该PyTorch版本匹配的torch_npu wheel包，最后严格按昇腾官方文档的兼容性矩阵去安装对应版本的CANN。顺序不能反，因为CANN的驱动层一旦装了，再回退的代价非常大。

安装完之后，第一件验证环境的事就是跑一个最简单的张量运算，确认torch_npu能正常加载且算子能真实下沉到NPU执行。注意我说的是"真实下沉"——因为有一些算子会静默回退到CPU执行，如果你不看日志，根本发现不了。这个在后面的算子适配里会详细说。

2.2 从.cuda()到.npu()：代码侧的真实改动量

网上很多教程把昇腾适配说得跟换一行代码一样简单：把model.cuda()换成model.npu()就行。实际改起来确实是这么个思路，但远不止这一处。

以我们HSTU训练脚本为例，代码改动集中在五个地方：

• 设备初始化：torch.cuda.set_device(rank)改成torch_npu.npu.set_device(rank)，device = 'cuda'改成device = 'npu'。
• 设备相关条件判断：代码里所有if torch.cuda.is_available():的地方，都要额外处理NPU分支，否则在昇腾环境下直接走False分支。
• Dataloader：GPU训练时为了加速数据传输，很多人会开pin_memory=True，但NPU下部分版本的torch_npu对pin_memory支持有问题，需要关掉或者改用自己的异步数据加载策略。
• amp混合精度：torch.cuda.amp.autocast()要改成torch.npu.amp.autocast()，GradScaler同理。
• 分布式初始化：init_process_group要指定backend='hccl'，这个后面单独展开讲。

整体看，纯代码改动量其实不大，大概一天的工时就能全部改完。真正的坑都在代码表面之下——比如数据加载的瓶颈、算子不支持的地方，这些改起来才是无底洞。

2.3 Dataloader与pin_memory的隐藏障碍

这里单独讲一下Dataloader的坑，因为推荐系统的训练流程和CV/NLP不一样：我们的输入不是单张图片或单段文本，而是变长行为序列，每个样本都有自己的长度，所以Collate要做大量Padding。GPU训练时我们可以依赖pin_memory加多进程Worker来掩盖Padding开销，但在NPU上这套逻辑会出问题。

具体的现象是：开启pin_memory后训练速度不升反降，而且偶发C10指针报错。排查后发现是torch_npu对pin_memory的默认实现并没有真正把内存固定在NPU可访问的物理页上，反而额外引入了内存拷贝开销。

我们的解法很朴素：关掉pin_memory，改用一个额外的线程池在CPU侧做序列Padding和预处理，再把处理好的Tensor异步拷贝到NPU上。这样改动后数据加载瓶颈反而比原来降低了不少。如果你的模型对数据加载敏感，建议也提前做这个改造，而不是等到Dataloader成为瓶颈才动手。

3. 算子适配：HSTU在昇腾上最难啃的硬骨头

3.1 先让模型跑起来：常见算子的映射关系

昇腾的算子适配遵循一套"先跑通、再优化"的逻辑。很多常用PyTorch算子在torch_npu上都有对应实现，可以直接跑，但实现效率和GPU不完全一样。我们梳理了HSTU中用到的算子，按适配难度把它们分成了三类：

第一类是直接映射的，包括Linear、LayerNorm、Embedding、GELU、Add、Mul、MatMul等基础算子。这些在昇腾上都有充分优化，不需要任何改动。

第二类是需要小改的，包括Softmax（关注mask参数传法）、Einsum（部分字符串pattern在NPU上不支持）、MaskedFill、Scatter等。这些算子本身在NPU上有实现，但参数支持范围比GPU窄，需要把一些复杂调用改写成分步操作。

第三类是难啃的，包括FlashAttention类的高效Attention实现、带有复杂Mask组合的Attention计算、以及某些动态Shape下的Gather/Index操作。这些在NPU上要么有性能问题，要么直接不支持，需要另想办法。

我的建议是：在动手改模型之前，先写一个脚本把模型跑一遍前向，把所有报错的算子收集齐，一次性解决，而不是跑一次改一个，那样效率太低。

3.2 相对位置编码与Attention Mask的适配方案

HSTU中相对位置编码的实现非常依赖高阶张量操作：它对序列中每个位置对计算一个相对距离，然后查表转换成一组合适的偏置向量，再加到Attention分数上。在GPU上我们可以直接用PyTorch的矩阵广播和torch.gather来实现，几个张量一拼就OK了。

但这个逻辑搬到昇腾NPU上，第一个版本跑下来发现Attention部分的前向速度慢得离谱——比GPU慢了近一个数量级。用npu-smi info和msprof分析后发现，问题出在三个方面：一是相对距离矩阵的动态Shape导致NPU无法提前编译最优kernel；二是Gather操作被降级到了AICPU上执行，没有走AI Core的高速通道；三是多个Mask操作没有做算子融合，导致反复读写中间结果。

我们的适配方案是：

1. 把序列长度固定为训练和推理时实际使用的值，不再动态变长，从而让相对距离矩阵的Shape变得静态可推导；
2. 用矩阵乘法和布尔掩码的组合方式来替代部分Gather操作，让计算落到AI Core上；
3. 把padding mask、causal mask、相对位置偏置三者融合成一个偏置矩阵，在Attention之前一次性叠加到Query-Key乘积上。

这三步做完，Attention部分的速度回来了，而且精度完全对齐。这也印证了一个经验：在NPU上做算子适配，核心思路就是把GPU上那些"灵活但隐晦"的高阶操作，改写为"机械但高效"的基础算子组合。别嫌丑，性能说话。

3.3 动态Shape问题：为什么NPU比GPU更敏感

动态Shape是这次迁移中花时间最多的一个坑，值得单独拿出来说。

GPU上的CUDA kernel对动态Shape容忍度很高，PyTorch的autograd会为每次新的输入Shape做一层JIT编译，开销虽然存在但不致命。但昇腾NPU的图编译机制更倾向于静态Shape：它提前把整张计算图编译成面向硬件调度单元的执行序列，一旦输入Shape变了，可能整个图的编译都要重来，代价非常昂贵。

在HSTU模型里，用户行为序列天然是变长的，每个Batch内都要padding到同一长度，而Batch和序列长度稍微一变，Attention部分的shape就跟着变。第一版跑训练的时候，每个step光图编译就要额外花掉很长时间，GPU上30毫秒的训练step在NPU上被拖到了300毫秒以上。

我们的解法是组合拳：

• 训练前根据数据统计，把序列长度分桶（比如取8种长度区间），每个桶内固定Shape；
• Batch内按桶组织数据，减少每个step的Shape变化频率；
• 关键计算路径上的Shape全部显式声明，不依赖PyTorch的shape推断。

做完这套改造后，图编译的额外开销基本消失了。老实说，这个改造过程比较痛苦，因为它需要动数据管线的设计，但收益也是实实在在的。任何一个要做NPU适配的团队，建议把动态Shape处理排进计划前两周。

4. 精度对齐：完整排查链路分享

4.1 分模块比对策略：从Embedding到Logits逐层核对

模型跑通之后，下一步就是精度对齐。我们的方法不是直接看最终loss和指标，而是分模块逐层比对。

具体做法是：在GPU和NPU上分别跑同一个训练好的checkpoint，加载完全相同的数据，逐个模块比对中间输出。我们会比较每层输出的余弦相似度、最大绝对误差和平均相对误差。HSTU模型的比对链路我建议按这个顺序来：

1. Embedding输出：这里最容易出问题的是Embedding层在不同框架下的参数布局不一致，导致输出张量本身错位；
2. Attention前的投影输出：MHLA的线性投影涉及大量Einsum操作，要确认NPU上的Einsum实现没有改变算子语义；
3. Attention得分和Softmax输出：这一层最容易出现数值误差，因为Softmax的指数运算在低精度下对输入范围很敏感；
4. Mixer/Fusion层输出；
5. 最终Logits和Loss。

如果某一层的相似度掉到0.999以下，就停下来深挖这一层，而不是继续往下比对。我们当时发现Attention输出层的余弦相似度只有0.97左右，标准差也明显偏高，最后定位到是混合精度的Loss Scaling策略不一致导致的，这个见下文。

4.2 混合精度和Loss Scaling在NPU上的差异

HSTU训练用到了混合精度。在GPU上我们用FP16+动态Loss Scaling，一切正常。切到NPU后，第一次跑训练直接遇到了Loss变为NaN的问题。

排查链路是这样的：先把混合精度全部关掉，全FP32跑，训练稳定，loss曲线正常。然后单独开启autocast，不开GradScaler，发现指数增长后loss慢慢开始波动。最后同时开autocast+GradScaler，发现loss频繁出现NaN。

进一步定位发现，NPU上torch.npu.amp.GradScaler对梯度的scale factor更新逻辑和GPU存在细微差异——动态调整时，某些极小梯度在FP16下直接下溢为0，scaler没有及时调整下限，导致梯度累计误差最终把loss推爆。我们的解决方式是：

• 改用BF16混合精度替代FP16。BF16的指数位和FP32一致，精度问题天然少很多；
• 如果业务必须用FP16，则把动态Loss Scaling改成固定值Loss Scaling，并且手动监控梯度数值范围。

这个经验让我意识到，在昇腾上跑混合精度，尽可能优先考虑BF16，能省掉大量跟数值稳定性缠斗的时间。

4.3 随机性因素排查：种子、初始化与累积误差

除了混合精度，还有一类精度问题来自随机性。GPU和NPU上即使设置相同的随机种子，很多操作的随机数生成顺序和内部实现也不同，这会导致训练过程无法逐bit复现，模型收敛点有细微差异。

但注意，这类随机性通常不叫"bug"，而是"正常波动"。真正需要排查的是确定性模式下依然出现的差异。我们踩过的具体场景是：在DDP训练时，因为和GPU环境一样没有设置set_deterministic，NPU上的某些reduce操作顺序不一致，导致多卡训练中每张卡的loss出现肉眼可见的抖动（从0.001量级到0.01量级）。后来我们显式设置了：

torch_npu.npu.set_deterministic(True)

并且把HCCL通信库的HCCL_DETERMINISTIC环境变量设为1，抖动就消失了。

如果你想复现我上面的排查过程，我建议你也保持同样的步骤：先固定种子和确定性模式，再关掉混合精度比对，然后逐层输出比对。不要一上来就怀疑框架底层实现，错误往往出在你自己代码的某个小小分支条件里。

5. 分布式训练与通信库切换：从NCCL到HCCL

5.1 DDP改造与HCCL环境变量要点

HSTU这种规模的模型，单卡训练几乎不可能满足业务迭代速度，所以分布式是必须的。在GPU上我们用的是NCCL作为后端，PyTorch DDP天然支持。切换到昇腾后，通信库换成了HCCL（华为集合通信库），DDP的初始化方式要改。

核心改动就两处：一是init_process_group的backend参数从'nccl'改成'hccl'；二是每张卡要先torch_npu.npu.set_device(local_rank)再初始化。但是分布式训练真正的问题不在这些明面上的API，而在于一些环境变量和经验值。

我们实际踩到的几个问题：

• HCCL连接超时：多机训练时，HCCL的建链过程对网络有一定要求，默认超时时间可能不够。需要设置HCCL_CONNECT_TIMEOUT为较大值（我们设了1800），否则会出现莫名其妙的初始化失败。
• 通信算子与计算算子的重叠：GPU上NCCL能够比较好地和CUDA算子做stream重叠，HCCL在torch_npu上默认也支持，但如果多个rank之间没有用torch_npu.npu.set_stream做显式控制，有时会出现通信等待计算的问题。训练吞吐掉10%~20%不等的坑，根源常常在这里。
• 共享内存提交流：多机训练时HCCL依赖共享内存和网卡，建议确认ulimit -l设置足够大，否则会报HCCL内存注册失败的错误。

5.2 多卡扩展性与显存策略调整

HSTU模型里Embedding表很大，占了显存的主要部分。GPU训练时我们习惯直接把整个Embedding Table加载到每张卡上显存里（即Data Parallel的做法），但在昇腾平台上，显存容量相对紧张，这种做法很容易OOM。

我们做了三件事来缓解：

1. Embedding分片：把大Embedding表按行切分到多张卡，通过聚合通信来做Gather。这个改动本来是为了多卡显存均衡，最后发现对性能影响不大，但对显存压力缓解非常明显。
2. 梯度合并：HSTU的embedding梯度非常稀疏，默认DDP的allreduce通信量很大。我们改为梯度累积多步后再做一次AllReduce，通信量降低为原来的1/4，训练吞吐反而提高了。
3. 上移模型并行：对HSTU Block部分，我们做了简单的Layer级Pipeline并行，卡间传输的只是中间激活值，降低了单卡显存峰值。

这些策略虽然听起来不复杂，但在昇腾上的调优空间比GPU更大——因为GPU的成熟生态里这些细节往往已经被上层框架自动处理了，而在NPU上需要你自己的系统里补回来。

6. 性能调优与实测量化：吃满NPU需要主动做些事

6.1 msprof分析和热点定位

代码跑通、精度对齐、分布式正常之后，真正的性能压测才开始。昇腾上常用的性能分析工具是msprof，它可以输出每个算子在NPU上的执行时间、占用的AI Core数量和AICPU负载情况。

我们第一次跑全量训练时，用msprof发现一个有意思的现象：模型总执行时间里，有近25%用在了AICPU上执行的非矩阵算子。而这些"非矩阵算子"恰恰是之前提到的动态Shape相关的Gather、Mask、Index操作。这说明我们虽然修了主链路的算子适配，但这些边边角角的算子拖了后腿。

定位热点的过程不复杂：msprof生成的文件里有每个算子的耗时统计表，按耗时排序后一目了然。真正要做的判断是：哪些算子可以合并，哪些可以替换，哪些可以整体挪出训练关键路径。

6.2 图模式、算子融合与静态Shape优化

在我们把所有能用手工优化的算子都优化完之后，发现还有接近15%的性能提升空间没有榨出来。这时候需要借助昇腾的图编译能力。

torch_npu上可以打开计算图模式，让CANN把整个前向+反向计算图做全局的算子融合和调度优化。打开后，我们能明显看到Attention里的多个小算子被融合成了一个大算子，减少了多次中间缓存的读写。配合前面提到的静态Shape改造，整体训练吞吐又提高了20%左右。

这个阶段的经验是：不要追求某个单独算子的极致优化，而是先把所有Shape静态化，再开图编译，最后才针对剩余热点做专项优化。这个顺序反了会事倍功半。

6.3 实测性能对比与调优前后数据

最终跑出的数据，给大家一个参考。我们环境是8卡昇腾910B，对比原来的8卡A100。调优前，训练吞吐只有GPU的40%左右，几乎没法用；调优后，训练吞吐大概到GPU的70%~80%。推理部分因为固定了Batch和序列长度，走图编译后吞吐基本打平GPU。

这个结果在我们预估范围内。昇腾910B的硬件算力不弱，瓶颈主要在图编译和算子调优的成熟度上。换句话说，HSTU昇腾NPU适配不是一台不能开的车，而是需要你熟悉它的发动机调校逻辑。一旦过了前面这几关，后面的常规迭代提升并不比GPU差多少。

7. 关于本次迁移的最终体会

把HSTU这套生成式推荐模型从GPU迁到昇腾NPU，说到底是三种能力的博弈：算子适配能力、精度工程能力、系统调优能力。缺一不可。

我们这次没有用任何外部供应商支持，完全靠团队内部分工协作，六周时间从零到能跑通训练和推理全流程。如果再让我做一次，我会在第一天就梳理算子和Shape，把动态Shape处理排在所有工作的最前面，因为这个坑的排查成本最高。

最后给几个很实际的建议：

• 如果你们的模型还在快速迭代阶段，先别急着上NPU。等模型结构稳定了再做适配，否则每次改模型结构都要重新过一遍算子适配流程，成本太高。
• 昇腾NPU适配不是一个人的活。算子适配和精度对齐需要懂模型的人，分布式和性能调优需要懂系统的人，这两拨人必须从一开始就一起工作，不要等前一个环节做完了再交接。
• 我建议你也别再观望了。先找个规模适中的模型在昇腾上跑通一遍全流程，把你的算子覆盖评估体系和性能基线建立起来，等哪天业务真的需要大算力国产化替代的时候，你手里已经有一张现成的适配地图了。