ai-for-science-ankh-ascend-npu-skill

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 validated, and do a minimal training loop on Huawei Ascend NPU.

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

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 to local HuggingFace weight directory loading
Complete Ascend environment documentation, validation 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 notes, see:

Quick Start

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

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 the 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 validation

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 validation, 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 that are not applicable or require 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 deep 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; use the following placeholders by default:

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 there is 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	Matches `torch_npu` version
Python	3.8 - 3.11
PyTorch	Compatible with CANN
torch_npu	Strictly matches 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 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 usable.

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 fixed 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, prefer injecting at the top of the entry script:

import torch_npu
from torch_npu.contrib import transfer_to_npu

Injection location priority:

Single-entry project: main entry file
Multi-entry project:train.py、inference.py、test.py for each entry such as
Reusable package: 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 a detailed loading example, see references/weight-loading.md。

For ankh3_xl additional attention is needed:

For the 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, then use T5ForConditionalGeneration

4.5 Inference and training entry adaptation

During migration, ensure:

Input tensors and the model are on the same device
The inference script no longer depends on use_gpu=True legacy parameters such as
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 run only compileall。

5.2 Runtime environment check

For runtime environment check, 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 execute:

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

Expected results:

resolved_device is npu:0 or automatically downgraded device
train_step_ok is true
loss is a finite value
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 weight 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 CUDA custom operators are introduced into the repository in the future, a dedicated migration will need to 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
Verification performed and results
Unfinished items and risks
Next-step recommendations

If the current environment cannot perform real-device verification, the following should also be clearly distinguished:

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

8. Suggested 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 object is clearly ankh3_xl, it is recommended to additionally provide:

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 provide:

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, the lightweight migration approach based solely on this Skill should be stopped, and a more complete dedicated solution should be used instead:

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, a more general Ascend migration or profiling Skill can be linked for further processing.

About this skill