文档
用户指南
NPU 支持

NPU 支持

Twinkle 支持华为昇腾 NPU 进行训练。本指南介绍 NPU 环境下的安装和使用方法。

环境要求

组件版本备注
Python>= 3.11, < 3.13推荐 3.11
昇腾 HDK最新版硬件驱动和固件
CANN Toolkit8.3.RC1+约 10GB 磁盘空间
PyTorch2.7.1必须与 torch_npu 一致
torch_npu2.7.1必须与 PyTorch 一致
注意
torch 和 torch_npu 版本必须完全一致（例如都是 2.7.1）

支持的硬件

  • 昇腾 910 系列
  • 其他兼容的昇腾加速卡

安装

安装 NPU 环境

按照 torch_npu 官方安装指南 安装：

  • 昇腾驱动（HDK）
  • CANN 工具包
  • PyTorch 和 torch_npu

安装 Twinkle

git clone https://github.com/modelscope/twinkle.git
cd twinkle
pip install -e ".[transformers,ray]"

安装 vLLM（可选）

用于 vLLMSampler 支持：

pip install vllm==0.11.0
pip install vllm-ascend==0.11.0rc3
注意
按上述顺序安装，忽略依赖冲突警告。安装前先激活 CANN：source /usr/local/Ascend/ascend-toolkit/set_env.sh

验证安装

import torch
import torch_npu

print(f"PyTorch version: {torch.__version__}")
print(f"torch_npu version: {torch_npu.__version__}")
print(f"NPU available: {torch.npu.is_available()}")
print(f"NPU device count: {torch.npu.device_count()}")

if torch.npu.is_available():
    x = torch.randn(3, 3).npu()
    y = torch.randn(3, 3).npu()
    z = x + y
    print(f"NPU computation test passed: {z.shape}")

快速开始示例

SFT LoRA 微调（4 卡 DP+FSDP）

示例：cookbook/transformers/fsdp2.py

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3
python cookbook/transformers/fsdp2.py

GRPO 强化学习（8 卡）

示例：cookbook/rl/grpo.py

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
python cookbook/rl/grpo.py

DP + FSDP 配置

import numpy as np
from twinkle import DeviceMesh

# 4 卡：DP=2, FSDP=2
device_mesh = DeviceMesh(
    device_type='npu',
    mesh=np.array([[0, 1], [2, 3]]),
    mesh_dim_names=('dp', 'fsdp')
)

并行策略支持

策略说明NPU 支持状态
DP数据并行✅已验证
FSDP全分片数据并行✅已验证
TP张量并行（Megatron）🚧待验证
PP流水线并行（Megatron）🚧待验证
CP上下文并行🚧待验证
EP专家并行（MoE）🚧待验证

功能支持矩阵

功能GPUNPU示例备注
SFT + LoRA✅✅cookbook/transformers/fsdp2.py已验证
GRPO✅✅cookbook/rl/grpo.py已验证
DP 并行✅✅cookbook/transformers/fsdp2.py已验证
FSDP 并行✅✅cookbook/transformers/fsdp2.py已验证
Ray 分布式✅✅cookbook/transformers/fsdp2.py已验证
TorchSampler✅✅cookbook/rl/grpo.py已验证
vLLMSampler✅✅cookbook/rl/grpo.py已验证
全参数微调✅🚧-待验证
QLoRA✅❌-量化不支持
DPO✅🚧-待验证
Megatron TP/PP✅🚧-待验证
Flash Attention✅⚠️-部分算子不支持

图例：✅ 已验证 | 🚧 待验证 | ⚠️ 部分支持 | ❌ 不支持

常见问题

torch_npu 版本不匹配

# 检查版本
python -c "import torch; import torch_npu; print(torch.__version__, torch_npu.__version__)"

# 重新安装匹配版本
pip uninstall torch torch_npu -y
pip install torch==2.7.1
pip install torch_npu-2.7.1-cp311-cp311-linux_aarch64.whl

CANN 兼容性

查看 昇腾社区版本兼容表

调试日志

export ASCEND_GLOBAL_LOG_LEVEL=1
python your_script.py

参考资源

docs