ai-for-science-generator

GENERator Ascend NPU Migration Skill

Prerequisites

Item	Requirement
Hardware	Ascend910 series (at least 1 card)
CANN	≥ 8.2 (verified version 8.3.RC1)
Python	3.11
PyTorch	2.5.1
torch_npu	2.5.1

Migration Steps

1. Environment Initialization

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/

2. Create Conda Environment

conda create -n GENERator python=3.11 -y

3. Install Dependencies

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. Code Adaptation

4.1 Add NPU Import (at the top of the file)

import torch_npu
from torch_npu.contrib import transfer_to_npu

4.2 Replace CUDA API Calls

Original Code	Replace with
`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 Fix from_pretrained Parameters

The original code uses dtype=dtype, needs to be changed to torch_dtype=dtype：

model = AutoModelForCausalLM.from_pretrained(
    args.model_path,
    trust_remote_code=True,
    torch_dtype=dtype  # 原为 dtype=dtype
).to(device)

4.4 Multiprocessing Subprocess Adaptation

GENERator uses ProcessPoolExecutor for multi-card inference. torch_npu needs to be re-imported within each subprocess function:

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 Enable HF Mirror Source (if unable to connect directly to huggingface.co)

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

5. Validation

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

Validation passing criteria:

Program exits normally (exit code 0)
Output ✅ Completed and 📊 Results saved
Generate ./sequence_recovery_results/GENERator-v2-eukaryote-1.2b-base_bfloat16.parquet
Results are non-empty (30000 lines), accuracy metrics are reasonable (Overall Accuracy ~0.515)

6. Notes

Ascend910 does not support fp64, torch_npu automatically downgrades to fp32
When running on multiple cards, use ASCEND_RT_VISIBLE_DEVICES to control the number of visible devices
Single-card inference speed is about 42 seq/s (batch_size=64, bf16, Ascend910)

Companion Scripts

Environment and optional local model path pre-check:python scripts/validate_generator_env.py --model-path /path/to/model

References

GENERator runtime adaptation key points:references/runtime-adaptation.md