File size: 1,384 Bytes
e65937c | 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 | import torch
from transformers.utils import is_torch_npu_available
def is_npu_available() -> bool:
"""Detect Ascend NPU availability"""
try:
return is_torch_npu_available()
except (ImportError, RuntimeError, ModuleNotFoundError, AttributeError):
return False
def get_current_device() -> str:
"""Get the current accelerator device string (e.g. 'cuda:0', 'npu:0')."""
acc = torch.accelerator.current_accelerator()
if acc is None:
return "cuda:0"
device_idx = torch.accelerator.current_device_index()
return f"{acc.type}:{device_idx}"
def get_device_name(idx: int) -> str:
acc = torch.accelerator.current_accelerator()
if acc is None:
return "NO ACCELERATOR"
mod = torch.get_device_module(acc)
if hasattr(mod, "get_device_name"):
return mod.get_device_name(idx)
else:
return f"{str(acc).upper()} DEVICE"
def mem_get_info():
acc = torch.accelerator.current_accelerator()
if acc is None:
return (0, 0)
mod = torch.get_device_module(acc)
if hasattr(mod, "mem_get_info"):
return mod.mem_get_info()
else:
return (0, 0)
def empty_cache():
acc = torch.accelerator.current_accelerator()
if acc is None:
return
mod = torch.get_device_module(acc)
if hasattr(mod, "empty_cache"):
mod.empty_cache()
return
|