用途:把 Ankh 系列蛋白质语言模型(ankh_base、ankh_large、ankh3_large、ankh3_xl 及基于 T5EncoderModel/T5ForConditionalGeneration 的同类模型)从默认 CUDA/GPU 运行方式迁移为可在华为 Ascend NPU 上运行、验证并完成最小训练闭环。适用于“标准 Transformers 模型主体 + 设备层写死 CUDA”的场景,不重写模型结构。
主要能力:统一设备层,去除写死的 torch.cuda、.cuda()、cuda:0;适配本地 HuggingFace 权重目录加载;补齐 Ascend 环境说明、验证脚本与训练 smoke test;明确已完成项与仍需用户补环境或硬件权限的部分。
使用方式:先确认仓库根目录、权重目录与 conda 环境,检查 torch_npu、CANN、ASCEND_RT_VISIBLE_DEVICES 是否可用,搜索仓库中的 CUDA 假设点,抽出统一设备层并在入口脚本注入 torch_npu / transfer_to_npu,再用验证脚本做最小验证。最小命令示例:
source "$(conda info --base)/etc/profile.d/conda.sh" conda activate <conda_env> source /usr/local/Ascend/ascend-toolkit/set_env.sh python scripts/verify_ankh_base_npu.py --device auto
验证包括静态校验(python -m compileall src、pytest -q)、运行时环境检查、最小前向验证与训练 smoke test(python scripts/train_smoke_test.py --model ankh_base --device auto --dtype auto)。
适用场景:用户提到 Ankh、蛋白 embedding、pseudo likelihood、T5 蛋白模型,或要把 HuggingFace/PyTorch 模型从 GPU 迁到昇腾 NPU,且模型主体是标准 PyTorch/Transformers。
限制:若仓库包含 .cu/.cuh、依赖 flash_attn/triton/xformers、使用 CUDA 自定义扩展,或需要 nccl 到 hccl 深度改造、ATC/OM 转换、Ascend C 自定义算子、profiling 分析,应改用更通用的深度迁移方案。当前迁移主要针对设备适配层,若未来引入 CUDA 自定义算子需追加专项迁移。