feat: NPU (Ascend 910B3) support for RFdiffusion inference - #459
Conversation
Adapt RFdiffusion to run on Huawei Ascend NPU (Ascend 910B3) with
torch_npu, enabling protein structure generation on NPU hardware.
Changes (8 files, +114 -16 lines):
1. rfdiffusion/__init__.py:
- Global torch.npu.config.allow_internal_format = False
- Patch torch.cdist for NPU (manual bmm + sqrt implementation)
2. rfdiffusion/inference/model_runners.py:
- Device selection: torch.cuda -> torch.npu
3. rfdiffusion/Track_module.py:
- autocast: torch.cuda.amp.autocast -> torch.amp.autocast(device_type='npu')
4. scripts/run_inference.py:
- NPU device detection (priority NPU > CUDA > CPU)
- torch.npu.empty_cache() support
- TRB metadata records actual device name
5-8. SE3Transformer (basis.py, attention.py, convolution.py, norm.py):
- nvtx range three-tier fallback: CUDA nvtx -> NPU record_function -> no-op
- dgl.ops.e_dot_v replaced with manual implementation (Ascend SDDMM
only supports lhs_target=0)
Verified: RFdiffusion NPU pLDDT=0.9885, CPU pLDDT=0.9885, diff=0.000005
NPU speedup: 15.35x over CPU (43.3s vs 664.7s for 20-step inference)
Environment: PyTorch 2.7.1+cpu, torch_npu 2.7.1.post4, CANN 8.2.RC1,
DGL 2.5 (Ascend), Ascend 910B3
|
Thank you for this PR. There are newer version of RFdiffusion (most notably RFD3) whose dependencies are more up to date. Do you run into similar issues trying to RFD3 on the NPU resources you have access to? |
rclune
left a comment
rclune
left a comment
There was a problem hiding this comment.
Thank you for this PR. I have some concerns around ensuring that the changes made will still allow run_inference to work with GPU/CUDA systems. Please make sure that any NPU-specific additions are wrapped in conditionals or try/except clauses to ensure that the GPU functionality is not changed.
| import re | ||
| import os, time, pickle | ||
| import torch | ||
| import torch_npu |
There was a problem hiding this comment.
Many users of RFD only have access to GPU/CPU resources. Adding this import statement not in a try/except will make it impossible to run on GPU systems. I suggest wrapping it like you do in __init__.py.
| self._log = logging.getLogger(__name__) | ||
| if torch.cuda.is_available(): | ||
| self.device = torch.device("cuda") | ||
| if torch.npu.is_available(): |
There was a problem hiding this comment.
Please back in the check for torch.cuda.is_available() as a fallback. Right now only NPU is checked, so GPU users will be silently placed on CPUs. You can do something similar to what you have on lines 46-50 of run_inference.py.
| nn.init.zeros_(self.embed_e2.bias) | ||
|
|
||
| @torch.cuda.amp.autocast(enabled=False) | ||
| @torch.amp.autocast(device_type="npu", enabled=False) |
There was a problem hiding this comment.
This hardcodes device_type="npu", so on CUDA systems this decorator no longer disables autocast, meaning torch.cuda.amp.autocast(enabled=False)'s effect is lost. Should be conditional on the active device.
|
|
||
| with nvtx_range('attention dot product + softmax'): | ||
| # Compute attention weights (softmax of inner product between key and query) | ||
| edge_weights = dgl.ops.e_dot_v(graph, key, query).squeeze(-1) |
There was a problem hiding this comment.
The removal of this for a manual implementation should be conditioned on the device type, similar to how the nvtx changes are.
2 participants
Description
This PR adapts RFdiffusion to run on Huawei Ascend NPU (Ascend 910B3) with
torch_npu, enabling protein structure generation on NPU hardware.Background
RFdiffusion depends on PyTorch + DGL + e3nn + SE3Transformer, all originally CUDA-only. This PR replaces CUDA-specific calls with NPU-compatible alternatives and patches operators not supported on NPU.
Changes (8 files, +114 -16 lines)
rfdiffusion/__init__.pyallow_internal_format=False+torch.cdistNPU patch (manualbmm + sqrt)rfdiffusion/inference/model_runners.pytorch.cuda→torch.npurfdiffusion/Track_module.pytorch.cuda.amp.autocast→torch.amp.autocast(device_type="npu")scripts/run_inference.pytorch.npu.empty_cache(), device metadataSE3Transformer/model/basis.pyrecord_function→ no-opSE3Transformer/model/layers/attention.pye_dot_vmanual implementationSE3Transformer/model/layers/convolution.pySE3Transformer/model/layers/norm.pyKey Adaptations
torch.cuda.*calls replaced withtorch.npu.*equivalentstorch.npu.config.allow_internal_format = Falseto fixtorch.cattensor format conflicts in SE3Transformertorch.cdistnatively; patched with||x1-x2|| = sqrt(||x1||² + ||x2||² - 2*x1·x2^T)torch.cuda.nvtx.rangeimport succeeds but throws at runtime on NPU; three-tier fallback withtorch.autograd.profiler.record_functionfor NPU profiling supportlhs_target=0; replaced with manual(key * query[dst]).sum(dim=-1)Prerequisites
copy_rhsfix (see dgl-ascend PR #24)Verification
Type