SuperAPI AI 工具集

ai-for-science-ankh

智能体与模型

查看项目

更新于 10-15 17:19

技能介绍

用途:本 Skill 面向 Ankh、Ankh base/large、Ankh3 large/XL 及同类基于 HuggingFace Transformers 与 PyTorch 的蛋白质语言模型仓库,把默认 CUDA/GPU 运行方式迁移为可在华为 Ascend NPU 上运行、验证并完成最小训练闭环的实现。适用于“标准 Transformers 模型主体 + 设备层写死 CUDA”的场景,重点覆盖环境检查、统一设备层、本地权重目录加载、验证脚本补齐和 README 文档沉淀,而非重写模型结构。

主要能力:统一设备层,去掉写死的 torch.cuda、.cuda()、cuda:0;适配本地 HuggingFace 权重目录加载;补齐 Ascend 环境说明、验证脚本与训练 smoke test;明确已完成项与仍需用户补环境或硬件权限的部分。典型对象包括 ankh_base、ankh_large、ankh3_large、ankh3_xl 及基于 T5EncoderModel / T5ForConditionalGeneration 的蛋白模型。

使用方式:先确认仓库根目录、权重目录与环境,检查 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。

适用场景:推理、embedding 抽取、pseudo likelihood 计算、最小训练闭环与纯 PyTorch 下游头。

限制:若仓库包含 .cu / .cuh、依赖 flash_attn、triton、xformers、使用 CUDA 自定义扩展,或需要 nccl 到 hccl 的分布式深度改造,则不适用本轻量迁移方式,应改用更通用的 CUDA 到昇腾 NPU 深度迁移 Skill。具体硬件、CANN、PyTorch 与 torch_npu 版本匹配要求以官方资料为准,本资料未提供完整版本对应表。