AI 纪元 追踪人工智能最新进展,深度解读前沿技术

PyTorch 迁到昇腾 910C/910D,先改哪几处

把 PyTorch 项目从 CUDA 搬到昇腾,先处理设备、HCCL 和混合精度,再用 CANN 8 Profiling 查看算子、通信、内存与下发延迟。代码能启动只是入口,自定义算子和版本对齐更容易拖慢进度。

0 次阅读 约 2960 字 6 分钟 AI 生成

把现有 PyTorch 项目搬到昇腾,第一轮工作往往很朴素:改设备、换分布式后端、核对混合精度,再把自定义算子逐个列出来。麻烦通常出现在程序启动之后——某个算子没有 NPU 实现,单机运行正常而多机通信变慢,或者版本组合不一致,导致同一段代码换台机器就报错。

CANN 位于 AI 框架与昇腾处理器之间。资料《大模型国产化适配1-华为昇腾AI全栈软硬件平台总结》将其概括为对上支持多种框架、对下服务 AI 处理器的软件层。对于仍要维护 PyTorch 代码库的团队,常见入口是 torch_npu。是否转到 MindSpore,可以放在算子覆盖和长期维护成本评估之后再决定。

设备字符串只是入口

下面这张表覆盖了最先要改的接口。它能帮助项目通过启动阶段,后面仍要继续处理版本、算子和分布式问题。

位置 CUDA 写法 昇腾侧写法 排查重点
插件加载 无 import torch_npu 确认 PyTorch、torch_npu、CANN 与驱动版本匹配
设备指定 device = 'cuda' device = 'npu' 多卡时明确设备编号
张量迁移 tensor.cuda() tensor.to('npu') 或 tensor.npu() .to('npu') 更便于保留通用代码
缓存清理 torch.cuda.empty_cache() torch.npu.empty_cache() 只能释放未使用缓存,不能替代内存问题定位
分布式后端 backend='nccl' backend='hccl' 多机项目还要检查拓扑和 rank 配置
混合精度 torch.cuda.amp.autocast() torch.npu.amp.autocast() FP16、BF16 需要按模型单独做精度验证

最小改动大致如下,具体接口仍要以项目使用的 torch_npu 版本为准:

import torch
import torch_npu
import torch.distributed as dist

device = torch.device("npu:0")
x = x.to(device)

dist.init_process_group(backend="hccl")

with torch.npu.amp.autocast():
    output = model(x)

改完这几行,还要全局搜索 torch.cuda、.cuda()、nccl、CUDA stream/event,以及项目里的 C++/CUDA 扩展。小样本能运行,只能说明程序通过了入口;训练跑到特定算子时,残留的 CUDA 调用仍可能暴露出来。

Profile 先回答“慢在哪里”

鲲鹏昇腾开发者社区对 CANN 8 的介绍提到,Profiling 模块可以记录算子耗时、通信时长、内存占用和下发延迟。迁移初期,这四项比一个总吞吐数字更有用。

  • 算子耗时高:先确认该算子是否已有 NPU 实现,再检查输入 shape、数据类型和融合情况。
  • 通信时长突增:把 HCCL、集群拓扑和 rank 映射单独检查,不要继续在计算算子上消耗时间。
  • 内存占用异常:查看张量生命周期、批次设置和缓存变化,频繁调用 empty_cache 不能解决所有 OOM。
  • 下发延迟偏高:留意大量小算子和计算图切分,再评估编译器优化或算子融合。

更稳妥的顺序是:

CANN 8 Profiling
→ 锁定主要耗时项
→ 核对算子覆盖与精度
→ 尝试融合和编译器优化
→ 必要时再实现自定义 TBE 算子

没有 Profile 就直接重写算子,很容易改到一个占时并不高的位置。原稿中“自动搜索可以解决固定比例瓶颈”一类说法也不适合当作工程结论,不同模型、输入形状和 CANN 版本之间差异很大。

Segmentation Fault 要回到算子边界

一篇昇腾系列参数梳理提到,相关架构将离散内存访问粒度从 512 字节缩减至 128 字节,并称访问效率提升 4 倍。这组参数可以说明硬件访问方式发生了变化,但不能直接推断某次 Segmentation Fault 就由内存对齐引起,也不能把 4 倍写成项目的实际性能收益。

如果报错来自自定义 C++ 算子,应先隔离最小输入,核查 shape、dtype、数据布局、内存边界和 TBE 定义。原有 CUDA 扩展也不会因为把 .cuda() 改成 .npu() 就自动获得 NPU 实现。Python 层设备调用已经替换干净、错误仍停留在算子内部时,继续修改上层业务代码通常没有帮助。

多机变慢,把 HCCL 单独拉出来

把 nccl 换成 hccl 只完成了接口切换。遇到单机或少卡运行正常、多机吞吐明显下降的情况,可以先对比 Profile 中的通信时长,再检查集群拓扑、rank 映射、初始化日志和 HCCL_BUFFSIZE。

HCCL_BUFFSIZE 没有适用于所有模型的固定值。更可靠的做法是保留调整前后的通信记录,在相同模型和批次设置下比较。这样才能判断变化来自通信参数,还是计算负载、内存占用或节点状态。

DeepSeek V4 更适合当规模参照

经济观察报 2026 年 5 月 1 日刊发的《华为算力“开链”》提到,DeepSeek V4 于 4 月 24 日发布。其中,V4-Pro 为 1.6 万亿参数、490 亿激活参数,V4-Flash 为 2840 亿参数、130 亿激活参数,上下文长度从 128K 扩展到 1M。

这些数字能帮助团队理解大参数量、长上下文会给内存、通信和计算图带来多大压力。但该来源摘要没有给出 CANN 适配代码或性能测试,不能据此把 DeepSeek V4 写成一份已经验证的迁移模板。规划类似项目时,仍要从可复现的小计算图开始,逐步记录算子覆盖、峰值内存和多机通信数据。

上线前留一张排查表

  • 版本矩阵:记录驱动、固件、CANN、PyTorch 和 torch_npu 的完整版本,避免只写“CANN 8”。
  • 设备调用:确认项目中没有遗漏的 torch.cuda、.cuda() 和 NCCL 配置。
  • 算子覆盖:列出全部自定义算子,标明已有 NPU 实现、需要替换或需要 TBE 开发的部分。
  • 精度基线:使用相同输入对比输出和 Loss,验收阈值按具体模型任务确定。
  • 内存记录:保存峰值占用和 OOM 出现位置,区分缓存、碎片和张量生命周期问题。
  • 通信记录:保存 HCCL 初始化日志、拓扑配置以及 Profile 中的通信时长。
  • 性能基线:归档算子耗时、通信时长、内存占用和下发延迟,保证每次调整都有前后对照。

只改设备字符串,最多证明程序过了启动阶段;在没有 Profile 的情况下重写全部算子,也很难知道时间究竟花在哪里。对 910C/910D 项目来说,最省时间的做法是把问题拆开:框架入口、算子覆盖、精度、内存和多机通信分别验证,每一步都留下可复现的记录。

参考资料

相关阅读