ai-for-science-ankh

Ankh Ascend NPU Migration Skill

This Skill is used to handle Ankh、Ankh base/large、Ankh3 large/XL as well as similar protein language model repositories based on HuggingFace Transformers and PyTorch, migrating the default CUDA/GPU runtime approach to an implementation that can run, be verified, and perform a minimal training loop on Huawei Ascend NPU.

Suitable for migration scenarios of "standard Transformers model body + device layer still hardcoded to CUDA", with focus on environment checks, unified device layer, local weight directory loading, verification script completion, and README documentation consolidation.

The focus of this Skill is not rewriting the model structure, but rather:

Unify the device layer, remove hardcoded torch.cuda、.cuda()、cuda:0
Adapt local HuggingFace weight directory loading
Complete Ascend environment instructions, verification scripts, training smoke test
Clarify what has been completed and what still requires the user to supplement the environment or hardware permissions

For detailed examples and extended instructions, see:

Quick Start

When the user requests "migrate Ankh from GPU to Ascend NPU", prioritize executing in the following order:

Confirm repository root directory <repo_root>, weight directory <weights_root>, environment <conda_env>
Check whether torch_npu、CANN、ASCEND_RT_VISIBLE_DEVICES is available
Search the repository for torch.cuda、.cuda()、cuda:0、torch.device("cuda")
Extract a unified device layer and inject it in the entry script torch_npu / transfer_to_npu
Use scripts/verify_ankh_base_npu.py or the corresponding model script for minimal verification

Minimal command example:

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

1. When to Use This Skill

The user mentions Ankh, protein embedding, pseudo likelihood, T5 protein model.
The user wants to migrate a HuggingFace/PyTorch model from GPU to Ascend NPU.
The repository contains torch.cuda、.cuda()、torch.device("cuda")、cuda:0。
The user requests completion of inference verification, training smoke test, README, weight loading script.
The model body is standard PyTorch / Transformers, rather than a large number of CUDA custom operators.

Typical applicable targets:

ankh_base
ankh_large
ankh3_large
ankh3_xl
Based on T5EncoderModel / T5ForConditionalGeneration protein models

Cases not applicable or requiring an upgraded process:

The repository contains .cu / .cuh
Depends on flash_attn、triton、xformers
Uses CUDA custom extensions or C++ Extension
Requires distributed nccl -> hccl in-depth modification

In these cases, switch to the more general CUDA to Ascend NPU deep migration Skill.

2. Path and Variable Conventions

This Skill should not hardcode paths as absolute paths of a specific machine; by default use the following placeholders:

Variable	Meaning	Recommended value
`<repo_root>`	Ankh repository root directory	Current working directory or detected repository root directory
`<weights_root>`	Local weight root directory	User-provided directory or unified model cache directory
`<weights_dir>`	Single model weight directory	such as `<weights_root>/Ankh_base`
`<conda_env>`	conda environment name or path	Prefer reusing an existing environment; otherwise `ankh-npu`
`<model_name>`	Model name	`ankh_base` / `ankh_large` / `ankh3_large` / `ankh3_xl`
`<device>`	Runtime device parameter	Recommended `auto`
`<dtype>`	Precision parameter	Recommended `auto`

Recommended directory mapping:

ankh_base  -> <weights_root>/Ankh_base
ankh_large -> <weights_root>/Ankh_large
ankh3_large -> <weights_root>/Ankh3_large
ankh3_xl -> <weights_root>/Ankh3_XL

Path inference order:

Whether the current working directory is the Ankh repository
Whether src/ankh、scripts/、README.md
Whether there is ANKH_BASE_PATH、ANKH_LARGE_PATH、ANKH3_LARGE_PATH、ANKH3_XL_PATH
Whether local validation scripts and a unified device layer already exist

Only ask the user for the path when it cannot be inferred from the repository structure or environment variables.

3. Prerequisites

Before performing the migration, confirm the following environment:

Item	Requirement
Hardware	Ascend 910/910B/910C
OS	Ubuntu / openEuler / KylinOS
CANN	with `torch_npu` version match
Python	3.8 - 3.11
PyTorch	compatible with CANN
torch_npu	strictly matching the PyTorch version

3.1 Conda isolated environment

conda create -y -n ankh-npu python=3.10
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate ankh-npu
python -m pip install -U pip setuptools wheel

If the user has already specified an environment name or prefix path, reuse that environment first.

3.2 CANN and NPU self-check

source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate <conda_env>

if [ -n "${ASCEND_TOOLKIT_HOME:-}" ] && [ -f "${ASCEND_TOOLKIT_HOME}/set_env.sh" ]; then
  source "${ASCEND_TOOLKIT_HOME}/set_env.sh"
else
  source /usr/local/Ascend/ascend-toolkit/set_env.sh
fi

npu-smi info
python3 -c "import torch; import torch_npu; print(torch.__version__); print(torch.npu.is_available())"

If torch.npu.is_available() is True, the base environment is available.

3.3 Automatically select an idle NPU card

if [ -z "${ASCEND_RT_VISIBLE_DEVICES:-}" ]; then
  export ASCEND_RT_VISIBLE_DEVICES=0
fi
echo "Using ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES}"

If the user or container has already pinned ASCEND_RT_VISIBLE_DEVICES, do not override it.

4. Migration process

1. 环境检查
-> 2. 搜索 GPU 假设点
-> 3. 提取统一设备层
-> 4. 适配本地权重加载
-> 5. 更新推理/似然/抽取入口
-> 6. 添加验证脚本
-> 7. 增加训练 smoke test
-> 8. 更新 README 与迁移文档

4.1 Code analysis

Check the following files first:

src/ankh/models/ankh_transformers.py
src/ankh/extract.py
src/ankh/likelihood.py
src/ankh/__init__.py
scripts/ inference, validation, and training scripts under
README.md

Search for GPU assumptions:

rg -n "torch\.cuda|\.cuda\(|cuda:0|torch\.device\(['\"]cuda|to\(['\"]cuda" .
rg -n "device=.*cuda|use_gpu|compute_pseudo_likelihood|extract" .

4.2 Unified device layer

It is recommended to extract a common interface instead of repeatedly writing CUDA checks in multiple files.

import torch


def is_npu_available() -> bool:
    npu = getattr(torch, "npu", None)
    return bool(npu is not None and getattr(npu, "is_available", lambda: False)())


def resolve_device(device=None):
    if device is None or str(device).lower() == "auto":
        if is_npu_available():
            return torch.device("npu:0")
        if torch.cuda.is_available():
            return torch.device("cuda:0")
        return torch.device("cpu")

    normalized = str(device).lower()
    if normalized == "npu":
        normalized = "npu:0"
    elif normalized == "cuda":
        normalized = "cuda:0"
    return torch.device(normalized)

Recommended exports:

is_npu_available
has_torch_npu
resolve_device
infer_torch_dtype
bootstrap_torch_npu

4.3 torch_npu injection

If the repository is mainly PyTorch-based, inject at the top of the entry script first:

import torch_npu
from torch_npu.contrib import transfer_to_npu

Injection location priority:

Single-entry projects: main entry file
Multi-entry projects:train.py、inference.py、test.py each entry such as
Reusable packages: top-level __init__.py

4.4 Local weight loading

For Ankh-like projects, prefer HuggingFace directory-level loading rather than a single .bin file. For detailed loading examples, see references/weight-loading.md。

For ankh3_xl , pay extra attention to:

For tokenizer, prefer T5Tokenizer
It is usually a sharded weight directory; pass the directory directly rather than a single weight file.
For embedding extraction, prefer T5EncoderModel
For generation or likelihood scenarios, use T5ForConditionalGeneration

4.5 Inference and training entry adaptation

During migration, ensure that:

Input tensors and the model are on the same device
Inference scripts no longer depend on use_gpu=True and other legacy parameters
compute_pseudo_likelihood() The default device uses auto
On the training side, at least one loss.backward() and optimizer.step()

5. Validation process

5.1 Static validation

python -m compileall src
pytest -q

If there is no test directory, you can only run compileall。

5.2 Runtime environment check

For runtime environment checks, minimal forward validation, and training smoke test references, see references/validation-workflow.md。

5.3 Recommended validation commands

python scripts/verify_ankh_base_npu.py --device auto
python scripts/verify_ankh_large_npu.py --device auto
python scripts/verify_ankh3_large_npu.py --device auto
python scripts/verify_ankh3_xl_npu.py --device auto --weights-dir <weights_root>/Ankh3_XL

5.4 Training smoke test

If the task scope includes training adaptation, at least run:

python scripts/train_smoke_test.py --model ankh_base --device auto --dtype auto

Expected results:

resolved_device is npu:0 or an automatically downgraded device
train_step_ok is true
loss is a finite value
The gradient parameter count is greater than 0

6. README and documentation requirements

After the migration is complete, update the README or add documentation covering at least:

Ascend environment preparation
Automatic device selection example
Local HuggingFace weights directory loading method
Minimal inference verification command
Training smoke test command
Current migration boundaries and known limitations

It is recommended to explain the following boundaries:

The current migration mainly targets the device adaptation layer
If the repository introduces CUDA custom operators in the future, a dedicated migration must be added
Currently applicable to inference, embedding extraction, likelihood, the minimal training loop, and pure PyTorch downstream heads

7. Output requirements

When using this Skill to complete a task, the final reply should include as much as possible:

Migration scope and target model
Path assumptions for the repository, weights, and environment
List of actually modified files
Executed verification and results
Unfinished items and risks
Next-step recommendations

If the current environment cannot perform on-device verification, it should also clearly distinguish:

Code-level migration completed
Syntax-level validation completed
NPU on-device verification not yet completed

8. Recommended artifacts

After completing the migration, it is recommended to also keep the following artifacts:

docs/ascend_npu_migration.md
scripts/verify_ankh_base_npu.py
scripts/verify_ankh_large_npu.py
scripts/verify_ankh3_large_npu.py
scripts/verify_ankh3_xl_npu.py
CodeReview_Results_YYYY-MM-DD.md

If the target is explicitly ankh3_xl, it is recommended to additionally add:

scripts/verify_ankh3_xl_npu.py
docs/ankh3_xl_local_weights.md
examples/ankh3_xl_npu_infer.py

If the task explicitly requires both training and inference adaptation to be completed, it is recommended to additionally add:

scripts/train_smoke_test.py
docs/ankh_train_npu.md
references/validation-workflow.md

9. Upgrade conditions

If any of the following situations is found during analysis, stop the lightweight migration approach that relies only on this Skill and switch to a more complete dedicated solution:

CUDA custom operators appear,.cu、.cuh
Appears flash_attn、triton、xformers
Appears nccl distributed training migration requirement
Need to use ATC, OM conversion, Ascend C custom operators
Need to do profiling, host bound, communication bottleneck analysis

At this point, you can link with a more general Ascend migration or profiling Skill to continue processing.

About this skill