ascend-ai-coding GENERator 昇腾 NPU 迁移 Skill
前置条件
| 项目 | 要求 |
|---|---|
| 硬件 | Ascend910 系列（至少 1 卡） |
| CANN | ≥ 8.2（验证版本 8.3.RC1） |
| Python | 3.11 |
| PyTorch | 2.5.1 |
| torch_npu | 2.5.1 |
迁移步骤
1. 环境初始化
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/
2. 创建 Conda 环境
conda create -n GENERator python=3.11 -y
3. 安装依赖
pip install torch==2.5.1 -i https://repo.huaweicloud.com/repository/pypi/simple/
pip install torch_npu==2.5.1 # 从本地 whl 或华为源安装
pip install numpy==1.26.4 pyyaml decorator attrs psutil absl-py cloudpickle ml-dtypes scipy tornado
pip install transformers==4.49.0 huggingface_hub 'datasets<3.0.0' scikit-learn pandas tqdm pyarrow
4. 代码适配
4.1 添加 NPU 导入（文件顶部）
import torch_npu
from torch_npu.contrib import transfer_to_npu
4.2 替换 CUDA API 调用
| 原始代码 | 替换为 |
|---|---|
torch.cuda.set_device(id) |
torch.npu.set_device(id) |
device = f"cuda:{id}" |
device = f"npu:{id}" |
torch.cuda.empty_cache() |
torch.npu.empty_cache() |
torch.cuda.device_count() |
torch.npu.device_count() |
4.3 修复 from_pretrained 参数
原始代码使用 dtype=dtype，需改为 torch_dtype=dtype：
model = AutoModelForCausalLM.from_pretrained(
args.model_path,
trust_remote_code=True,
torch_dtype=dtype # 原为 dtype=dtype
).to(device)
4.4 多进程子进程适配
GENERator 使用 ProcessPoolExecutor 进行多卡推理。 每个子进程函数内需重新导入 torch_npu：
def process_data_shard(shard_id, ...):
import torch_npu
from torch_npu.contrib import transfer_to_npu
torch.npu.set_device(shard_id)
device = f"npu:{shard_id}"
...
4.5 启用 HF 镜像源（如无法直连 huggingface.co）
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
5. 验证
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export ASCEND_RT_VISIBLE_DEVICES=0
conda activate GENERator
cd /root/GENERator
python src/tasks/downstream/sequence_recovery.py --bf16
验证通过标准：
- 程序正常退出（exit code 0）
- 输出
✅ Completed和📊 Results saved - 生成
./sequence_recovery_results/GENERator-v2-eukaryote-1.2b-base_bfloat16.parquet - 结果非空（30000 行），精度指标合理（Overall Accuracy ~0.515）
6. 注意事项
- Ascend910 不支持 fp64，torch_npu 自动降级为 fp32
- 多卡运行时通过
ASCEND_RT_VISIBLE_DEVICES控制可见设备数量 - 单卡推理速度约 42 seq/s (batch_size=64, bf16, Ascend910)
配套脚本
- 环境与可选本地模型路径预检：
python scripts/validate_generator_env.py --model-path /path/to/model
参考资料
- GENERator 运行时适配要点：
references/runtime-adaptation.md