Raise if the nvcc in cuda_home is too old for the current GPU.
(cuda_home: str)
| 32 | |
| 33 | |
| 34 | def _check_nvcc_supports_gpu(cuda_home: str) -> None: |
| 35 | """Raise if the nvcc in cuda_home is too old for the current GPU.""" |
| 36 | nvcc = os.path.join(cuda_home, "bin", "nvcc") |
| 37 | if not os.path.isfile(nvcc): |
| 38 | return |
| 39 | try: |
| 40 | import torch |
| 41 | |
| 42 | major, minor = torch.cuda.get_device_capability() |
| 43 | sm = f"compute_{major}{minor}" |
| 44 | result = subprocess.run( |
| 45 | [nvcc, f"-arch={sm}", "-x", "cu", "/dev/null", "-o", "/dev/null"], |
| 46 | capture_output=True, |
| 47 | text=True, |
| 48 | ) |
| 49 | if "Unsupported gpu architecture" in result.stderr: |
| 50 | raise RuntimeError( |
| 51 | f"nvcc at {nvcc} does not support this GPU (sm_{major}{minor}). " |
| 52 | f"Please install a newer CUDA toolkit (>= 12.0 for H100) and set " |
| 53 | f"CUDA_HOME to its path, e.g.: export CUDA_HOME=/usr/local/cuda-12.2" |
| 54 | ) |
| 55 | except ImportError: |
| 56 | pass |
| 57 | |
| 58 | |
| 59 | cuda_home = _find_cuda_home() |