MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / get_default_compute_capabilities

Function get_default_compute_capabilities

op_builder/builder.py:62–78  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

60
61
62def get_default_compute_capabilities():
63 compute_caps = DEFAULT_COMPUTE_CAPABILITIES
64 # Update compute capability according to: https://en.wikipedia.org/wiki/CUDA#GPUs_supported
65 import torch.utils.cpp_extension
66 if torch.utils.cpp_extension.CUDA_HOME is not None:
67 if installed_cuda_version()[0] == 11:
68 if installed_cuda_version()[1] >= 0:
69 compute_caps += ";8.0"
70 if installed_cuda_version()[1] >= 1:
71 compute_caps += ";8.6"
72 if installed_cuda_version()[1] >= 8:
73 compute_caps += ";9.0"
74 elif installed_cuda_version()[0] == 12:
75 compute_caps += ";8.0;8.6;9.0"
76 if installed_cuda_version()[1] >= 8:
77 compute_caps += ";10.0;12.0"
78 return compute_caps
79
80
81# list compatible minor CUDA versions - so that for example pytorch built with cuda-11.0 can be used

Callers 2

setup.pyFile · 0.90

Calls 1

installed_cuda_versionFunction · 0.85

Tested by

no test coverage detected