# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai # SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium import contextlib import importlib import os import time from contextlib import contextmanager from enum import Enum from functools import lru_cache from typing import Callable, List, Optional, Union import torch from packaging import version from torch.cpu import StreamContext from ..utils.logger import setup_logger from ..utils.safe import GC from . import gte_python_3_13_3, gte_python_3_14, has_gil_disabled, log_gil_requirements_for # pytorch 2.6.0 fixes many compilation errors TORCH_HAS_COMPILE = version.parse(torch.__version__).release >= version.Version('2.6').release TORCH_GTE_28 = version.parse(torch.__version__).release >= version.Version('2.8').release TORCH_GTE_210 = version.parse(torch.__version__).release >= version.Version('2.10').release TORCH_HAS_FUSED_OPS = version.parse(torch.__version__).release >= version.Version('2.8').release HAS_CUDA = False HAS_XPU = False HAS_MPS = False HAS_MLX = False HAS_NPU = False CPU = torch.device("cpu") META = torch.device("meta") class BalanceStrategy(str, Enum): MEMORY = "memory", # make vram more spread out GPU = "gpu" # vram is less balanced (gpu0) but gpu0 is also used for quantization DEFAULT_BALANCE_STRATEGY = BalanceStrategy.GPU # TODO FIX ME...this should be removed STREAM = None # cache log = setup_logger() def timed_gc_collect() -> int: """Run ``gc.collect`` and log the elapsed time along with reclaimed object count.""" start = time.perf_counter() # Python 3.14 removed gen1 so there is only gen0 and gen2 collected = GC.collect() duration = time.perf_counter() - start log.info(f"gc.collect() reclaimed {collected} objects in {duration:.3f}s") return collected # reset dynamo cache on each model load since during ci loop model inference may exhuast cache try: torch._dynamo.reset() # Increase the dynamo cache size limit, default of 8 is too low if torch._dynamo.config.cache_size_limit < 128: torch._dynamo.config.cache_size_limit = 128 except BaseException: # triton built from source maybe incompatible with _dynamo private api pass try: HAS_CUDA = torch.cuda.is_available() except Exception: HAS_CUDA = False try: HAS_XPU = torch.xpu.is_available() except Exception: HAS_XPU = False try: HAS_MPS = torch.mps.is_available() except Exception: HAS_MPS = False def _ascend_runtime_env_ready() -> bool: # torch_npu may report available before the CANN environment is fully # sourced. Requiring the standard Ascend paths avoids import-time lazy # initialization failures when the Python package is installed but the # runtime is not usable in this shell. return any( os.environ.get(name) for name in ("ASCEND_HOME_PATH", "ASCEND_TOOLKIT_HOME", "ASCEND_OPP_PATH") ) try: importlib.import_module("torch_npu") HAS_NPU = _ascend_runtime_env_ready() and torch.npu.is_available() except Exception: HAS_NPU = False # mlx check try: import mlx.core # type: ignore # noqa: F401 import mlx.core.metal # noqa: F401 HAS_MLX = True except BaseException: pass try: torch.backends.fp32_precision BACKENDS_HAS_FP32_PRECISION = True except AttributeError: BACKENDS_HAS_FP32_PRECISION = False def _set_tf32_state(enabled: bool) -> None: if BACKENDS_HAS_FP32_PRECISION: mode = "tf32" if enabled else "ieee" torch.backends.fp32_precision = mode torch.backends.cuda.matmul.fp32_precision = mode torch.backends.cudnn.fp32_precision = mode torch.backends.cudnn.conv.fp32_precision = mode torch.backends.cudnn.rnn.fp32_precision = mode return torch.backends.cuda.matmul.allow_tf32 = enabled torch.backends.cudnn.allow_tf32 = enabled def _snapshot_tf32_state(): if BACKENDS_HAS_FP32_PRECISION: return ( torch.backends.fp32_precision, torch.backends.cuda.matmul.fp32_precision, torch.backends.cudnn.fp32_precision, torch.backends.cudnn.conv.fp32_precision, torch.backends.cudnn.rnn.fp32_precision, ) return ( torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32, ) def _restore_tf32_state(state) -> None: if BACKENDS_HAS_FP32_PRECISION: torch.backends.fp32_precision = state[0] torch.backends.cuda.matmul.fp32_precision = state[1] torch.backends.cudnn.fp32_precision = state[2] torch.backends.cudnn.conv.fp32_precision = state[3] torch.backends.cudnn.rnn.fp32_precision = state[4] return torch.backends.cuda.matmul.allow_tf32 = state[0] torch.backends.cudnn.allow_tf32 = state[1] def torch_compile(module: Union[torch.nn.Module, Callable], backend:str ="inductor", mode: str = None, fullgraph=False): # requires torch >2.8 for proper torch.compile + Python 3.13.3t (freethreading) if has_gil_disabled() and not gte_python_3_13_3(): log_gil_requirements_for("Torch Compile") return module if gte_python_3_14() and not TORCH_GTE_210: log_gil_requirements_for("Torch Compile") return module if not TORCH_HAS_COMPILE: return module if HAS_MPS and not TORCH_GTE_28: if not torch._dynamo.config.suppress_errors: log.warn("To use compile() with MPS, you need to have torch version >= 2.8.0, " "please upgrade it by `pip install -U torch torchaudio torchvision`") torch._dynamo.config.suppress_errors = True return module try: return torch.compile(module, backend=backend, mode=mode, fullgraph=fullgraph) except BaseException as e: log.warn.once(f"Failed to compile `{module}`, {e}") return module def torch_new_stream(): global STREAM if STREAM is not None: return STREAM if HAS_CUDA: STREAM = torch.cuda.Stream() return STREAM if HAS_XPU: STREAM = torch.xpu.Stream() return STREAM if HAS_NPU: STREAM = torch.npu.Stream() return STREAM return None def torch_new_stream_ctx(): if HAS_CUDA: return torch.cuda.stream(torch_new_stream()) if HAS_XPU: return torch.xpu.stream(torch_new_stream()) if HAS_NPU: return torch.npu.stream(torch_new_stream()) return contextlib.nullcontext() def torch_sync(device: torch.device = None): """Synchronize accelerator queues. When no device is provided we synchronize every detected accelerator index so replication work staged on multiple GPUs/NPUs completes before issuing more kernels. """ if device is None: synchronized_any = False if HAS_CUDA: dev_count = torch.cuda.device_count() if dev_count: synchronized_any = True for idx in range(dev_count): torch.cuda.synchronize(idx) if HAS_XPU: dev_count = torch.xpu.device_count() if dev_count: synchronized_any = True for idx in range(dev_count): torch.xpu.synchronize(idx) if HAS_MPS: synchronized_any = True torch.mps.synchronize() if HAS_NPU: dev_count = torch.npu.device_count() if dev_count: synchronized_any = True for idx in range(dev_count): torch.npu.synchronize(idx) if not synchronized_any: torch.cpu.synchronize() return if device.type == "cuda": torch.cuda.synchronize(device=device) elif device.type == "xpu": torch.xpu.synchronize(device=device) elif device.type == "mps": torch.mps.synchronize() elif device.type == "npu": torch.npu.synchronize(device=device) elif device.type == "cpu": torch.cpu.synchronize() def _normalize_device(device: Union[torch.device, str, int, None]) -> Optional[torch.device]: if device is None: return None if isinstance(device, torch.device): return device if isinstance(device, str): return torch.device(device) if isinstance(device, int): if HAS_CUDA and torch.cuda.device_count() > device: return torch.device("cuda", device) if HAS_XPU and torch.xpu.device_count() > device: return torch.device("xpu", device) if HAS_NPU and torch.npu.device_count() > device: return torch.device("npu", device) raise ValueError(f"Unable to map integer `{device}` to a known accelerator device.") raise TypeError(f"Unsupported device specifier type: {type(device)}") @lru_cache(maxsize=12) def resolve_empty_cache_callable(device_type: str) -> Optional[Callable[[], None]]: try: if device_type == "cuda" and HAS_CUDA: candidate = torch.cuda.empty_cache elif device_type == "xpu" and HAS_XPU: candidate = torch.xpu.empty_cache elif device_type == "mps" and HAS_MPS: candidate = torch.mps.empty_cache elif device_type == "npu" and HAS_NPU: candidate = torch.npu.empty_cache else: candidate = None except Exception: candidate = None if callable(candidate): return candidate return None def empty_cache_for_device(device: torch.device) -> bool: if device.type == "cuda": fn = resolve_empty_cache_callable("cuda") if fn is None: return False target = device if device.index is not None else torch.device("cuda") with torch.cuda.device(target): fn() return True if device.type == "xpu": fn = resolve_empty_cache_callable("xpu") if fn is None: return False target = device if device.index is not None else torch.device("xpu") with torch.xpu.device(target): fn() return True if device.type == "mps": fn = resolve_empty_cache_callable("mps") if fn is None: return False fn() if HAS_MLX: mlx.core.clear_cache() return True if device.type == "npu": fn = resolve_empty_cache_callable("npu") if fn is None: return False target = device if device.index is not None else torch.device("npu") with torch.npu.device(target): fn() return True if device.type == "cpu" and HAS_MLX: mlx.core.clear_cache() return True return False def torch_empty_cache_any( device: Union[torch.device, str, int, None] = None, gc: bool = True, sync: bool = False, ) -> bool: normalized = _normalize_device(device) if gc: timed_gc_collect() if sync: torch_sync(device=normalized) success = False if normalized is None: if HAS_CUDA: fn = resolve_empty_cache_callable("cuda") if fn is not None: fn() success = True if HAS_XPU: fn = resolve_empty_cache_callable("xpu") if fn is not None: fn() success = True if HAS_MPS: fn = resolve_empty_cache_callable("mps") if fn is not None: fn() success = True if HAS_NPU: fn = resolve_empty_cache_callable("npu") if fn is not None: fn() success = True if HAS_MLX: mlx.core.clear_cache() success = True return success return empty_cache_for_device(normalized) def torch_empty_cache(device: torch.device = None, gc: bool = True, sync: bool = False) -> bool: return torch_empty_cache_any(device=device, gc=gc, sync=sync) def auto_select_torch_device(index: int = 0): assert index >= 0, f"device index should be a positive integer: actual = `{index}`" if HAS_CUDA: # defensive check if index > 0 and torch.cuda.device_count() <= index : index = 0 device = torch.device(f"cuda:{index}") elif HAS_XPU: # defensive check if index > 0 and torch.xpu.device_count() <= index: index = 0 device = torch.device(f"xpu:{index}") elif HAS_NPU: # defensive check if index > 0 and torch.npu.device_count() <= index: index = 0 device = torch.device(f"npu:{index}") elif HAS_MPS: device = torch.device("mps") # mps has no index else: device = CPU # cpu has no index return device # some device types can have multiple gpus cuda/rocm + xpu def torch_devices() -> List[torch.device]: if HAS_CUDA: return [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())] elif HAS_XPU: return [torch.device(f"xpu:{i}") for i in range(torch.xpu.device_count())] elif HAS_NPU: return [torch.device(f"npu:{i}") for i in range(torch.npu.device_count())] elif HAS_MPS: return [torch.device("mps")] else: return [CPU] def npu_devices_by_pci_bus_order() -> List[torch.device]: """Return visible NPU devices in torch runtime order. Ascend exposes process-level NPU visibility through ASCEND_RT_VISIBLE_DEVICES. torch-npu remaps that visible set to logical indices, so callers should set the env var before process start and then use the resulting torch order. """ if not HAS_NPU: return [] try: count = int(torch.npu.device_count()) except Exception: return [] if count <= 0: return [] devices: List[torch.device] = [] for logical_index in range(count): try: devices.append(torch.device("npu", logical_index)) except (RuntimeError, ValueError): return [] return devices def last_npu_device_by_pci_bus_order() -> Optional[torch.device]: """Return the last visible NPU in torch runtime order, or None when unavailable.""" devices = npu_devices_by_pci_bus_order() if not devices: return None return devices[-1] ALL_DEVICES = torch_devices() class _LazyAcceleratorStreams: """Create accelerator streams only when a caller actually needs them.""" def __init__(self, devices: List[torch.device]): self._devices = devices self._streams: List[Optional[object]] = [None] * len(devices) def __len__(self): return max(1, len(self._devices)) def __iter__(self): for index in range(len(self)): yield self[index] def __getitem__(self, index: int): if not self._devices: if index == 0: return contextlib.nullcontext() raise IndexError(index) stream = self._streams[index] if stream is not None: return stream device = self._devices[index] if device.type == "cuda": stream = torch.cuda.Stream(device=device) elif device.type == "xpu": stream = torch.xpu.Stream(device=device) elif device.type == "npu": stream = torch.npu.Stream(device=device) else: stream = contextlib.nullcontext() self._streams[index] = stream return stream class _LazyAcceleratorStreamRef: """Reference one lazy stream without materializing it at import time.""" def __init__(self, streams: _LazyAcceleratorStreams, index: int): self._streams = streams self._index = index def get(self): return self._streams[self._index] ALL_STREAMS = _LazyAcceleratorStreams(ALL_DEVICES) DEVICE_0 = auto_select_torch_device(index=0) # device_1 may be same as device_0 if there is only 1 visible/active device DEVICE_1 = auto_select_torch_device(index=1) DEVICE_0_STREAM = _LazyAcceleratorStreamRef(ALL_STREAMS, 0) NEXT_DEVICE_INDEX = 0 # def device_next_reset(): # global NEXT_DEVICE_INDEX # NEXT_DEVICE_INDEX = 0 # # def device_next(balance_strategy: BalanceStrategy = DEFAULT_BALANCE_STRATEGY) -> torch.device: # global NEXT_DEVICE_INDEX # # if len(ALL_DEVICES) <= 1: # return ALL_DEVICES[0] # # device = ALL_DEVICES[NEXT_DEVICE_INDEX] # if NEXT_DEVICE_INDEX < len(ALL_DEVICES) - 1: # NEXT_DEVICE_INDEX += 1 # else: # if balance_strategy == BalanceStrategy.MEMORY: # NEXT_DEVICE_INDEX = 1 # else: # NEXT_DEVICE_INDEX = 0 # # return device def torch_streamCtx(stream) -> StreamContext: if isinstance(stream, _LazyAcceleratorStreamRef): stream = stream.get() if HAS_CUDA: return torch.cuda.stream(stream) if HAS_XPU: return torch.xpu.stream(stream) if HAS_NPU: return torch.npu.stream(stream) return contextlib.nullcontext() @contextmanager def tf32_high_precision_guard(): if not HAS_CUDA: yield return previous_state = _snapshot_tf32_state() _set_tf32_state(False) try: yield finally: _restore_tf32_state(previous_state) @contextmanager def tf32_disable_guard(): if not HAS_CUDA: yield return previous_state = _snapshot_tf32_state() _set_tf32_state(False) try: yield finally: _restore_tf32_state(previous_state) @contextmanager def tf32_enable_guard(): if not HAS_CUDA: yield return previous_state = _snapshot_tf32_state() _set_tf32_state(True) try: yield finally: _restore_tf32_state(previous_state)