Spaces:
Paused
Paused
File size: 2,278 Bytes
e319a7b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Device-agnostic accelerator helpers.
``torch.accelerator`` (``is_available``/``current_accelerator``/``synchronize``)
was only added in PyTorch 2.6; on older versions we fall back to checking CUDA
then MPS directly, in that order.
"""
import platform
from packaging.version import Version
import torch
_HAS_TORCH_ACCELERATOR = Version(torch.__version__.split("+")[0]) >= Version("2.6")
def _mps_available() -> bool:
"""Whether MPS is available and worth auto-selecting.
torch <= 2.2 also reports MPS as available on Intel Macs with AMD GPUs, a
backend that was never solid and has since been abandoned. Passing
``device="mps"`` explicitly still works there for those who want to try
it; this only affects auto-detection.
"""
return torch.backends.mps.is_available() and platform.machine() == "arm64"
def is_available() -> bool:
"""Whether an accelerator (GPU) is available."""
if _HAS_TORCH_ACCELERATOR:
return torch.accelerator.is_available()
return torch.cuda.is_available() or _mps_available()
def current_accelerator() -> torch.device:
"""The current accelerator device.
Raises ``RuntimeError`` if no accelerator is available; check
:func:`is_available` first.
"""
if _HAS_TORCH_ACCELERATOR:
return torch.accelerator.current_accelerator()
if torch.cuda.is_available():
return torch.device("cuda")
if _mps_available():
return torch.device("mps")
raise RuntimeError("No available accelerator detected.")
def synchronize() -> None:
"""Wait for all kernels on the current accelerator to complete.
No-op if no accelerator is available.
"""
if _HAS_TORCH_ACCELERATOR:
# torch.accelerator.synchronize() still tries to init CUDA even on CPU-only systems
# Only call it if we actually have a non-CPU accelerator
try:
current_device = torch.accelerator.current_device_index()
if current_device >= 0:
torch.accelerator.synchronize()
except RuntimeError:
# No accelerator available (CUDA not found, etc.)
pass
return
if torch.cuda.is_available():
torch.cuda.synchronize()
elif _mps_available():
torch.mps.synchronize()
|