ai-for-science-boltz2

Boltz2 Ascend NPU Migration and Reproduction Guide

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

---

1. Prerequisites

1.1 Hardware and System Requirements

Item	Requirement
NPU Chip	Ascend 910/910B/910C (at least 1 card)
Driver	npu-smi info Normal
CANN Toolkit	8.2.RC1
Operating System	Linux aarch64
Python Environment	conda/miniconda

1.2 Verified Environment

Item	Version/Information
npu-smi	25.0.rc1.3
CANN Toolkit	8.2.RC1
CANN Path	/usr/local/Ascend/ascend-toolkit/set_env.sh
Python	3.11.14
torch	2.5.1 (matches the CANN version)
torch_npu	2.5.1 (matches the CANN version)
conda environment	Boltz-2

Note: The default path for CANN is/usr/local/Ascend/ascend-toolkit/。

2. Weights and Data Files

Boltz2 inference requires the following files to be placed in ~/.boltz/：

File Name	Purpose	Suggested Source
boltz2_conf.ckpt	Boltz2 main structure prediction weights	boltz-community/boltz-2
boltz2_aff.ckpt	Boltz2 affinity head weights	boltz-community/boltz-2
boltz1_conf.ckpt	Boltz1 compatible weights (optional)	boltz-community/boltz-1
ccd.pkl	CCD molecule dictionary	boltz-community/boltz-1
mols.tar	Molecule data package	boltz-community/boltz-2

Suggested directory structure:

~/.boltz/
  boltz1_conf.ckpt
  boltz2_conf.ckpt
  boltz2_aff.ckpt
  ccd.pkl
  mols.tar

Check command:

ls -lh ~/.boltz

---

3. Environment Setup

3.1 Create Environment

source /root/anaconda3/etc/profile.d/conda.sh
conda create -n Boltz-2 python=3.11 -y
conda activate Boltz-2

Note: You may first check whether miniconda or anaconda is already installed on the machine, and then create the virtual environment.

3.2 Install torch / torch_npu

pip install torch==2.5.1 torch_npu==2.5.1
pip install pyyaml numpy decorator cloudpickle ml-dtypes tornado

3.3 Source CANN Environment Variables

source /usr/local/Ascend/ascend-toolkit/set_env.sh

3.4 Basic Check Commands

npu-smi info

cat  /usr/local/Ascend/ascend-toolkit/latest/version.cfg | head -20

source /root/anaconda3/etc/profile.d/conda.sh
conda activate Boltz-2
python -c "import sys, torch, torch_npu; print(sys.version); print(torch.__version__, torch_npu.__version__)"

---

Note: If you encounter torch_npu-related errors, it may be due to a version mismatch between CANN and torch_npu. You can check the current CANN version with cat /home/Ascend/ascend-toolkit/latest/version.cfg | head -20, and check whether other CANN versions matching the current environment are installed locally.

---

4. Source Code Migration Changes

4.1 Overall Approach

1. Use transfer_to_npuat the entry layer to automatically remap CUDA APIs to NPU.
2. Use npu_adapter for device-agnostic encapsulation, replacing CUDA hardcoding in the code.
3. For CUDA-only kernels (cuequivariance/trifast), use a fallback or NPU-equivalent implementation on NPU.
4. Register for PyTorch Lightning npu accelerator。

Download source code

cd /home/workspace
git clone https://github.com/jwohlwend/boltz.git
cd boltz

4.2 Key New Files

File	Purpose
src/boltz/model/npu_adapter.py	NPU compatibility layer (device detection, autocast device type, capability detection, etc.)
src/boltz/model/npu_accelerator.py	Lightning's npu accelerator registration
src/boltz/model/layers/npu_kernels.py	NPU-equivalent implementation of triangle attention / triangular mult
tests/test_npu_migration.py	Migration regression tests

4.3 Key Modified Files and Specific Modification Methods

4.3.1 Modify src/boltz/main.py: enable at the entry torch_npu and transfer_to_npu

Add at the top of the file:

try:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
except ImportError:
    pass

Purpose: automatically map common torch.cuda.* APIs to NPU.

4.3.2 Modify src/boltz/model/models/boltz2.py and src/boltz/model/models/boltz1.py: fix the major attribute issue

In both files' setup() , replace the original capability check:

if stage == "predict" and not (
    torch.cuda.is_available()
    and torch.cuda.get_device_properties(torch.device("cuda")).major >= 8.0
):
    self.use_kernels = False

with:

if stage == "predict":
    try:
        has_capability = (
            torch.cuda.is_available()
            and torch.cuda.get_device_properties(torch.device("cuda")).major >= 8.0
        )
    except (AttributeError, RuntimeError):
        has_capability = torch.cuda.is_available()
    if not has_capability:
        self.use_kernels = False

Reason: NPU device properties do not have CUDA's major field.

4.3.3 Add src/boltz/model/layers/npu_kernels.py: provide NPU-equivalent implementation of CUDA kernel

Create the file and write:

import torch
import torch.nn.functional as F


def npu_triangle_attention(q, k, v, tri_bias, mask, scale):
    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)
    return torch.matmul(attn, v)


def npu_triangle_multiplicative_update(
    x,
    direction,
    mask,
    norm_in_weight,
    norm_in_bias,
    p_in_weight,
    g_in_weight,
    norm_out_weight,
    norm_out_bias,
    p_out_weight,
    g_out_weight,
    eps=1e-5,
):
    x_normed = F.layer_norm(x, [x.shape[-1]], norm_in_weight, norm_in_bias, eps)
    proj = F.linear(x_normed, p_in_weight)
    gate = F.linear(x_normed, g_in_weight).sigmoid()
    x_gated = (proj * gate) * mask.unsqueeze(-1)
    a, b = x_gated.float().chunk(2, dim=-1)
    if direction == "outgoing":
        x_tri = torch.einsum("bikd,bjkd->bijd", a, b)
    else:
        x_tri = torch.einsum("bkid,bkjd->bijd", a, b)
    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)
    return x_out * F.linear(x_normed, g_out_weight).sigmoid()

4.3.4 Modify src/boltz/model/layers/triangular_mult.py: dispatch kernel by device

In kernel_triangular_mult , change to the following structure (the core is NPU-first + ImportError fallback):

@torch.compiler.disable
def kernel_triangular_mult(x, direction, mask, ..., eps):
    if x.device.type == "npu":
        from boltz.model.layers.npu_kernels import npu_triangle_multiplicative_update
        return npu_triangle_multiplicative_update(x, ...)

    try:
        from cuequivariance_torch.primitives.triangle import triangle_multiplicative_update
    except ImportError:
        from boltz.model.layers.npu_kernels import npu_triangle_multiplicative_update
        return npu_triangle_multiplicative_update(x, ...)

    return triangle_multiplicative_update(x, ...)

4.3.5 Modify src/boltz/model/layers/triangular_attention/primitives.py: dispatch attention kernel by device

In kernel_triangular_attn , change to:

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

    try:
        from cuequivariance_torch.primitives.triangle import triangle_attention
        return triangle_attention(q, k, v, tri_bias, mask=mask, scale=scale)
    except ImportError:
        from boltz.model.layers.npu_kernels import npu_triangle_attention
        return npu_triangle_attention(q, k, v, tri_bias=tri_bias, mask=mask, scale=scale)

4.3.6 Add src/boltz/model/npu_accelerator.py: register Lightning NPU accelerator

Create the file and write:

import torch
from pytorch_lightning.accelerators import AcceleratorRegistry
from pytorch_lightning.accelerators.accelerator import Accelerator


class NPUAccelerator(Accelerator):
    def setup_device(self, device):
        if device.type != "npu":
            device = torch.device("npu", 0)
        torch.npu.set_device(device)

    def teardown(self):
        torch.npu.empty_cache()

    @staticmethod
    def parse_devices(devices):
        if isinstance(devices, int):
            return list(range(devices))
        if isinstance(devices, str):
            return [int(d) for d in devices.split(",")]
        return devices

    @staticmethod
    def get_parallel_devices(devices):
        return [torch.device("npu", i) for i in devices]

    @staticmethod
    def auto_device_count():
        return torch.npu.device_count()

    @staticmethod
    def is_available():
        return torch.npu.device_count() > 0


def register_npu_accelerator():
    if "npu" not in AcceleratorRegistry:
        AcceleratorRegistry.register("npu", NPUAccelerator, description="Ascend NPU")

4.3.7 Modify again src/boltz/main.py: CLI and Trainer adaptation

1. --accelerator option add npu：

type=click.Choice(["gpu", "cpu", "tpu", "npu"])

1. At the beginning of predict() , add registration and automatic detection:

from boltz.model.npu_accelerator import register_npu_accelerator
register_npu_accelerator()

if accelerator == "gpu" and not torch.cuda.is_available():
    try:
        import torch_npu
        if torch.npu.is_available():
            accelerator = "npu"
    except ImportError:
        pass

1. Trainer sets up dedicated logic for NPU single-card and bf16:

if accelerator == "npu":
    from pytorch_lightning.strategies import SingleDeviceStrategy
    if isinstance(devices, int) and devices == 1:
        strategy = SingleDeviceStrategy(device=torch.device("npu", 0))

precision_val = 32 if model == "boltz1" else "bf16-mixed"
plugins = []
if accelerator == "npu" and precision_val == "bf16-mixed":
    from pytorch_lightning.plugins.precision import MixedPrecision
    plugins.append(MixedPrecision(precision_val, device="npu"))
    precision_val = None

trainer_kwargs = dict(
    default_root_dir=out_dir,
    strategy=strategy,
    callbacks=[pred_writer],
    accelerator=accelerator,
    devices=devices,
)
if precision_val is not None:
    trainer_kwargs["precision"] = precision_val
if plugins:
    trainer_kwargs["plugins"] = plugins
trainer = Trainer(**trainer_kwargs)

4.4 Install Boltz2 source code

cd /home/workspace/boltz
pip install -e .

5. Inference reproduction (final acceptance)

5.1 Run command

source /usr/local/Ascend/ascend-toolkit/set_env.sh
conda activate Boltz-2
cd /home/workspace/boltz

boltz predict examples/prot_no_msa.yaml \
  --devices 1 \
  --out_dir ./boltz_output_e2e \
  --num_workers 0

Note:--num_workers 0 is a safe configuration on NPU, avoiding crashes caused by DataLoader multiprocessing.

5.2 Expected log keywords

• Auto-detected Ascend NPU(if using the default accelerator)
• Number of failed examples: 0

5.3 Expected output files

ls ./boltz_output_e2e/boltz_results_prot_no_msa/predictions/prot_no_msa/

Should contain at least:

• prot_no_msa_model_0.cif
• plddt_prot_no_msa_model_0.npz
• pae_prot_no_msa_model_0.npz
• pde_prot_no_msa_model_0.npz
• confidence_prot_no_msa_model_0.json

---

6. Common issues

6.1 ValueError: invalid accelerator name: npu

Lightning NPU accelerator not registered. Check whether src/boltz/model/npu_accelerator.py exists, and confirm that predict() executes register_npu_accelerator()。

6.2 ModuleNotFoundError: cuequivariance_torch

This is a missing CUDA-only dependency. The NPU fallback path should be used. Check whether triangular_mult.py and triangular_attention/primitives.py have been added to the NPU dispatch logic.

6.3 DataLoader worker crash / segfault

Add to the inference command --num_workers 0。

6.4 torch.cuda.* related property error

This indicates there is still CUDA hardcoding that has not been migrated. Use the following command to troubleshoot:

rg -n "torch\.cuda|autocast\(\"cuda\"\)" src/boltz

---

7. Criteria for determining reproduction completion

Migration is considered successful if the following 4 conditions are met:

1. Version check passes: CANN/torch/torch_npu/Python match the documentation.
2. boltz predict can run on NPU and Number of failed examples: 0。
3. The output directory contains CIF + pLDDT + PAE + PDE + confidence files.
4. The same input can be inferred repeatedly and produce stable results (small numerical fluctuations are allowed).

Supporting scripts

• Asset and runtime check:python scripts/check_boltz2_assets.py --boltz-home ~/.boltz --check-runtime

References

• Boltz2 migration checklist:references/migration-checklist.md