diff --git a/env/SE3Transformer/se3_transformer/model/basis.py b/env/SE3Transformer/se3_transformer/model/basis.py index 74f04a0f..46d54e04 100644 --- a/env/SE3Transformer/se3_transformer/model/basis.py +++ b/env/SE3Transformer/se3_transformer/model/basis.py @@ -29,7 +29,21 @@ import torch import torch.nn.functional as F from torch import Tensor -from torch.cuda.nvtx import range as nvtx_range +try: + from torch.cuda.nvtx import range as _nvtx_range_cuda + try: + with _nvtx_range_cuda("init"): + pass + nvtx_range = _nvtx_range_cuda + except Exception: + raise ImportError("NVTX not available at runtime") +except (ImportError, RuntimeError): + from torch.autograd.profiler import record_function + from contextlib import contextmanager + @contextmanager + def nvtx_range(msg, *args, **kwargs): + with record_function(msg): + yield from se3_transformer.runtime.utils import degree_to_dim diff --git a/env/SE3Transformer/se3_transformer/model/layers/attention.py b/env/SE3Transformer/se3_transformer/model/layers/attention.py index 091525e9..ed7b3808 100644 --- a/env/SE3Transformer/se3_transformer/model/layers/attention.py +++ b/env/SE3Transformer/se3_transformer/model/layers/attention.py @@ -34,7 +34,21 @@ from se3_transformer.model.layers.convolution import ConvSE3, ConvSE3FuseLevel from se3_transformer.model.layers.linear import LinearSE3 from se3_transformer.runtime.utils import degree_to_dim, aggregate_residual, unfuse_features -from torch.cuda.nvtx import range as nvtx_range +try: + from torch.cuda.nvtx import range as _nvtx_range_cuda + try: + with _nvtx_range_cuda("init"): + pass + nvtx_range = _nvtx_range_cuda + except Exception: + raise ImportError("NVTX not available at runtime") +except (ImportError, RuntimeError): + from torch.autograd.profiler import record_function + from contextlib import contextmanager + @contextmanager + def nvtx_range(msg, *args, **kwargs): + with record_function(msg): + yield class AttentionSE3(nn.Module): @@ -78,7 +92,10 @@ def forward( 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) + # Use manual implementation for Ascend NPU compatibility (e_dot_v not supported) + dst_nodes = graph.edges()[1] + query_dst = query[dst_nodes] + edge_weights = (key * query_dst).sum(dim=-1) edge_weights /= np.sqrt(self.key_fiber.num_features) edge_weights = edge_softmax(graph, edge_weights) edge_weights = edge_weights[..., None, None] diff --git a/env/SE3Transformer/se3_transformer/model/layers/convolution.py b/env/SE3Transformer/se3_transformer/model/layers/convolution.py index c3de0153..a5f6fec7 100644 --- a/env/SE3Transformer/se3_transformer/model/layers/convolution.py +++ b/env/SE3Transformer/se3_transformer/model/layers/convolution.py @@ -31,7 +31,21 @@ import torch.nn as nn from dgl import DGLGraph from torch import Tensor -from torch.cuda.nvtx import range as nvtx_range +try: + from torch.cuda.nvtx import range as _nvtx_range_cuda + try: + with _nvtx_range_cuda("init"): + pass + nvtx_range = _nvtx_range_cuda + except Exception: + raise ImportError("NVTX not available at runtime") +except (ImportError, RuntimeError): + from torch.autograd.profiler import record_function + from contextlib import contextmanager + @contextmanager + def nvtx_range(msg, *args, **kwargs): + with record_function(msg): + yield from se3_transformer.model.fiber import Fiber from se3_transformer.runtime.utils import degree_to_dim, unfuse_features diff --git a/env/SE3Transformer/se3_transformer/model/layers/norm.py b/env/SE3Transformer/se3_transformer/model/layers/norm.py index acbe23d7..0732e5a3 100644 --- a/env/SE3Transformer/se3_transformer/model/layers/norm.py +++ b/env/SE3Transformer/se3_transformer/model/layers/norm.py @@ -27,7 +27,21 @@ import torch import torch.nn as nn from torch import Tensor -from torch.cuda.nvtx import range as nvtx_range +try: + from torch.cuda.nvtx import range as _nvtx_range_cuda + try: + with _nvtx_range_cuda("init"): + pass + nvtx_range = _nvtx_range_cuda + except Exception: + raise ImportError("NVTX not available at runtime") +except (ImportError, RuntimeError): + from torch.autograd.profiler import record_function + from contextlib import contextmanager + @contextmanager + def nvtx_range(msg, *args, **kwargs): + with record_function(msg): + yield from se3_transformer.model.fiber import Fiber diff --git a/rfdiffusion/Track_module.py b/rfdiffusion/Track_module.py index 27511e5d..749e901a 100644 --- a/rfdiffusion/Track_module.py +++ b/rfdiffusion/Track_module.py @@ -233,7 +233,7 @@ def reset_parameter(self): nn.init.zeros_(self.embed_e1.bias) nn.init.zeros_(self.embed_e2.bias) - @torch.cuda.amp.autocast(enabled=False) + @torch.amp.autocast(device_type="npu", enabled=False) def forward(self, msa, pair, R_in, T_in, xyz, state, idx, motif_mask, cyclic_reses=None, top_k=64, eps=1e-5): B, N, L = msa.shape[:3] diff --git a/rfdiffusion/__init__.py b/rfdiffusion/__init__.py index e69de29b..24acbd53 100644 --- a/rfdiffusion/__init__.py +++ b/rfdiffusion/__init__.py @@ -0,0 +1,29 @@ +import torch + +try: + import torch_npu + torch.npu.config.allow_internal_format = False + + # Patch torch.cdist for NPU (NPU does not support cdist natively) + _orig_cdist = torch.cdist + def _npu_cdist(x1, x2, p=2.0, compute_mode='use_mm_for_euclid_dist_if_necessary', **kwargs): + if x1.device.type == 'npu' or (x2 is not None and hasattr(x2, 'device') and x2.device.type == 'npu'): + if x1.dim() == 2: + x1 = x1.unsqueeze(0) + x2 = x2.unsqueeze(0) + squeeze = True + else: + squeeze = False + x1_sq = (x1 * x1).sum(dim=-1, keepdim=True) + x2_sq = (x2 * x2).sum(dim=-1, keepdim=True) + cross = torch.bmm(x1, x2.transpose(-1, -2)) + dist_sq = x1_sq + x2_sq.transpose(-1, -2) - 2 * cross + dist_sq = dist_sq.clamp(min=0) + result = torch.sqrt(dist_sq + 1e-12) + if squeeze: + result = result.squeeze(0) + return result + return _orig_cdist(x1, x2, p=p, compute_mode=compute_mode, **kwargs) + torch.cdist = _npu_cdist +except ImportError: + pass diff --git a/rfdiffusion/inference/model_runners.py b/rfdiffusion/inference/model_runners.py index 2bd0d530..6c97b3f6 100644 --- a/rfdiffusion/inference/model_runners.py +++ b/rfdiffusion/inference/model_runners.py @@ -48,8 +48,8 @@ def initialize(self, conf: DictConfig) -> None: """ self._log = logging.getLogger(__name__) - if torch.cuda.is_available(): - self.device = torch.device("cuda") + if torch.npu.is_available(): + self.device = torch.device("npu") else: self.device = torch.device("cpu") needs_model_reload = ( diff --git a/scripts/run_inference.py b/scripts/run_inference.py index 3ebb5e3f..221e26e0 100755 --- a/scripts/run_inference.py +++ b/scripts/run_inference.py @@ -18,6 +18,8 @@ import re import os, time, pickle import torch +import torch_npu +torch.npu.config.allow_internal_format = False from omegaconf import OmegaConf import hydra import logging @@ -41,13 +43,16 @@ def main(conf: HydraConfig) -> None: if conf.inference.deterministic: make_deterministic() - # Check for available GPU and print result of check - if torch.cuda.is_available(): + # Check for available NPU/GPU and print result of check + if torch.npu.is_available(): + device_name = torch.npu.get_device_name(torch.npu.current_device()) + log.info(f"Found NPU with device_name {device_name}. Will run RFdiffusion on {device_name}") + elif torch.cuda.is_available(): device_name = torch.cuda.get_device_name(torch.cuda.current_device()) log.info(f"Found GPU with device_name {device_name}. Will run RFdiffusion on {device_name}") else: log.info("////////////////////////////////////////////////") - log.info("///// NO GPU DETECTED! Falling back to CPU /////") + log.info("///// NO GPU/NPU DETECTED! Falling back to CPU /////") log.info("////////////////////////////////////////////////") # Initialize sampler and target/contig. @@ -148,9 +153,11 @@ def main(conf: HydraConfig) -> None: trb = dict( config=OmegaConf.to_container(sampler._conf, resolve=True), plddt=plddt_stack.cpu().numpy(), - device=torch.cuda.get_device_name(torch.cuda.current_device()) - if torch.cuda.is_available() - else "CPU", + device=torch.npu.get_device_name(torch.npu.current_device()) + if torch.npu.is_available() + else (torch.cuda.get_device_name(torch.cuda.current_device()) + if torch.cuda.is_available() + else "CPU"), time=time.time() - start_time, ) if hasattr(sampler, "contig_map"): @@ -188,8 +195,11 @@ def main(conf: HydraConfig) -> None: chain_ids=sampler.chain_idx, ) - if conf.inference.empty_cache_per_design and torch.cuda.is_available(): - torch.cuda.empty_cache() + if conf.inference.empty_cache_per_design: + if torch.npu.is_available(): + torch.npu.empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() log.info(f"Finished design in {(time.time()-start_time)/60:.2f} minutes")