Skip to content

Commit fda9cea

Browse files
FlamefireFlamefire
committed
Don't detect local GPU if $DS_IGNORE_CUDA_DETECTION is set
Make this consistent over all OPs: For cross-compilation we should not check the local GPU version.
1 parent b6346bf commit fda9cea

File tree

7 files changed

+85
-75
lines changed

7 files changed

+85
-75
lines changed

‎op_builder/fp_quantizer.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
# DeepSpeed Team
55

6+
import os
67
try:
78
from packaging import version as pkg_version
89
except ImportError:
@@ -31,19 +32,20 @@ def is_compatible(self, verbose=False):
3132
return False
3233

3334
cuda_okay = True
34-
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
35-
sys_cuda_major, _ = installed_cuda_version()
36-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
37-
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
38-
if cuda_capability < 8:
39-
if verbose:
40-
self.warning("NVIDIA Inference is only supported on Ampere and newer architectures")
41-
cuda_okay = False
42-
if cuda_capability >= 8:
43-
if torch_cuda_major < 11 or sys_cuda_major < 11:
35+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
36+
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
37+
sys_cuda_major, _ = installed_cuda_version()
38+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
39+
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
40+
if cuda_capability < 8:
4441
if verbose:
45-
self.warning("On Ampere and higher architectures please use CUDA 11+")
42+
self.warning("NVIDIA Inference is only supported on Ampere and newer architectures")
4643
cuda_okay = False
44+
if cuda_capability >= 8:
45+
if torch_cuda_major < 11 or sys_cuda_major < 11:
46+
if verbose:
47+
self.warning("On Ampere and higher architectures please use CUDA 11+")
48+
cuda_okay = False
4749

4850
try:
4951
import triton

‎op_builder/inference_core_ops.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,19 +28,20 @@ def is_compatible(self, verbose=False):
2828
return False
2929

3030
cuda_okay = True
31-
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
32-
sys_cuda_major, _ = installed_cuda_version()
33-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
34-
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
35-
if cuda_capability < 6:
36-
if verbose:
37-
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
38-
cuda_okay = False
39-
if cuda_capability >= 8:
40-
if torch_cuda_major < 11 or sys_cuda_major < 11:
31+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
32+
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
33+
sys_cuda_major, _ = installed_cuda_version()
34+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
35+
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
36+
if cuda_capability < 6:
4137
if verbose:
42-
self.warning("On Ampere and higher architectures please use CUDA 11+")
38+
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
4339
cuda_okay = False
40+
if cuda_capability >= 8:
41+
if torch_cuda_major < 11 or sys_cuda_major < 11:
42+
if verbose:
43+
self.warning("On Ampere and higher architectures please use CUDA 11+")
44+
cuda_okay = False
4445
return super().is_compatible(verbose) and cuda_okay
4546

4647
def filter_ccs(self, ccs):

‎op_builder/inference_cutlass_builder.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,19 +27,20 @@ def is_compatible(self, verbose=False):
2727
return False
2828

2929
cuda_okay = True
30-
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
31-
sys_cuda_major, _ = installed_cuda_version()
32-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
33-
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
34-
if cuda_capability < 6:
35-
if verbose:
36-
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
37-
cuda_okay = False
38-
if cuda_capability >= 8:
39-
if torch_cuda_major < 11 or sys_cuda_major < 11:
30+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
31+
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
32+
sys_cuda_major, _ = installed_cuda_version()
33+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
34+
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
35+
if cuda_capability < 6:
4036
if verbose:
41-
self.warning("On Ampere and higher architectures please use CUDA 11+")
37+
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
4238
cuda_okay = False
39+
if cuda_capability >= 8:
40+
if torch_cuda_major < 11 or sys_cuda_major < 11:
41+
if verbose:
42+
self.warning("On Ampere and higher architectures please use CUDA 11+")
43+
cuda_okay = False
4344
return super().is_compatible(verbose) and cuda_okay
4445

4546
def filter_ccs(self, ccs):

‎op_builder/ragged_ops.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,19 +28,20 @@ def is_compatible(self, verbose=False):
2828
return False
2929

3030
cuda_okay = True
31-
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
32-
sys_cuda_major, _ = installed_cuda_version()
33-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
34-
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
35-
if cuda_capability < 6:
36-
if verbose:
37-
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
38-
cuda_okay = False
39-
if cuda_capability >= 8:
40-
if torch_cuda_major < 11 or sys_cuda_major < 11:
31+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
32+
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
33+
sys_cuda_major, _ = installed_cuda_version()
34+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
35+
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
36+
if cuda_capability < 6:
4137
if verbose:
42-
self.warning("On Ampere and higher architectures please use CUDA 11+")
38+
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
4339
cuda_okay = False
40+
if cuda_capability >= 8:
41+
if torch_cuda_major < 11 or sys_cuda_major < 11:
42+
if verbose:
43+
self.warning("On Ampere and higher architectures please use CUDA 11+")
44+
cuda_okay = False
4445
return super().is_compatible(verbose) and cuda_okay
4546

4647
def filter_ccs(self, ccs):

‎op_builder/ragged_utils.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,19 +28,20 @@ def is_compatible(self, verbose=False):
2828
return False
2929

3030
cuda_okay = True
31-
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
32-
sys_cuda_major, _ = installed_cuda_version()
33-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
34-
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
35-
if cuda_capability < 6:
36-
if verbose:
37-
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
38-
cuda_okay = False
39-
if cuda_capability >= 8:
40-
if torch_cuda_major < 11 or sys_cuda_major < 11:
31+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
32+
if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda
33+
sys_cuda_major, _ = installed_cuda_version()
34+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
35+
cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda
36+
if cuda_capability < 6:
4137
if verbose:
42-
self.warning("On Ampere and higher architectures please use CUDA 11+")
38+
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
4339
cuda_okay = False
40+
if cuda_capability >= 8:
41+
if torch_cuda_major < 11 or sys_cuda_major < 11:
42+
if verbose:
43+
self.warning("On Ampere and higher architectures please use CUDA 11+")
44+
cuda_okay = False
4445
return super().is_compatible(verbose) and cuda_okay
4546

4647
def filter_ccs(self, ccs):

‎op_builder/spatial_inference.py‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
# DeepSpeed Team
55

6+
import os
67
from .builder import CUDAOpBuilder, installed_cuda_version
78

89

@@ -26,15 +27,16 @@ def is_compatible(self, verbose=False):
2627
return False
2728

2829
cuda_okay = True
29-
if not self.is_rocm_pytorch() and torch.cuda.is_available():
30-
sys_cuda_major, _ = installed_cuda_version()
31-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
32-
cuda_capability = torch.cuda.get_device_properties(0).major
33-
if cuda_capability >= 8:
34-
if torch_cuda_major < 11 or sys_cuda_major < 11:
35-
if verbose:
36-
self.warning("On Ampere and higher architectures please use CUDA 11+")
37-
cuda_okay = False
30+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
31+
if not self.is_rocm_pytorch() and torch.cuda.is_available():
32+
sys_cuda_major, _ = installed_cuda_version()
33+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
34+
cuda_capability = torch.cuda.get_device_properties(0).major
35+
if cuda_capability >= 8:
36+
if torch_cuda_major < 11 or sys_cuda_major < 11:
37+
if verbose:
38+
self.warning("On Ampere and higher architectures please use CUDA 11+")
39+
cuda_okay = False
3840
return super().is_compatible(verbose) and cuda_okay
3941

4042
def sources(self):

‎op_builder/transformer_inference.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
# DeepSpeed Team
55

6+
import os
67
from .builder import CUDAOpBuilder, installed_cuda_version
78

89

@@ -26,19 +27,20 @@ def is_compatible(self, verbose=False):
2627
return False
2728

2829
cuda_okay = True
29-
if not self.is_rocm_pytorch() and torch.cuda.is_available():
30-
sys_cuda_major, _ = installed_cuda_version()
31-
torch_cuda_major = int(torch.version.cuda.split('.')[0])
32-
cuda_capability = torch.cuda.get_device_properties(0).major
33-
if cuda_capability < 6:
34-
if verbose:
35-
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
36-
cuda_okay = False
37-
if cuda_capability >= 8:
38-
if torch_cuda_major < 11 or sys_cuda_major < 11:
30+
if not os.environ.get("DS_IGNORE_CUDA_DETECTION"):
31+
if not self.is_rocm_pytorch() and torch.cuda.is_available():
32+
sys_cuda_major, _ = installed_cuda_version()
33+
torch_cuda_major = int(torch.version.cuda.split('.')[0])
34+
cuda_capability = torch.cuda.get_device_properties(0).major
35+
if cuda_capability < 6:
3936
if verbose:
40-
self.warning("On Ampere and higher architectures please use CUDA 11+")
37+
self.warning("NVIDIA Inference is only supported on Pascal and newer architectures")
4138
cuda_okay = False
39+
if cuda_capability >= 8:
40+
if torch_cuda_major < 11 or sys_cuda_major < 11:
41+
if verbose:
42+
self.warning("On Ampere and higher architectures please use CUDA 11+")
43+
cuda_okay = False
4244
return super().is_compatible(verbose) and cuda_okay
4345

4446
def filter_ccs(self, ccs):

0 commit comments

Comments
 (0)