ai-for-science-ai4s-basic

General Ascend NPU Model Migration Skill

This Skill provides a general standardized process for migrating from CUDA to Ascend NPU, suitable for model projects using mainstream frameworks such as PyTorch, TensorFlow, and vLLM. After migration is complete, a dedicated migration Skill for that model should be generated based on the actual working steps, for reuse.

Prerequisites

Before performing migration, confirm that the following environment is ready:

Item	Requirement
Hardware	Ascend910 series (at least 1 card)
OS	openEuler / Ubuntu / KylinOS (aarch64 or x86_64)
CANN	≥ 8.0 (8.2+ or 8.3.RC1 recommended)
Python	3.8 – 3.10 (3.10 recommended)
PyTorch	Matches the CANN version (refer to Huawei version compatibility table)
torch_npu	Consistent with the PyTorch version

Migration Process Overview

0. 常用技巧与环境初始化（高优先级）
→ 1. 环境检查与代码分析
→ 2. 自动迁移注入
→ 3. 手动修改 CUDA 依赖
→ 4. 分布式适配
→ 5. 非 torch 框架兼容（如有）
→ 6. CUDA 内核算子 .cu 文件处理
→ 7. 第三方依赖库适配
→ 8. 适配验证
→ 9. 生成模型专属迁移 Skill

Execute in the order of the following sections, and proceed to the next step only after each step is completed.

---

0. Common Tips and Environment Initialization (execute with high priority before all steps)

0.1 Device and CANN Environment

# 指定可见 NPU 卡
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3

# CANN 默认环境脚本
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 查看 CANN 版本
cat /usr/local/Ascend/ascend-toolkit/latest/version.cfg

# 查看 NPU 设备状态
npu-smi info

0.2 Basic torch_npu Verification

python3 -c "import torch; import torch_npu; a = torch.randn(3, 4).npu(); print(a + a)"

If it can normally output device='npu:0' Tensor, it indicates that the PyTorch + NPU runtime is available.

If an error is reported, troubleshoot in the following order:

1. set_env.sh whether it has been sourced
2. decorator whether runtime dependencies such as are installed (see 0.3)
3. Whether the CANN and torch_npu versions match

0.3 Dependencies and Image Sources

# 华为镜像
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/

# numpy 使用 1.26.4（规避 numpy 2.x 与低版本 CANN 兼容问题）
pip install numpy==1.26.4

# torch_npu 运行常见依赖
pip install decorator attrs psutil absl-py cloudpickle ml-dtypes scipy tornado

0.4 Text Line Endings and Repository Cleanup

find ./ -type f -exec dos2unix {} \;

Used to fix line ending issues brought from a Windows environment.

0.5 Logging and Debugging Environment Variables

# 开启 CANN 详细日志（调试时使用，正常运行关闭以避免性能损失）
export ASCEND_GLOBAL_LOG_LEVEL=1   # 0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR
export ASCEND_SLOG_PRINT_TO_STDOUT=1

# 算子编译缓存，加速二次运行
export ASCEND_AICPU_PATH=/usr/local/Ascend/ascend-toolkit/latest

---

1. Environment Check and Code Analysis

Before making changes, conduct a comprehensive analysis of the project to clarify the migration workload.

1.1 Framework Identification

Determine the framework used by the project; migration strategies differ for different frameworks:

Framework	Migration Method
PyTorch	torch_npu + transfer_to_npu(automatic migration preferred)
TensorFlow	npu_device plugin or npu_bridge
vLLM	Need to confirm the support status of the Ascend version of vLLM, replace the backend
JAX	Currently not directly supported by Ascend; needs to be rewritten as PyTorch or feasibility assessed
PaddlePaddle	Natively supported by Ascend, use paddle_npu plugin

1.2 CUDA Dependency Analysis

Search the project for the following CUDA-specific APIs, and mark locations that need manual handling:

# 搜索 CUDA 特有 API
rg -n "torch\.cuda\.get_device_capability" .
rg -n "torch\.cuda\.get_device_properties.*\.major" .
rg -n "torch\.cuda\.amp" .
rg -n "@torch\.cuda\.amp\.autocast" .
rg -n "torch\.cuda\.mem_get_info" .
rg -n "cuda\.CUDAGraph\|torch\.cuda\.graph" .

# 搜索 .cu / .cuh 文件
find . -name "*.cu" -o -name "*.cuh" | head -20

1.3 Dependency Package Analysis

Check whether requirements.txt / setup.py / pyproject.toml contains CUDA-specific packages:

CUDA-specific Package	Ascend Alternative
flash_attn	torch_npu.npu_fusion_attention
jax[cuda]	Not currently supported; framework replacement required
xformers	Some operators can use torch_npu native attention replacement
apex	torch_npu built-in AMP support, or use PyTorch native AMP
triton	Ascend currently does not support Triton; needs to be rewritten as standard PyTorch operators or AscendC
bitsandbytes	Ascend currently has no equivalent implementation; quantization needs to be disabled or replaced with Ascend quantization tools
deepspeed	Need to use the Ascend-adapted version of DeepSpeed
cupy	Need to replace with numpy/scipy or AscendCL API

1.4 Distributed Analysis

# 检查分布式相关代码
rg -n "DataParallel\b" .           # DP 模式，需改为 DDP
rg -n "DistributedDataParallel" .   # DDP 模式，改 backend 为 hccl
rg -n "nccl" .                      # 需替换为 hccl
rg -n "init_process_group" .        # 检查 backend 参数

Key: If the project uses DataParallel(DP), it must be changed to DistributedDataParallel(DDP), because Ascend NPU does not support DP mode.

1.5 Generate Analysis Report

It is recommended to organize the analysis conclusions into a migration checklist, recording:

• Framework type
• Files and line numbers that need manual modification
• Incompatible third-party libraries and alternatives
• Distributed mode and required changes
• Number of .cu files and complexity assessment

---

2. Automatic Migration Injection (Execute First)

At the top of the entry scripttop(before all other imports), add:

import torch_npu
from torch_npu.contrib import transfer_to_npu

transfer_to_npu will automatically complete the following mappings, with no need for manual line-by-line replacement:

Original CUDA API	Automatic Mapping Target
torch.cuda.is_available()	Returns True (when NPU is available)
torch.Tensor.cuda() / nn.Module.cuda()	.npu()
torch.device('cuda')	torch.device('npu')
DDP backend nccl	hccl
torch.cuda.* series API	torch.npu.*

Injection Location Selection

• If the project has asingle entry(such as main.py、run.py), inject only in that file
• If there aremultiple entries(such as train.py、inference.py、test.py), injection is required in every entry file
• It can also be injected in the project root __init__.py to ensure it takes effect for any import path

---

3. Manually Modify CUDA Dependencies

transfer_to_npu cannot cover all scenarios; the following APIs need manual handling.

3.1 torch.cuda.get_device_properties(...).major

NPU does not have the major attribute (the CUDA compute capability concept); direct access will throw AttributeError。

Fix Pattern: wrap with try-except:

# 原始代码
if stage == "predict" and not (
    torch.cuda.is_available()
    and torch.cuda.get_device_properties(torch.device("cuda")).major >= 8.0
):
    self.use_kernels = False

# 修复：try-except 包裹
if stage == "predict":
    try:
        has_capability = (
            torch.cuda.is_available()
            and torch.cuda.get_device_properties(torch.device("cuda")).major >= 8.0
        )
    except (AttributeError, RuntimeError):
        has_capability = torch.cuda.is_available()
    if not has_capability:
        self.use_kernels = False

3.2 torch.cuda.get_device_capability()

Similarly, NPU has no such concept; it needs to be wrapped or replaced with a fixed value:

# 原始
cap = torch.cuda.get_device_capability()
if cap[0] >= 8:
    use_bf16 = True

# 修复：NPU 默认支持 bf16（Ascend910 系列）
try:
    cap = torch.cuda.get_device_capability()
    use_bf16 = cap[0] >= 8
except (AttributeError, RuntimeError):
    use_bf16 = True  # Ascend910 支持 bf16

3.3 torch.cuda.amp / autocast

NPU supports torch.npu.amp, but it is recommended to uniformly use PyTorch native AMP:

# 推荐写法（兼容 CUDA 和 NPU）
from torch.amp import autocast
with autocast(device_type="npu"):  # 或根据实际设备动态设置
    output = model(input)

3.4 torch.cuda.mem_get_info()

Replace with torch_npu equivalent API:

# 原始
free, total = torch.cuda.mem_get_info()

# 修复
try:
    free, total = torch.cuda.mem_get_info()
except (AttributeError, RuntimeError):
    total = torch.npu.get_device_properties(0).total_memory
    free = total - torch.npu.memory_allocated(0)

---

4. Distributed Adaptation

4.1 DP → DDP Transformation

Ascend NPU does not support torch.nn.DataParallel(DP), it must be changed to DistributedDataParallel（DDP）。

# 原始 DP 写法
model = torch.nn.DataParallel(model)

# 改为 DDP
import torch.distributed as dist
dist.init_process_group(backend="hccl")  # 昇腾使用 hccl
local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.npu.set_device(local_rank)
model = model.npu(local_rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

4.2 Backend Replacement

All backend="nccl" need to be changed to backend="hccl"。 transfer_to_npu is usually handled automatically, but in explicit configuration files (such as DeepSpeed config, Accelerate config) it needs to be modified manually.

4.3 Distributed Startup Command

# 单机多卡
torchrun --nproc_per_node=8 train.py

# 多机多卡（需配置 HCCL 通信）
torchrun --nnodes=2 --nproc_per_node=8 \
    --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train.py

---

5. Non-torch Framework Compatibility

5.1 TensorFlow Projects

# 安装昇腾 TF 插件
# pip install npu_device  (或 npu_bridge，取决于 TF 版本)

from npu_device.compat.v1.npu import npu_config
npu_config()

5.2 vLLM Projects

Confirm whether there is an Ascend-adapted version of vLLM; if so, replace the installation; if not, evaluate switching to another inference framework (such as MindIE, vLLM Ascend fork).

5.3 Mixed-Framework Projects

Some projects may mix PyTorch + other libraries (such as JAX/TF for data processing), and the device mapping for each framework needs to be handled separately, ensuring that all paths route to NPU or CPU.

---

6. Handling .cu Files for CUDA Kernel Operators

6.1 Evaluation Strategy

First check the logic complexity in .cu files:

Complexity	Judgment Criteria	Handling Method
Low	Simple element-wise operations, reduction	Rewrite as PyTorch native operators or AscendC
Medium	Involves shared memory, warp operations	Assess AscendC rewrite cost, or fall back to CPU
High	Deeply dependent on the CUDA ecosystem (cuBLAS, etc.)	Prioritize CPU fallback, use AscendCL when necessary

6.2 CPU Fallback Mode

If adaptation to the NPU is difficult, have this operator switch to CPU execution, and immediately switch back to the NPU after execution completes:

def cuda_kernel_fallback(func, *args, **kwargs):
    # 将输入移到 CPU
    cpu_args = [a.cpu() if isinstance(a, torch.Tensor) else a for a in args]
    cpu_kwargs = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
    # 在 CPU 上执行
    result = func(*cpu_args, **cpu_kwargs)
    # 结果移回 NPU
    if isinstance(result, torch.Tensor):
        return result.npu()
    return result

Note: CPU fallback introduces data transfer overhead and is only suitable for non-hot paths. For performance-sensitive operators, prioritize AscendC rewriting or finding an equivalent torch_npu native operator.

6.3 Disabling Compiled Extensions

If the .cu file is built via setup.py of CUDAExtension , it can be skipped during installation:

# 设置环境变量跳过 CUDA 扩展编译
export FORCE_CUDA=0
export TORCH_CUDA_ARCH_LIST=""
pip install -e . --no-build-isolation

Then ensure there is a fallback path in the code (usually try: import cuda_ext except: use_fallback = True）。

---

7. Third-Party Dependency Library Adaptation

7.1 Flash Attention

On Ascend, use torch_npu.npu_fusion_attention instead:

# 原始
from flash_attn import flash_attn_func
out = flash_attn_func(q, k, v, causal=True)

# 替换
import torch_npu
out = torch_npu.npu_fusion_attention(
    q, k, v,
    head_num=num_heads,
    input_layout="BNSD",  # 根据实际 layout 调整
    atten_mask=causal_mask,
)[0]

Note：npu_fusion_attention the input layout and parameters of are not fully consistent with those of flash_attn , and input_layout(supports "BSH", "BNSD" , etc.) needs to be adjusted according to the project's actual tensor shape.

7.2 PyG Extension Libraries (torch_scatter / torch_sparse / torch_cluster)

Mandatory requirement: all PyG extension libraries must be compiled and installed from source; precompiled wheel methods such as pip install torch_scatter -f https://data.pyg.org/whl/... are prohibited. PyPI / PyG precompiled packages are built for CUDA or x86 CPU, and may have ABI incompatibility or missing operator issues in an aarch64 + NPU environment, so they must be compiled from source:

 source /usr/local/Ascend/ascend-toolkit/set_env.sh
 
 git clone https://github.com/rusty1s/pytorch_scatter.git
 git clone https://github.com/rusty1s/pytorch_sparse.git
 git clone https://github.com/rusty1s/pytorch_cluster.git
 
 # 分别进入目录执行
 cd pytorch_scatter && python setup.py bdist_wheel && pip install dist/*.whl && cd ..
 cd pytorch_sparse  && python setup.py bdist_wheel && pip install dist/*.whl && cd ..
 cd pytorch_cluster && python setup.py bdist_wheel && pip install dist/*.whl && cd ..

8. Adaptation Verification

8.1 Verification Strategy

Verification priority:

1. User-specified: If the prompt explicitly specifies a verification method, execute the user's requirement first
2. Official README: Verify according to the steps in the model's official GitHub README
3. Minimal inference test: Run the simplest inference example to confirm there are no end-to-end errors

8.2 General Verification Steps

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0

# 运行推理/测试命令（根据项目实际情况调整）
python inference.py --num_workers 0

--num_workers 0 can avoid multi-process and NPU driver conflicts in some environments.

8.3 Verification Pass Criteria

• The program exits normally (exit code 0), with no NPU-related errors
• The output result file is non-empty and correctly formatted
• The precision deviation of the results from the CPU/CUDA baseline is within a reasonable range (fp32 scenario difference < 1e-5, fp16/bf16 scenario difference < 1e-3)

8.4 Precision Considerations

• Ascend910 fp64 (double) precision is not supported, torch_npu automatically downgrades to fp32
• If the model involves random sampling (such as diffusion models), per-sample comparison is meaningless; distribution-level metrics should be compared instead
• It is recommended to fix the random seed during the verification phase to improve comparability:

import torch
import numpy as np
import random
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

---

9. Generate a Model-Specific Migration Skill

After the migration runs successfully, generate a complete model migration adaptation Skill based on the actual steps, with the following requirements:

• Reproducible: On any server with the same configuration, any Agent can quickly complete the migration according to this Skill

• Self-contained: Includes complete environment requirements, dependency installation commands, code modifications, and verification steps

• Unified format: Follows the YAML frontmatter + Markdown structure of this Skill

• Record differences: Clearly indicate the adaptation points specific to the model (such as special operator replacement, special configuration modification).

• Huawei mirror source: All pip install steps in the generated Skill must be configured with the Huawei mirror source to ensure that dependencies can be installed quickly in the domestic environment:

  export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/
  

  Or append to each pip install command -i https://repo.huaweicloud.com/repository/pypi/simple/. The generated Skill should set this mirror source at the very beginning of the dependency installation section.

• PyG library source installation: If the project depends on PyG extension libraries (torch_scatter / torch_sparse / torch_cluster, etc.), the generated Skill must use source compilation and installation, and precompiled wheels must not be used. Refer to the source compilation steps in Section 7.2 of this Skill.

The generated Skill should be placed in ~/.kernelcat/skills/<model-name>/SKILL.md。

---

Migration checklist

When migrating to NPU, at least check the following items:

• Has the entry script been injected with transfer_to_npu
• Whether it avoids torch.cuda.get_device_properties(...).major NPU-incompatible access
• Whether it avoids torch.cuda.get_device_capability() NPU-incompatible access
• torch.cuda.amp / autocast Whether it has been adapted
• Whether CUDA-proprietary fused kernels have an NPU fallback
• .cu Whether the files have been handled (rewrite / CPU fallback / skip compilation)
• Whether the distributed mode has been changed from DP to DDP (if applicable)
• Whether the distributed backend has been changed from nccl to hccl
• flash_attn Whether third-party libraries such as these have been replaced
• Whether the inference framework has registered npu accelerator
• Whether the inference command uses --num_workers 0(to avoid multiprocessing issues in some environments)
• __init__.py Whether it is complete (to avoid ModuleNotFoundError）
• Whether CUDA Graph-related code has been disabled
• Environment variables (ASCEND_RT_VISIBLE_DEVICES、set_env.sh) whether configured

---

Common issues

Issue	Cause	Solution
No module named 'decorator'	torch_npu runtime dependency missing	pip install decorator
SetPrecisionMode ... error code 500001	CANN environment not loaded	source set_env.sh
ModuleNotFoundError: xxx	Missing __init__.py	Add an empty __init__.py
torch_scatter Compilation failed	Missing compilation tools	pip install setuptools wheel
double precision warning	Ascend910 does not support fp64	No action needed, automatically downgraded to fp32
Multi-card training hangs	HCCL communication issue	Check ASCEND_RT_VISIBLE_DEVICES and network configuration
RuntimeError: ... FORCE_CUDA	Attempting to compile CUDA extensions	Set export FORCE_CUDA=0
OOM (out of memory)	NPU memory management is different	Reduce batch_size or enable gradient checkpointing
Operator not supported error not supported on NPUAscend	Specific operator not adapted	Use CPU fallback or find an equivalent operator
HCCL ... timeout	Multi-card communication timeout	Increase HCCL_CONNECT_TIMEOUT, check inter-card communication
Performance far below expectations	Operators frequently fall back to CPU	Use torch_npu.profiler to locate hotspots and optimize operators

Supporting scripts

• Basic NPU environment pre-check:python scripts/check_npu_basics.py --device npu:0

Reference materials

• Pre-migration analysis template:references/pre-migration-analysis.md
• PyTorch migration API replacement table:references/pytorch-migration-api-replacement.md
• Migration completion definition and checklist:references/validation-checklist.md