ai-for-science-boltzgen

BoltzGen Ascend NPU Adaptation Guide

This Skill provides the complete steps for migrating BoltzGen from CUDA to Ascend NPU.

---

1. Prerequisites

1.1 Hardware and Drivers

Item	Requirement
NPU Chip	Ascend 910 (at least 1 card)
NPU Driver	Installed (npu-smi info outputs normally)
CANN Toolkit	8.2.RC1 or a compatible version, installed in /usr/local/Ascend/ascend-toolkit or a similar path
Operating System	Linux aarch64 (Ubuntu or EulerOS)
conda	Miniconda / Anaconda installed

Verification command:

# 验证 NPU 驱动
npu-smi info

# 验证 CANN Toolkit（路径按实际安装位置调整）
ls /usr/local/Ascend/ascend-toolkit

1.2 Required Installation Packages

Package	Version	Description
Python	3.11	conda create
PyTorch	2.1.0	aarch64 whl (must match the torch_npu version)
torch_npu	2.1.0.post17	Compatible with CANN 8.2.RC1

Important: The versions of torch and torch_npu must strictly match, and the torch_npu version must be compatible with the CANN Toolkit version.

1.3 Model Weight Files

The following weight files need to be downloaded in advance and placed in the ~/.cache/ directory: if the directory already contains the following weights, there is no need to download them again.

File Name	Purpose	Source
boltzgen1_diverse.ckpt	Diffusion design model (diversity)	HuggingFace boltzgen/boltzgen-1
boltzgen1_adherence.ckpt	Diffusion design model (fit)	HuggingFace boltzgen/boltzgen-1
boltzgen1_ifold.ckpt	Inverse folding model	HuggingFace boltzgen/boltzgen-1
boltz2_conf_final.ckpt	Folding/confidence model	HuggingFace boltzgen/boltzgen-1
mols.zip	Molecule library (CCD data)	HuggingFace boltzgen/inference-data

If the machine can access the external network, BoltzGen will automatically download from HuggingFace; otherwise, manual transfer is required.

mols.zip How to create it (if you only have the extracted mols directory):

# 进入 mols 目录所在的父目录
cd /path/to/mols_parent
# 打包（注意：直接打包 pkl 文件，不包含父目录名）
cd mols && zip -qr ~/.cache/mols.zip *.pkl

---

2. Environment Setup

2.1 Create conda Environment

conda create -n boltzgen_npu python=3.11 -y
conda activate boltzgen_npu

2.2 Install PyTorch, torch_npu, and runtime dependencies

# 安装 PyTorch（使用本地 whl 或 pip）
pip install torch-2.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

# 安装 torch_npu（使用本地 whl）
pip install torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

# numpy 使用 1.26.4（规避 numpy 2.x 与低版本 CANN 兼容问题）
pip install numpy==1.26.4

# torch_npu 运行常见依赖
pip install decorator attrs psutil absl-py cloudpickle ml-dtypes scipy tornado

2.3 Set CANN Environment Variables

# 路径按实际 CANN Toolkit 安装位置调整
source /usr/local/Ascend/ascend-toolkit/set_env.sh

2.4 Verify NPU Environment

python -c "
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
print(\"torch:\", torch.__version__)
print(\"torch_npu:\", torch_npu.__version__)
print(\"GPU available (remapped to NPU):\", torch.cuda.is_available())
print(\"device count:\", torch.cuda.device_count())
t = torch.randn(3,3).cuda()
print(\"tensor device:\", t.device)
"

Expected output:

torch: 2.1.0
torch_npu: 2.1.0.post17
GPU available (remapped to NPU): True
device count: <你的NPU卡数>
tensor device: npu:0

---

2.5 Text Line Endings and Repository Cleanup

find ./ -type f -exec dos2unix {} \;

Used to fix line ending issues introduced from a Windows environment.

2.6 Handling PyG Component Installation Errors (Source Installation)

When torch_scatter、torch_sparse、torch_cluster installation or runtime errors occur:

git clone https://github.com/rusty1s/pytorch_sparse.git
git clone https://github.com/rusty1s/pytorch_cluster.git
git clone https://github.com/rusty1s/pytorch_scatter.git

# 分别进入目录执行
python setup.py bdist_wheel
pip install dist/*.whl

If PyG does not need to be installed, ignore this section

3. Obtain Source Code and Adapt

3.1 Clone Repository

cd /home/workspace
git clone https://github.com/HannesStark/boltzgen.git
cd boltzgen

If you encounter an inability to access the external network, you can set a mirror source via environment variables and then try accessing again

# 华为镜像
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/

3.2 Adaptation Overview

A total of 6 existing files + need to be modified, and 1 file needs to be added, with a total change volume of +58/-19 lines。

Core approach:

1. transfer_to_npu: Add import torch_npu and from torch_npu.contrib import transfer_to_npuat the top of the two Python entry files, automatically hijacking all torch.cuda.* APIs to torch.npu.*。
2. cuEquivariance NPU equivalent operators: Add npu_kernels.py, in use_kernels=True In the path, dispatch to the NPU native implementation by device type.
3. Fix compatibility issues：get_device_capability return value patch, numpy/torch type comparison fix, CUDA-specific dependency comments.

3.3 Detailed Explanation of File-by-File Modifications

---

Modification 1:src/boltzgen/cli/boltzgen.py(CLI entry point)

Insert the following code at the very top of the file (before line 1)：

try:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
    import torch
    _orig_get_device_capability = torch.cuda.get_device_capability
    def _patched_get_device_capability(device=None):
        cap = _orig_get_device_capability(device)
        return cap if cap is not None else (9, 0)
    torch.cuda.get_device_capability = _patched_get_device_capability
except ImportError:
    pass

Purpose：

• transfer_to_npu: hijack all CUDA APIs to NPU
• get_device_capability patch: torch_npu 2.1.0 returns None for this interface, and PyTorch Lightning in _check_cuda_matmul_precision attempting to unpack it raises a TypeError; after the patch it returns (9, 0) indicating a high-compute-capability device

In the get_artifact_path function, modify the weight loading logic(around line 1375), changing:

        result = huggingface_hub.hf_hub_download(
            repo_id,
            filename,
            repo_type=repo_type,
            library_name="boltzgen",
            force_download=args.force_download,
            token=args.models_token,
            cache_dir=args.cache,
        )
        result = Path(result)

to:

        cache_dir = args.cache or Path.home() / ".cache"
        local_path = Path(cache_dir) / filename
        if local_path.exists() and not args.force_download:
            result = local_path
        else:
            result = huggingface_hub.hf_hub_download(
                repo_id,
                filename,
                repo_type=repo_type,
                library_name="boltzgen",
                force_download=args.force_download,
                token=args.models_token,
                cache_dir=args.cache,
            )
            result = Path(result)

Purpose: preferentially look for local weights from ~/.cache/<filename> , and if not found, fall back to HuggingFace download, solving offline environment issues.

---

Modification 2:src/boltzgen/resources/main.py(subprocess entry point)

Insert the following code at the very top of the file (before line 1)(exactly the same as the CLI entry point):

try:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
    import torch
    _orig_get_device_capability = torch.cuda.get_device_capability
    def _patched_get_device_capability(device=None):
        cap = _orig_get_device_capability(device)
        return cap if cap is not None else (9, 0)
    torch.cuda.get_device_capability = _patched_get_device_capability
except ImportError:
    pass

Why both files need it: BoltzGen's CLI (boltzgen.py) is responsible for configuring the pipeline, then starts subprocess via main.py as an independent subprocess to execute each step. The subprocess has an independent Python runtime and needs to independently initialize transfer_to_npu。

---

Modification 3: Add src/boltzgen/model/layers/npu_kernels.py

Create this file with the following complete contents:

"""NPU-native implementations of cuequivariance_torch operations.

Provides functionally equivalent replacements for:
  1. cuequivariance_torch.primitives.triangle.triangle_attention
  2. cuequivariance_torch.primitives.triangle.triangle_multiplicative_update

These run entirely on Ascend NPU using standard PyTorch ops (no CPU fallback).
"""

import math

import torch
import torch.nn.functional as F
from torch import Tensor


def npu_triangle_attention(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    tri_bias: Tensor,
    mask: Tensor,
    scale: float,
) -> Tensor:
    """NPU-native triangle attention.

    Equivalent to cuequivariance_torch.primitives.triangle.triangle_attention.

    Parameters
    ----------
    q : Tensor  [*, H, Q, C]
    k : Tensor  [*, H, K, C]
    v : Tensor  [*, H, K, C]
    tri_bias : Tensor  [*, 1, H, Q, K] or [*, H, Q, K]
    mask : Tensor  bool [*, Q, K] or broadcastable
    scale : float

    Returns
    -------
    Tensor  [*, H, Q, C]
    """
    q = q * scale

    attn = torch.matmul(q, k.transpose(-1, -2))

    if tri_bias.dim() > attn.dim():
        tri_bias = tri_bias.squeeze(-4)
    attn = attn + tri_bias

    if mask is not None:
        if mask.dim() < attn.dim():
            mask = mask.unsqueeze(-3)
        attn = attn.masked_fill(~mask, float("-inf"))

    attn = F.softmax(attn, dim=-1)
    out = torch.matmul(attn, v)
    return out


def npu_triangle_multiplicative_update(
    x: Tensor,
    direction: str,
    mask: Tensor,
    norm_in_weight: Tensor,
    norm_in_bias: Tensor,
    p_in_weight: Tensor,
    g_in_weight: Tensor,
    norm_out_weight: Tensor,
    norm_out_bias: Tensor,
    p_out_weight: Tensor,
    g_out_weight: Tensor,
    eps: float = 1e-5,
) -> Tensor:
    """NPU-native triangle multiplicative update.

    Equivalent to cuequivariance_torch.primitives.triangle.triangle_multiplicative_update.

    Parameters
    ----------
    x : Tensor  [B, N, N, D]
    direction : str  "outgoing" or "incoming"
    mask : Tensor  [B, N, N]
    norm_in_weight, norm_in_bias : LayerNorm parameters
    p_in_weight : Tensor  [2D, D] projection weight
    g_in_weight : Tensor  [2D, D] gating weight
    norm_out_weight, norm_out_bias : LayerNorm parameters
    p_out_weight : Tensor  [D, D] output projection weight
    g_out_weight : Tensor  [D, D] output gating weight
    eps : float

    Returns
    -------
    Tensor  [B, N, N, D]
    """
    x_normed = F.layer_norm(x, [x.shape[-1]], norm_in_weight, norm_in_bias, eps)
    x_in = x_normed

    proj = F.linear(x_normed, p_in_weight)
    gate = F.linear(x_normed, g_in_weight).sigmoid()
    x_gated = proj * gate

    x_gated = x_gated * mask.unsqueeze(-1)

    a, b = x_gated.float().chunk(2, dim=-1)

    if direction == "outgoing":
        x_tri = torch.einsum("bikd,bjkd->bijd", a, b)
    elif direction == "incoming":
        x_tri = torch.einsum("bkid,bkjd->bijd", a, b)
    else:
        raise ValueError(f"direction must be outgoing or incoming, got {direction}")

    x_out = F.layer_norm(x_tri, [x_tri.shape[-1]], norm_out_weight, norm_out_bias, eps)
    x_out = F.linear(x_out, p_out_weight)

    gate_out = F.linear(x_in, g_out_weight).sigmoid()
    return x_out * gate_out

Purpose: uses PyTorch standard operators (matmul/einsum/softmax/layer_norm/linear) to implement triangle attention and triangle multiplicative update fully equivalent in functionality to the CUDA-specific library cuequivariance_torch, all executed natively on NPU via torch_npu, without involving CPU fallback.

---

Modification 4:src/boltzgen/model/layers/triangular.py

In the _kernel_triangular_mult function body's very beginning（eps: float, after the parameter,try: before) insert the NPU dispatch logic:

    if x.device.type == "npu":
        from boltzgen.model.layers.npu_kernels import npu_triangle_multiplicative_update
        return npu_triangle_multiplicative_update(
            x, direction=direction, mask=mask,
            norm_in_weight=norm_in_weight, norm_in_bias=norm_in_bias,
            p_in_weight=p_in_weight, g_in_weight=g_in_weight,
            norm_out_weight=norm_out_weight, norm_out_bias=norm_out_bias,
            p_out_weight=p_out_weight, g_out_weight=g_out_weight, eps=eps,
        )

Modified function structure：

def _kernel_triangular_mult(x, *, direction, mask, ...):
    # --- 新增：NPU 分发 ---
    if x.device.type == "npu":
        from boltzgen.model.layers.npu_kernels import npu_triangle_multiplicative_update
        return npu_triangle_multiplicative_update(...)
    # --- 原有：CUDA 路径 ---
    try:
        from cuequivariance_torch.primitives.triangle import ...
    except ModuleNotFoundError:
        raise RuntimeError(...)
    return _triangle_multiplicative_update(...)

---

Modification 5:src/boltzgen/model/layers/triangular_attention/primitives.py

In the kernel_triangular_attn function body's very beginning（@torch.compiler.disable under the decorator,from cuequivariance_torch before) insert the NPU dispatch logic:

    if q.device.type == "npu":
        from boltzgen.model.layers.npu_kernels import npu_triangle_attention
        return npu_triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)

Modified function：

@torch.compiler.disable
def kernel_triangular_attn(q, k, v, tri_bias, mask, scale):
    if q.device.type == "npu":
        from boltzgen.model.layers.npu_kernels import npu_triangle_attention
        return npu_triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)
    from cuequivariance_torch.primitives.triangle import triangle_attention
    return triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)

---

Modification 6:src/boltzgen/task/analyze/analyze.py

Around line 928, change the original mixed numpy/torch comparison:

                if (
                    not len(folded["res_type"].squeeze())
                    == len(feat_design["res_type"])
                    or not (
                        folded["res_type"].squeeze() == feat_design["res_type"]
                    ).all()
                ):

to:

                _folded_res = torch.from_numpy(folded["res_type"].squeeze())
                if (
                    _folded_res.shape != feat_design["res_type"].shape
                    or not (_folded_res == feat_design["res_type"]).all()
                ):

Reason：folded["res_type"] is a numpy array,feat_design["res_type"] is a torch tensor. Under numpy 1.26 + torch 2.1, comparing the two directly with == returns the Python scalar bool instead of an array, causing the subsequent .all() call to fail. Converting both to torch tensors before comparing resolves this.

---

Modification 7:pyproject.toml

Comment out 4 CUDA-specific dependencies (these packages cannot be installed in an NPU environment):

    # "nvidia-ml-py>=12.535.133",  # CUDA-only
    # "cuequivariance_ops_cu12>=0.5.0",  # CUDA-only
    # "cuequivariance_ops_torch_cu12>=0.5.0",  # CUDA-only
    # "cuequivariance_torch>=0.5.0",  # CUDA-only

---

IV. Installing the Project

conda activate boltzgen_npu
source /usr/local/Ascend/ascend-toolkit/set_env.sh
cd /home/workspace/boltzgen
pip install -e .

---

V. Inference Testing

5.1 Run Command

conda activate boltzgen_npu
source /usr/local/Ascend/ascend-toolkit/set_env.sh
cd /home/workspace/boltzgen

boltzgen run example/vanilla_protein/1g13prot.yaml \
  --output workbench/test_run \
  --protocol protein-anything \
  --num_designs 100 \
  --budget 5 \
  --devices 1 \
  --num_workers 0

Parameter Description：

• --num_workers 0: In an NPU environment, DataLoader multiprocessing conflicts with device initialization, so it must be set to 0 to load data using the main process
• --num_designs 10: generate 10 candidate designs
• --budget 2: final output of the Top-2 ranked designs
• All other parameters remain consistent with the original README; there is no need to specify a weight path (automatically from ~/.cache/ Discovery)

5.2 Expected Results

All 6 Pipeline steps should complete successfully:

Step design completed successfully in ~134s
Step inverse_folding completed successfully in ~34s
Step folding completed successfully in ~141s
Step design_folding completed successfully in ~122s
Step analysis completed successfully in ~10s
Step filtering completed successfully in ~16s

5.3 Verify Output

ls workbench/test_run/final_ranked_designs/

Expected to contain:

all_designs_metrics.csv           # 全部设计指标
final_designs_metrics_2.csv       # Top-2 设计指标
final_2_designs/                  # Top-2 排名设计 (.cif)
  rank01_xxx.cif
  rank02_xxx.cif
intermediate_ranked_10_designs/   # 全部排名设计
results_overview.pdf              # 结果概览图

5.4 Unit Tests

pip install pytest
python -m pytest tests/ -v

Expected:44 passed, 0 failed

---

VI. Known Limitations and Notes

1. --num_workers 0: Under NPU, DataLoader multiprocessing will cause worker crashes, must be set to 0
2. transfer_to_npu Disable TorchScript：torch.jit.script and torch.jit.script_method will be disabled
3. version compatibility: The versions of torch / torch_npu / CANN Toolkit must strictly match
4. use_kernels: By default, auto mode, NPU automatically uses the npu_kernels.py path; you can also explicitly pass --use_kernels true or --use_kernels false

Supporting Scripts

• Cached weights and runtime check:python scripts/check_boltzgen_assets.py --cache-dir ~/.cache --check-runtime

References

• BoltzGen runtime adaptation notes:references/runtime-notes.md