Aurora 华为昇腾 NPU 适配指导书

原始仓库：https://github.com/decisionintelligence/Aurora
目标硬件：华为昇腾 910B / Atlas 800T A2
CANN 版本：8.2.RC1
PyTorch：2.5.1
torchvision：0.20.1
torch_npu：2.5.1

关于 PyTorch 版本

Aurora 原始 requirements.txt 声明 torch==2.4.0，但昇腾官方对 CANN 8.2.RC1 不再发布 torch 2.4.0 配套的 torch_npu（torch 2.4.0 的 torch_npu 配套版本最高仅到 CANN 8.0.0）。 因此必须把 PyTorch 一并升级到 2.5.1。torch 2.5 对 Aurora 用到的 nn.Linear / nn.TransformerEncoderLayer / torch.fft / torch.linalg.pinv / F.scaled_dot_product_attention 等 API 完全向后兼容，无破坏性改动。

1. 适配总览

1.1 Aurora 代码 GPU 依赖分析

1.2 适配策略

策略：torch_npu transfer_to_npu 全局重定向
    + 运行时小粒度 monkey-patch（apply_all_patches）
    + launcher 入口包装（无需修改 runner.py）

核心原理：

1. torch_npu.contrib.transfer_to_npu 把 tensor.cuda() / torch.cuda.* 重定向到 NPU。
2. apply_all_patches() 在运行时替换 aurora.util_functions.resample / causal_attention_mask，并在 NPU 不支持 FFT 时全局替换 torch.fft.rfft/irfft。
3. aurora_launcher.py 在 runner.py 之前完成所有上述初始化，不需要修改任何 Aurora 上游代码。

2. 环境准备（CANN 8.2+）

2.1 硬件要求

2.2 软件版本矩阵（昇腾官方 CANN 8.2.RC1 配套）

三者版本必须严格对应：CANN 8.2.RC1 ↔ torch 2.5.1 ↔ torch_npu 2.5.1。

2.3 安装步骤

# 1) 安装 CANN 8.2
# https://www.hiascend.com/developer/download/community/result?module=cann
chmod +x Ascend-cann-toolkit_8.2.RC1_linux-aarch64.run
./Ascend-cann-toolkit_8.2.RC1_linux-aarch64.run --install
source /usr/local/Ascend/ascend-toolkit/set_env.sh
echo 'source /usr/local/Ascend/ascend-toolkit/set_env.sh' >> ~/.bashrc

# 2) 下载昇腾官方 whl（不在 PyPI 上）
mkdir -p aurora-npu-adaptation/whl_cache
# - torch 2.5.1       : https://www.hiascend.com/developer/download/community/result?module=pt
# - torchvision 0.20.1: 同上页面（若无 ARM 预编译包，install_env.sh 会自动 pip 安装）
# - torch_npu 2.5.1   : https://gitee.com/ascend/pytorch/releases  → 选择 v2.5.1
# 把所有 whl 放进 whl_cache/

# 3) 一键安装
bash aurora-npu-adaptation/scripts/install_env.sh

# 4) 验证
source aurora-npu-adaptation/venv_npu/bin/activate
python aurora-npu-adaptation/scripts/verify_npu.py

3. 代码改动清单

3.1 必须改动（仅 1 个文件）

用 configs/deepspeed_npu.json 直接替换即可，无需修改其他文件。

3.2 通过运行时补丁完成的「无侵入」改动

3.3 可选的「源码级」改动（最大化性能）

以下改动可消除运行时探测开销，但不是必需：

• 在 aurora/modeling_aurora.py 中将
  decay_weights = 0.5 ** torch.arange(predict_token_num)
  改为
  decay_weights = (0.5 ** torch.arange(predict_token_num, device=hidden_states.device)).to(hidden_states.dtype)
• 在 aurora/util_functions.py 调用 causal_attention_mask 的地方，
  把 causal_attention_mask(n).to(x.device) 改为 causal_attention_mask(n, device=x.device)

4. 推荐适配流程：launcher + 运行时补丁

4.1 三步上手

# 1) 安装环境（见 §2.3）
bash aurora-npu-adaptation/scripts/install_env.sh

# 2) 自检
python aurora-npu-adaptation/scripts/verify_npu.py

# 3) 启动训练 —— 不需要修改任何 Aurora 代码！
bash aurora-npu-adaptation/scripts/run_train_npu.sh

4.2 launcher 工作原理

torchrun runner ─┐
                 │
                 ▼
        aurora_launcher.py
                 │
                 ├─ ① import npu_compat        ← transfer_to_npu 注册 + 算子探测
                 ├─ ② setup_device(local_rank) ← NPU 设备绑定（HCCL 前置）
                 ├─ ③ apply_all_patches()      ← 给 aurora.util_functions / torch.fft 打补丁
                 │
                 └─→ runpy.run_path("Aurora/runner.py")  ← 跑上游代码

关键保证：npu_compat 在 import aurora 之前完成加载，否则 aurora 内的 nn.Module 实例化（含 nn.Parameter）就已经在 CPU 上，后续 .cuda() 调用会触发 transfer_to_npu 但参数已存在 — 行为不再可控。

4.3 备选：手动改 runner.py

如果不想用 launcher，请参考 patches/runner_npu.py 在原 runner.py 顶部添加：

import os, sys
sys.path.insert(0, "/path/to/aurora-npu-adaptation")
import npu_compat
from patches.apply_all_patches import apply_all_patches
apply_all_patches()

然后把 model.cuda() 改成 model.to(npu_compat.get_device(local_rank))。

5. 逐文件适配说明

5.1 runner.py

改动方式 （推荐，零改动）：使用 aurora_launcher.py。

5.2 aurora/util_functions.py

上游文件含 RoPE_decoder、sinusoidal_position_embedding 等被 modeling_aurora.py 直接 import 的函数，整体覆盖会触发 ImportError。

正确做法：调用 patches/util_functions_npu.py 中的两个小粒度补丁函数：

from patches.util_functions_npu import patch_resample, patch_causal_attention_mask
patch_resample()                  # 仅替换 resample（pinv 加 CPU 回退）
patch_causal_attention_mask()     # 仅替换 causal_attention_mask（接受 device 参数）

这两个补丁已被 apply_all_patches() 自动调用。

5.3 aurora/modeling_aurora.py

不要尝试包裹整个 AuroraPatchEmbedding.forward——这会把 NPU 上的 nn.Linear 权重与临时搬到 CPU 的输入混在一起，触发 device mismatch。

正确做法：CANN 8.2+ 已原生支持 torch.fft.rfft/irfft，无需任何 patch。

退一步：若使用更低版本 CANN，apply_all_patches() 会自动调用 patch_fft_to_cpu()，把 torch.fft.rfft/irfft 替换为「输入若在 NPU 则内部走 CPU、输出搬回原设备」的版本。由于 FFT 没有可学习参数，CPU↔NPU 拷贝不会引入 device mismatch。

5.4 aurora/flow_loss.py / prototype_retriever.py / modality_connector.py

无需改动。这些文件已正确使用 tensor.device 进行设备传播。

5.5 trainer/hf_trainer.py

无需改动。HuggingFace Trainer 在 transformers >= 4.40 中已通过 accelerate 支持 NPU。

6. 兼容性问题与解决方案

6.1 算子兼容性矩阵

6.2 FFT 回退策略对比

选择「全局 patch」是因为它不依赖对上游代码的认知，对 CANN 8.0 / 8.2 都安全。CANN 8.2+ 用户该 patch 不会触发。

6.3 NPU + FP16 警告

WARNING [NPU] 检测到 FP16 精度。昇腾 910B 原生支持 BF16 且无需 loss scaling，
强烈建议改用 --precision bf16 以获得更好的稳定性与性能。

如果非要用 FP16：amp_grad_scaler("fp16") 会返回 torch_npu.npu.amp.GradScaler()，与 CUDA 用法一致。

7. DeepSpeed 分布式训练适配

7.1 必须安装昇腾版 DeepSpeed

标准 PyPI 的 deepspeed 不识别 HCCL 后端。即使 JSON 中写了 "communication_backend": "hccl" 也会被忽略（这不是官方字段）。

正确做法：

git clone https://gitee.com/ascend/DeepSpeed.git
cd DeepSpeed
pip install -e .

7.2 JSON 配置改动

只改两处（其他字段保持上游不变）：

{
  "bf16": { "enabled": true },
  "fp16": { "enabled": false }
}

完整文件见 configs/deepspeed_npu.json，已删除上一版本中错误的 "communication_backend" 和误导性的 aio 段。

7.3 HCCL 后端的真正启用方式

# launcher 中已自动完成以下三步：
import torch_npu                                    # ① 注入 hccl backend
import torch.distributed as dist
dist.init_process_group(backend="hccl")             # ② 显式 hccl
torch.npu.set_device(local_rank)                    # ③ 绑卡

启动命令：

bash aurora-npu-adaptation/scripts/run_train_npu.sh \
  --max_steps 1000 \
  --logging_steps 10

启动脚本中已自动设置：

export PYTHONPATH="${ADAPTATION_DIR}:${AURORA_DIR}:${PYTHONPATH}"  # sitecustomize 自动加载 npu_compat
export AURORA_NPU_AUTOLOAD=1
export HCCL_TIMEOUT=1800
export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"

8. 验证与测试

8.1 运行验证脚本

python aurora-npu-adaptation/scripts/verify_npu.py

新版输出包含 9 节：基础环境 / 设备状态 / 必需算子 / 可选算子 / safe_* 自检 / Aurora 级 smoke test / 分布式后端 / DeepSpeed / 汇总。

safe_* 自检会构造一个 mini patch embedding（nn.Linear 在 NPU + FFT 频域掩码），验证输入输出设备一致性 — 复现上一版本 _patched_forward 触发 device mismatch 的场景。

8.2 单步推理 / 单步训练验证

NPROC_PER_NODE=2 bash aurora-npu-adaptation/scripts/run_train_npu.sh \
    --max_steps 5 --logging_steps 1

9. 性能调优建议

9.1 精度

昇腾 910B 强烈推荐 BF16：硬件原生支持、无需 GradScaler、数值范围与 FP32 一致。

9.2 通信

export HCCL_BUFFSIZE=200          # MB
# DeepSpeed config: "overlap_comm": true  (已默认)

9.3 内存

export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"

9.4 算子编译

import torch_npu
torch_npu.npu.set_compile_mode(jit_compile=False)   # 二进制复用，首次后无编译开销

npu_compat.set_npu_options() 已在导入时自动设置。

10. 常见错误排查

10.1 RuntimeError: Expected all tensors to be on the same device

• 原因 A：launcher 未启用 → npu_compat 未在 import aurora 之前加载。 → 改用 scripts/run_train_npu.sh，或检查 PYTHONPATH 是否包含 aurora-npu-adaptation 目录（让 sitecustomize.py 生效）。
• 原因 B：上游代码中常量张量（如 decay_weights）未显式指定 device。 → 查看启动日志中是否有 [NPU-PATCH] 检测到 modeling_aurora.py 中 decay_weights ... 告警，按提示手改。

10.2 HCCL backend not supported

• 没安装昇腾版 DeepSpeed → git clone https://gitee.com/ascend/DeepSpeed && pip install -e .
• 或者 launcher 没有 import torch_npu → 用 aurora_launcher.py 启动。

10.3 torch.fft.rfft is not implemented on NPU

CANN < 8.2 才会出现。两种修复：

• 升级 CANN 到 8.2.RC1；或
• 确保 apply_all_patches() 在 import aurora 之前调用（launcher 默认会做）。

10.4 ImportError: cannot import name 'RoPE_decoder' from 'aurora.util_functions'

你把 patches/util_functions_npu.py 整文件复制覆盖了上游 aurora/util_functions.py。 → 还原上游文件，改用 apply_all_patches() 做运行时补丁。

10.5 pip install torch-npu 失败 / 404

PyPI 上没有昇腾版 torch_npu，必须从 Gitee Releases 下载 whl：
https://gitee.com/ascend/pytorch/releases → 选择 v2.5.1（CANN 8.2.RC1 配套）

10.6 OSError: libtorch_npu.so: version 'GLIBCXX_X.X' not found 或导入 torch_npu 段错误

通常是 torch / torch_npu / CANN 三者版本错配 导致：

• 装了 torch 2.4.0 + torch_npu 2.5.1（不兼容）
• 或 CANN 7.x + torch_npu 2.5.1（CANN 太旧）

解决：严格按 §2.2 矩阵安装 — CANN 8.2.RC1 + torch 2.5.1 + torch_npu 2.5.1。

pip uninstall -y torch torch_npu torchvision
# 重新按 install_env.sh 流程从昇腾官方 whl 安装
bash aurora-npu-adaptation/scripts/install_env.sh