Repository: aiqm/torchani (language: python)

Issue:
CUAEV RuntimeError: CUDA error: an illegal memory access was encountered
Hi, when using CUAEV on <code>torch.device('cuda:1')</code>, it gave me the following error information:

```bash
Traceback (most recent call last):
  File "ani_e.py", line 208, in <module>
    rmse = validate()
  File "ani_e.py", line 194, in validate
    _, predicted_energies = model((species, coordinates))
  File "/home/jack/miniconda3/envs/CUAEV/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jack/miniconda3/envs/CUAEV/lib/python3.7/site-packages/torchani-2.3.dev60+ga85e330-py3.7-linux-x86_64.egg/torchani/nn.py", line 111, in forward
    input_ = module(input_, cell=cell, pbc=pbc)
  File "/home/jack/miniconda3/envs/CUAEV/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/jack/miniconda3/envs/CUAEV/lib/python3.7/site-packages/torchani-2.3.dev60+ga85e330-py3.7-linux-x86_64.egg/torchani/nn.py", line 60, in forward
    atomic_energies = self._atomic_energies((species, aev))
  File "/home/jack/miniconda3/envs/CUAEV/lib/python3.7/site-packages/torchani-2.3.dev60+ga85e330-py3.7-linux-x86_64.egg/torchani/nn.py", line 72, in _atomic_energies
    output = aev.new_zeros(species_.shape)
**RuntimeError: CUDA error: an illegal memory access was encountered**
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
```


If I use <code>torch.device('cuda:0')</code> with CUAEV, the the code runs well. Otherwise, using <code>torch.device('cuda:1')</code> without CUAEV also works. This looks like a issue related to CUAEV on a device with multiple GPUs. Could anyone provide some debug tips for that? Thanks!

To sum up:
- <code>torch.device('cuda:0')</code>  with CUAEV: runs well
- <code>torch.device('cuda:1')</code> with CUAEV: yields the above error
- <code>torch.device('cuda:1')</code> without CUAEV: runs well


Discussion / hints:
Hi, thanks for the report!
Could you please provide a minimal reproducible example?
The aev should be created on the same device as the input coordinates: https://github.com/aiqm/torchani/blob/master/torchani/cuaev/aev.cu#L755
Hi, thanks for the reply. I created a [folder](https://cmu.box.com/s/48d6e26rpej1m2owwjaedde8sxu8bcb8) containing a reproducible example.

This folder contains the following three files:
- **sampe.h5**: a dataset containing 50 data points;
- **ani_2x_f.py**: the script that training an ani_2x model with force field;
- **env.yml**: the yaml file for my conda environment.

You can get the same error message when executing the following in the folder:
 `python ani_2x_f.py`


The GPU device can be adjusted in the python file. 

Sorry for the late sample😂I got obsessed with another script the past week. Please let me know how things are going, thank you very much!
Thanks! I Was able to reproduce by the following, will try to fix it soon

```python
import os
import torch
import torchani
import unittest
import pickle
import copy
from torchani.testing import TestCase, make_tensor


class TestCUAEV(TestCase):

    def setUp(self):
        self.tolerance = 5e-5
        i = 1
        self.device = torch.device(f'cuda:{i}' if torch.cuda.is_available() else 'cpu')
        print('{}: {}'.format(i, torch.cuda.get_device_name(f'cuda:{i}')))
        Rcr = 5.2000e+00
        Rca = 3.5000e+00
        EtaR = torch.tensor([1.6000000e+01], device=self.device)
        ShfR = torch.tensor([9.0000000e-01, 1.1687500e+00, 1.4375000e+00, 1.7062500e+00, 1.9750000e+00, 2.2437500e+00, 2.5125000e+00, 2.7812500e+00, 3.0500000e+00, 3.3187500e+00, 3.5875000e+00, 3.8562500e+00, 4.1250000e+00, 4.3937500e+00, 4.6625000e+00, 4.9312500e+00], device=self.device)
        Zeta = torch.tensor([3.2000000e+01], device=self.device)
        ShfZ = torch.tensor([1.9634954e-01, 5.8904862e-01, 9.8174770e-01, 1.3744468e+00, 1.7671459e+00, 2.1598449e+00, 2.5525440e+00, 2.9452431e+00], device=self.device)
        EtaA = torch.tensor([8.0000000e+00], device=self.device)
        ShfA = torch.tensor([9.0000000e-01, 1.5500000e+00, 2.2000000e+00, 2.8500000e+00], device=self.device)
        num_species = 4
        self.aev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species)
        self.cuaev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species, use_cuda_extension=True)
        self.nn = torch.nn.Sequential(torch.nn.Linear(384, 1, False)).to(self.device)
        self.radial_length = self.aev_computer.radial_length

    def testSimple(self):
        coordinates = torch.tensor([
            [[0.03192167, 0.00638559, 0.01301679],
             [-0.83140486, 0.39370209, -0.26395324],
             [-0.66518241, -0.84461308, 0.20759389],
             [0.45554739, 0.54289633, 0.81170881],
             [0.66091919, -0.16799635, -0.91037834]],
            [[-4.1862600, 0.0575700, -0.0381200],
             [-3.1689400, 0.0523700, 0.0200000],
             [-4.4978600, 0.8211300, 0.5604100],
             [-4.4978700, -0.8000100, 0.4155600],
             [0.00000000, -0.00000000, -0.00000000]]
        ], device=self.device)
        species = torch.tensor([[1, 0, 0, 0, 0], [2, 0, 0, 0, -1]], device=self.device)

        _, aev = self.aev_computer((species, coordinates))
        _, cu_aev = self.cuaev_computer((species, coordinates))
        self.assertEqual(cu_aev, aev)

if __name__ == '__main__':
    unittest.main()
```

Fix patch:
diff --git a/torchani/cuaev/aev.cu b/torchani/cuaev/aev.cu
index 4edb511e1..3f02e9a01 100644
--- a/torchani/cuaev/aev.cu
+++ b/torchani/cuaev/aev.cu
@@ -1,12 +1,15 @@
 #include <aev.h>
 #include <torch/extension.h>
 #include <cuaev_cub.cuh>
-#include <vector>
 
 #include <ATen/Context.h>
 #include <THC/THC.h>
 #include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
 #include <THC/THCThrustAllocator.cuh>
+#include <vector>
 
 #define PI 3.141592653589793
 using torch::Tensor;
@@ -742,9 +745,13 @@ Result cuaev_forward(const Tensor& coordinates_t, const Tensor& species_t, const
   TORCH_CHECK(
       (species_t.dtype() == torch::kInt32) && (coordinates_t.dtype() == torch::kFloat32), "Unsupported input type");
   TORCH_CHECK(
-      aev_params.EtaR_t.size(0) == 1 || aev_params.EtaA_t.size(0) == 1 || aev_params.Zeta_t.size(0) == 1,
+      aev_params.EtaR_t.size(0) == 1 && aev_params.EtaA_t.size(0) == 1 && aev_params.Zeta_t.size(0) == 1,
       "cuda extension is currently not supported for the specified "
       "configuration");
+  TORCH_CHECK(
+      coordinates_t.device() == species_t.device() && coordinates_t.device() == aev_params.EtaR_t.device() &&
+          coordinates_t.device() == aev_params.EtaA_t.device(),
+      "coordinates, species, and aev_params should be on the same device");
 
   float Rcr = aev_params.Rcr;
   float Rca = aev_params.Rca;
@@ -759,7 +766,8 @@ Result cuaev_forward(const Tensor& coordinates_t, const Tensor& species_t, const
         aev_t, Tensor(), Tensor(), Tensor(), 0, 0, 0, Tensor(), Tensor(), Tensor(), 0, 0, 0, coordinates_t, species_t};
   }
 
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  at::cuda::CUDAGuard device_guard(coordinates_t.device().index());
+  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
   auto& allocator = *c10::cuda::CUDACachingAllocator::get();
 
   // buffer to store all the pairwise distance (Rij)
@@ -790,6 +798,7 @@ Result cuaev_forward(const Tensor& coordinates_t, const Tensor& species_t, const
         coordinates_t.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
         d_Rij,
         max_natoms_per_mol);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
   } else {
     dim3 block(8, 8, 1);
     // Compute pairwise distance (Rij) for all atom pairs in a molecule
@@ -800,6 +809,7 @@ Result cuaev_forward(const Tensor& coordinates_t, const Tensor& species_t, const
         coordinates_t.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
         d_Rij,
         max_natoms_per_mol);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
   }
 
   // Extract Rijs that is needed for RadialAEV comptuation i.e. all the Rij <= Rcr
@@ -822,6 +832,7 @@ Result cuaev_forward(const Tensor& coordinates_t, const Tensor& species_t, const
       aev_params.radial_length,
       aev_params.radial_sublength,
       nRadialRij);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   // reuse buffer allocated for all Rij
   // d_angularRij will store all the Rij required in Angular AEV computation
@@ -890,6 +901,7 @@ Result cuaev_forward(const Tensor& coordinates_t, const Tensor& species_t, const
         maxnbrs_per_atom_aligned,
         angular_length_aligned,
         ncenter_atoms);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
 
     return {
         aev_t,
@@ -917,7 +929,8 @@ Tensor cuaev_backward(const Tensor& grad_output, const AEVScalarParams& aev_para
 
   const int n_molecules = coordinates_t.size(0);
   const int max_natoms_per_mol = coordinates_t.size(1);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  at::cuda::CUDAGuard device_guard(coordinates_t.device().index());
+  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
 
   auto grad_coord = torch::zeros(coordinates_t.sizes(), coordinates_t.options().requires_grad(false)); // [2, 5, 3]
 
@@ -943,6 +956,7 @@ Tensor cuaev_backward(const Tensor& grad_output, const AEVScalarParams& aev_para
       aev_params.radial_length,
       aev_params.radial_sublength,
       result.nRadialRij);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   // For best result, block_size should match average molecule size (no padding) to avoid atomicAdd
   nblocks = (result.nRadialRij + block_size - 1) / block_size;
@@ -952,6 +966,7 @@ Tensor cuaev_backward(const Tensor& grad_output, const AEVScalarParams& aev_para
       grad_coord.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
       d_radialRij,
       result.nRadialRij);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   auto smem_size = [&aev_params](int max_nbrs, int ncatom_per_tpb) {
     int sxyz = sizeof(float) * max_nbrs * 3;
@@ -991,6 +1006,7 @@ Tensor cuaev_backward(const Tensor& grad_output, const AEVScalarParams& aev_para
       result.maxnbrs_per_atom_aligned,
       result.angular_length_aligned,
       result.ncenter_atoms);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   return grad_coord;
 }
@@ -1002,7 +1018,8 @@ Tensor cuaev_double_backward(const Tensor& grad_force, const AEVScalarParams& ae
 
   const int n_molecules = coordinates_t.size(0);
   const int max_natoms_per_mol = coordinates_t.size(1);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  at::cuda::CUDAGuard device_guard(coordinates_t.device().index());
+  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
 
   int aev_length = aev_params.radial_length + aev_params.angular_length;
 
@@ -1027,6 +1044,7 @@ Tensor cuaev_double_backward(const Tensor& grad_force, const AEVScalarParams& ae
       grad_force.packed_accessor32<float, 3, torch::RestrictPtrTraits>(),
       d_radialRij,
       result.nRadialRij);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   nblocks = (result.nRadialRij * 8 + block_size - 1) / block_size;
   cuRadialAEVs_backward_or_doublebackward<true, int, float, 8><<<nblocks, block_size, 0, stream>>>(
@@ -1040,6 +1058,7 @@ Tensor cuaev_double_backward(const Tensor& grad_force, const AEVScalarParams& ae
       aev_params.radial_length,
       aev_params.radial_sublength,
       result.nRadialRij);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   auto smem_size = [&aev_params](int max_nbrs, int ncatom_per_tpb) {
     int sxyz = sizeof(float) * max_nbrs * 3;
@@ -1078,6 +1097,7 @@ Tensor cuaev_double_backward(const Tensor& grad_force, const AEVScalarParams& ae
       result.maxnbrs_per_atom_aligned,
       result.angular_length_aligned,
       result.ncenter_atoms);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
 
   return grad_grad_aev;
 }


Test patch:
diff --git a/tests/test_cuaev.py b/tests/test_cuaev.py
index 3426823fd..3f96efce2 100644
--- a/tests/test_cuaev.py
+++ b/tests/test_cuaev.py
@@ -8,8 +8,8 @@
 
 path = os.path.dirname(os.path.realpath(__file__))
 
-skipIfNoGPU = unittest.skipIf(not torch.cuda.is_available(),
-                              'There is no device to run this test')
+skipIfNoGPU = unittest.skipIf(not torch.cuda.is_available(), 'There is no device to run this test')
+skipIfNoMultiGPU = unittest.skipIf(not torch.cuda.device_count() >= 2, 'There is not enough GPU devices to run this test')
 skipIfNoCUAEV = unittest.skipIf(not torchani.aev.has_cuaev, "only valid when cuaev is installed")
 
 
@@ -39,9 +39,9 @@ def testAEVComputer(self):
 @skipIfNoCUAEV
 class TestCUAEV(TestCase):
 
-    def setUp(self):
+    def setUp(self, device='cuda:0'):
         self.tolerance = 5e-5
-        self.device = 'cuda'
+        self.device = device
         Rcr = 5.2000e+00
         Rca = 3.5000e+00
         EtaR = torch.tensor([1.6000000e+01], device=self.device)
@@ -131,6 +131,15 @@ def testSimple(self):
         _, cu_aev = self.cuaev_computer((species, coordinates))
         self.assertEqual(cu_aev, aev)
 
+    @skipIfNoMultiGPU
+    def testMultiGPU(self):
+        self.setUp(device='cuda:1')
+        self.testSimple()
+        self.testSimpleBackward()
+        self.testSimpleDoubleBackward_1()
+        self.testSimpleDoubleBackward_2()
+        self.setUp(device='cuda:0')
+
     def testSimpleBackward(self):
         coordinates = torch.tensor([
             [[0.03192167, 0.00638559, 0.01301679],