国产加速卡 PyTorch 适配与迁移：从算子差异到性能对齐
说明：本篇为综述/无第一手实验数据，采用公开适配框架资料与迁移实践路线。
实验目标
国产加速卡（昇腾、寒武纪、沐曦、天数智芯等）在超算与智算中心的使用越来越普遍，但生态成熟度与 CUDA 仍有差距。本篇目标：讲清PyTorch 程序迁移到国产加速卡的完整路径——适配层到底做了什么、你的训练脚本哪些地方会"静默变慢"或"直接报错"、如何用迁移清单和基准对比把性能对齐。这不是"换一行代码"的事，而是理解适配栈后的一次系统性改造。
代码逐段拆解
第一段：适配层的三层结构
text
你的训练脚本 (torch.nn / torch.optim / torch.distributed)
│
▼
PyTorch 主干 (保持原 API 不变)
│
▼
适配层: 设备插件 (torch_npu / torch_mlu / ...) + 自定义算子实现
│ │
├── 算子注册: aten 算子 → 国产卡 kernel
├── 设备管理: DeviceGuard, 显存分配器, 事件/流
└── 编译后端: 图模式 (torch.compile / AI 编译器) 做算子融合
│
▼
国产加速卡驱动 + 运行时 (ACL / CNRT / 各自 runtime)理解要点：
- 训练脚本几乎不动：适配层以"设备插件"形式挂进 PyTorch，
model.cuda()换成model.to('npu')（或通过PYTORCH_DEVICE环境变量统一），前向/反向/优化器代码保持原样——这是迁移成本低的根本原因。 - 算子注册是关键差异点：PyTorch 的 aten 算子（matmul、conv、softmax、LayerNorm 等）在 CUDA 上有原生 kernel；国产卡适配层要为每个算子提供自己的 kernel 实现。覆盖率决定兼容性——没覆盖到的算子会 fallback 到 CPU，表现为"能跑但极慢"。
- 图模式是性能追赶手段：国产卡的 kernel 质量初期不如 CUDA 库，靠
torch.compile/AI 编译器把多个算子融合成一个大 kernel，减少启动开销与显存往返，往往能追回 20%-40% 性能差距。
第二段：迁移第一步——设备无关化改造
python
# 迁移前：写死 cuda
device = 'cuda'
model = model.to(device)
# 迁移后：设备无关
import os
device = os.environ.get('AI_DEVICE', 'cuda') # 支持 cuda/npu/mlu/cambricon
model = model.to(device)
# 关键兼容点：
# 1. tensor 显存 API：tensor.cuda() → tensor.to(device)
# 2. 分布式后端：dist.init_process_group(backend='nccl') → 按卡适配（hccl/cncl）
# 3. 混合精度：torch.cuda.amp 部分 API 需换适配层等价物
# 4. 自定义 CUDA 算子（.cu 源码）无法直接迁移 → 需重写或走算子库逐项说明：
to(device)是万能入口：所有显存分配、设备切换都走它；写死cuda的地方是迁移第一排查点。- 分布式后端不同：NCCL 换成本地集合通信库（昇腾 HCCL、寒武纪 CNCL），
init_process_group参数要换；好在 PyTorch 的 DDP 上层逻辑不变。 - 混合精度 API 有差异：
torch.cuda.amp.GradScaler在部分适配层有等价实现，但自动混合精度（AMP）的算子白名单不同——需要按适配层文档核对。 - 自定义 CUDA kernel 是迁移最大成本：只要有手写
.cu算子，就必须用适配层的算子库重写或用 Triton 重写（08 章 Triton 篇的基础就在这派上用场）。
原理讲解
为什么"换 device 就能跑"不等于"迁移完成"
三个层次的坑，从浅到深：
- 兼容层（能跑）：API 能调用，算子有覆盖。但 fallback 到 CPU 的算子可能让整段训练慢 10 倍——能跑 ≠ 能训练。
- 性能层（能训）：算子覆盖 + 图融合到位，训练速度接近 CUDA 基线。判断标准：端到端吞吐差距 < 20%。
- 精度层（能比）：浮点实现差异（如不同厂商的 FlashAttention 实现、累加顺序）可能导致 loss 曲线细微不同——需要同 seed 对比、loss 对齐验证。
性能对齐方法
基线：同模型同数据，CUDA 卡上记录吞吐/显存/每步耗时
迁移：国产卡上跑同样配置
差距定位：逐算子 profile（适配层工具/ torch.profiler）
优化手段：算子融合（torch.compile）→ 显存分配器调优 → 集合通信参数调优 → 数据加载 pipeline
验收：吞吐差距 ≤20%、loss 曲线对齐、48h 稳定性关键认知：性能差距通常不是"一个算子慢"，而是"若干算子 fallback + 图未融合 + 显存分配策略差异"三者的叠加。先修 fallback（最伤）、再开图模式、最后调显存/通信参数，按此顺序收益最大。
实验结果
（估算口径，非实测；不同厂商/型号差异大）
| 迁移阶段 | 相对 CUDA 基线吞吐 | 主要问题 |
|---|---|---|
| 仅换 device | 20%-50% | 算子 fallback、无图融合 |
| + 修 fallback 算子 | 60%-80% | 单算子 kernel 效率差距 |
| + torch.compile 图融合 | 80%-100% | 接近基线；个别算子仍慢 |
| + 显存/通信调优 | 90%-110% | 与 CUDA 基线相当或略优 |
工程启示
- 迁移前先做算子覆盖审计：用适配层工具扫描脚本用到的算子清单，对比适配层支持表，把 fallback 算子提前列为工作量——这是排期的主要依据。
- 自定义 CUDA kernel 是分水岭：有手写算子的项目，迁移成本翻倍；没有的话，纯 PyTorch 项目迁移通常 1-2 周能到"能训"。
- 验收标准要写进合同/交付：吞吐差距、loss 对齐、稳定性三件套缺一不可——"能跑起来"不是交付标准。
- 与异构调度配套：07 章的 heterogeneous-resource-management 讲的是"混布怎么管"，本篇是"单个作业怎么迁"——混布集群里国产卡跑不了 CUDA 算子，迁移质量直接决定集群利用率。
下一篇预告：单卡迁移讲完，回到训练的全局稳定性——"预训练稳定性工程：loss spike 检测与断点重启"。