Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,307 @@
|
||||
r"""
|
||||
This package introduces support for the current :ref:`accelerator<accelerators>` in python.
|
||||
"""
|
||||
|
||||
from functools import cache
|
||||
from typing import Any
|
||||
from typing_extensions import deprecated
|
||||
|
||||
import torch
|
||||
|
||||
from ._utils import _device_t, _get_device_index
|
||||
from .graphs import Graph
|
||||
from .memory import (
|
||||
empty_cache,
|
||||
empty_host_cache,
|
||||
get_memory_info,
|
||||
max_memory_allocated,
|
||||
max_memory_reserved,
|
||||
memory_allocated,
|
||||
memory_reserved,
|
||||
memory_stats,
|
||||
reset_accumulated_memory_stats,
|
||||
reset_peak_memory_stats,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Graph",
|
||||
"current_accelerator",
|
||||
"current_device_idx", # deprecated
|
||||
"current_device_index",
|
||||
"get_device_capability",
|
||||
"current_stream",
|
||||
"device_count",
|
||||
"device_index",
|
||||
"empty_cache",
|
||||
"empty_host_cache",
|
||||
"get_memory_info",
|
||||
"is_available",
|
||||
"max_memory_allocated",
|
||||
"max_memory_reserved",
|
||||
"memory_allocated",
|
||||
"memory_reserved",
|
||||
"memory_stats",
|
||||
"reset_accumulated_memory_stats",
|
||||
"reset_peak_memory_stats",
|
||||
"set_device_idx", # deprecated
|
||||
"set_device_index",
|
||||
"set_stream",
|
||||
"synchronize",
|
||||
]
|
||||
|
||||
|
||||
def device_count() -> int:
|
||||
r"""Return the number of current :ref:`accelerator<accelerators>` available.
|
||||
|
||||
Returns:
|
||||
int: the number of the current :ref:`accelerator<accelerators>` available.
|
||||
If there is no available accelerators, return 0.
|
||||
|
||||
.. note:: This API delegates to the device-specific version of `device_count`.
|
||||
On CUDA, this API will NOT poison fork if NVML discovery succeeds.
|
||||
Otherwise, it will. For more details, see :ref:`multiprocessing-poison-fork-note`.
|
||||
"""
|
||||
acc = current_accelerator()
|
||||
if acc is None:
|
||||
return 0
|
||||
|
||||
mod = torch.get_device_module(acc)
|
||||
return mod.device_count()
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
r"""Check if the current accelerator is available at runtime: it was build, all the
|
||||
required drivers are available and at least one device is visible.
|
||||
See :ref:`accelerator<accelerators>` for details.
|
||||
|
||||
Returns:
|
||||
bool: A boolean indicating if there is an available :ref:`accelerator<accelerators>`.
|
||||
|
||||
.. note:: This API delegates to the device-specific version of `is_available`.
|
||||
On CUDA, when the environment variable ``PYTORCH_NVML_BASED_CUDA_CHECK=1`` is set,
|
||||
this function will NOT poison fork. Otherwise, it will. For more details, see
|
||||
:ref:`multiprocessing-poison-fork-note`.
|
||||
|
||||
Example::
|
||||
|
||||
>>> assert torch.accelerator.is_available() "No available accelerators detected."
|
||||
"""
|
||||
# Why not just check "device_count() > 0" like other is_available call?
|
||||
# Because device like CUDA have a python implementation of is_available that is
|
||||
# non-poisoning and some features like Dataloader rely on it.
|
||||
# So we are careful to delegate to the Python version of the accelerator here
|
||||
acc = current_accelerator()
|
||||
if acc is None:
|
||||
return False
|
||||
|
||||
mod = torch.get_device_module(acc)
|
||||
return mod.is_available()
|
||||
|
||||
|
||||
def current_accelerator(check_available: bool = False) -> torch.device | None:
|
||||
r"""Return the device of the accelerator available at compilation time.
|
||||
If no accelerator were available at compilation time, returns None.
|
||||
See :ref:`accelerator<accelerators>` for details.
|
||||
|
||||
Args:
|
||||
check_available (bool, optional): if True, will also do a runtime check to see
|
||||
if the device :func:`torch.accelerator.is_available` on top of the compile-time
|
||||
check.
|
||||
Default: ``False``
|
||||
|
||||
Returns:
|
||||
torch.device: return the current accelerator as :class:`torch.device`.
|
||||
|
||||
.. note:: The index of the returned :class:`torch.device` will be ``None``, please use
|
||||
:func:`torch.accelerator.current_device_index` to know the current index being used.
|
||||
This API does NOT poison fork. For more details, see :ref:`multiprocessing-poison-fork-note`.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest:
|
||||
>>> # If an accelerator is available, sent the model to it
|
||||
>>> model = torch.nn.Linear(2, 2)
|
||||
>>> if (current_device := current_accelerator(check_available=True)) is not None:
|
||||
>>> model.to(current_device)
|
||||
"""
|
||||
if (acc := torch._C._accelerator_getAccelerator()) is not None:
|
||||
if (not check_available) or (check_available and is_available()):
|
||||
return acc
|
||||
return None
|
||||
|
||||
|
||||
def current_device_index() -> int:
|
||||
r"""Return the index of a currently selected device for the current :ref:`accelerator<accelerators>`.
|
||||
|
||||
Returns:
|
||||
int: the index of a currently selected device.
|
||||
"""
|
||||
return torch._C._accelerator_getDeviceIndex()
|
||||
|
||||
|
||||
current_device_idx = deprecated(
|
||||
"Use `current_device_index` instead.",
|
||||
category=FutureWarning,
|
||||
)(current_device_index)
|
||||
|
||||
current_device_idx.__doc__ = r"""
|
||||
(Deprecated) Return the index of a currently selected device for the current :ref:`accelerator<accelerators>`.
|
||||
|
||||
Returns:
|
||||
int: the index of a currently selected device.
|
||||
|
||||
.. warning::
|
||||
|
||||
:func:`torch.accelerator.current_device_idx` is deprecated in favor of :func:`torch.accelerator.current_device_index`
|
||||
and will be removed in a future PyTorch release.
|
||||
"""
|
||||
|
||||
|
||||
@cache
|
||||
def get_device_capability(device: _device_t = None, /) -> dict[str, Any]:
|
||||
r"""Return the capability of the currently selected device.
|
||||
|
||||
Args:
|
||||
device (:class:`torch.device`, str, int, optional): The device to query capabilities for
|
||||
:ref:`accelerator<accelerators>` device type. If not given,
|
||||
use :func:`torch.accelerator.current_device_index` by default.
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: A dictionary containing device capability information. The dictionary includes:
|
||||
- ``supported_dtypes`` (set(torch.dtype)): Set of PyTorch data types for which
|
||||
tensors can be allocated on the accelerator and type conversion across
|
||||
supported dtypes are supported. Any operator support outside of that
|
||||
is not guaranteed
|
||||
|
||||
Examples:
|
||||
>>> # xdoctest: +SKIP("requires cuda")
|
||||
>>> # Query capabilities for current device
|
||||
>>> capabilities = torch.accelerator.get_device_capability("cuda:0")
|
||||
>>> print("Supported dtypes:", capabilities["supported_dtypes"])
|
||||
"""
|
||||
device_index = _get_device_index(device, optional=True)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return torch._C._accelerator_getDeviceCapability(device_index)
|
||||
|
||||
|
||||
def set_device_index(device: _device_t, /) -> None:
|
||||
r"""Set the current device index to a given device.
|
||||
|
||||
Args:
|
||||
device (:class:`torch.device`, str, int): a given device that must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
.. note:: This function is a no-op if this device index is negative.
|
||||
"""
|
||||
device_index = _get_device_index(device, optional=False)
|
||||
torch._C._accelerator_setDeviceIndex(device_index)
|
||||
|
||||
|
||||
set_device_idx = deprecated(
|
||||
"Use `set_device_index` instead.",
|
||||
category=FutureWarning,
|
||||
)(set_device_index)
|
||||
|
||||
set_device_idx.__doc__ = r"""
|
||||
(Deprecated) Set the current device index to a given device.
|
||||
|
||||
Args:
|
||||
device (:class:`torch.device`, str, int): a given device that must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
.. warning::
|
||||
|
||||
:func:`torch.accelerator.set_device_idx` is deprecated in favor of :func:`torch.accelerator.set_device_index`
|
||||
and will be removed in a future PyTorch release.
|
||||
"""
|
||||
|
||||
|
||||
def current_stream(device: _device_t = None, /) -> torch.Stream:
|
||||
r"""Return the currently selected stream for a given device.
|
||||
|
||||
Args:
|
||||
device (:class:`torch.device`, str, int, optional): a given device that must match the current
|
||||
:ref:`accelerator<accelerators>` device type. If not given,
|
||||
use :func:`torch.accelerator.current_device_index` by default.
|
||||
|
||||
Returns:
|
||||
torch.Stream: the currently selected stream for a given device.
|
||||
"""
|
||||
device_index = _get_device_index(device, optional=True)
|
||||
return torch._C._accelerator_getStream(device_index)
|
||||
|
||||
|
||||
def set_stream(stream: torch.Stream) -> None:
|
||||
r"""Set the current stream to a given stream.
|
||||
|
||||
Args:
|
||||
stream (torch.Stream): a given stream that must match the current :ref:`accelerator<accelerators>` device type.
|
||||
|
||||
.. note:: This function will set the current device index to the device index of the given stream.
|
||||
"""
|
||||
torch._C._accelerator_setStream(stream)
|
||||
|
||||
|
||||
def synchronize(device: _device_t = None, /) -> None:
|
||||
r"""Wait for all kernels in all streams on the given device to complete.
|
||||
|
||||
Args:
|
||||
device (:class:`torch.device`, str, int, optional): device for which to synchronize. It must match
|
||||
the current :ref:`accelerator<accelerators>` device type. If not given,
|
||||
use :func:`torch.accelerator.current_device_index` by default.
|
||||
|
||||
.. note:: This function is a no-op if the current :ref:`accelerator<accelerators>` is not initialized.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> assert torch.accelerator.is_available() "No available accelerators detected."
|
||||
>>> start_event = torch.Event(enable_timing=True)
|
||||
>>> end_event = torch.Event(enable_timing=True)
|
||||
>>> start_event.record()
|
||||
>>> tensor = torch.randn(100, device=torch.accelerator.current_accelerator())
|
||||
>>> sum = torch.sum(tensor)
|
||||
>>> end_event.record()
|
||||
>>> torch.accelerator.synchronize()
|
||||
>>> elapsed_time_ms = start_event.elapsed_time(end_event)
|
||||
"""
|
||||
device_index = _get_device_index(device, optional=True)
|
||||
torch._C._accelerator_synchronizeDevice(device_index)
|
||||
|
||||
|
||||
class device_index:
|
||||
r"""Context manager to set the current device index for the current :ref:`accelerator<accelerators>`.
|
||||
Temporarily changes the current device index to the specified value for the duration
|
||||
of the context, and automatically restores the previous device index when exiting
|
||||
the context.
|
||||
|
||||
Args:
|
||||
device (Optional[int]): a given device index to temporarily set. If None,
|
||||
no device index switching occurs.
|
||||
|
||||
Examples:
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> # Set device 0 as the current device temporarily
|
||||
>>> with torch.accelerator.device_index(0):
|
||||
... # Code here runs with device 0 as the current device
|
||||
... pass
|
||||
>>> # Original device is now restored
|
||||
>>> # No-op when None is passed
|
||||
>>> with torch.accelerator.device_index(None):
|
||||
... # No device switching occurs
|
||||
... pass
|
||||
"""
|
||||
|
||||
def __init__(self, device: int | None, /) -> None:
|
||||
self.idx = device
|
||||
self.prev_idx = -1
|
||||
|
||||
def __enter__(self) -> None:
|
||||
if self.idx is not None:
|
||||
self.prev_idx = torch._C._accelerator_exchangeDevice(self.idx)
|
||||
|
||||
def __exit__(self, *exc_info: object) -> None:
|
||||
if self.idx is not None:
|
||||
torch._C._accelerator_maybeExchangeDevice(self.prev_idx)
|
||||
@@ -0,0 +1,26 @@
|
||||
import torch
|
||||
from torch.types import Device as _device_t
|
||||
|
||||
|
||||
def _get_device_index(device: _device_t, optional: bool = False) -> int:
|
||||
if isinstance(device, int):
|
||||
return device
|
||||
if isinstance(device, str):
|
||||
device = torch.device(device)
|
||||
device_index: int | None = None
|
||||
if isinstance(device, torch.device):
|
||||
acc = torch.accelerator.current_accelerator()
|
||||
if acc is None:
|
||||
raise RuntimeError("Accelerator expected")
|
||||
if acc.type != device.type:
|
||||
raise ValueError(
|
||||
f"{device.type} doesn't match the current accelerator {acc}."
|
||||
)
|
||||
device_index = device.index
|
||||
if device_index is None:
|
||||
if not optional:
|
||||
raise ValueError(
|
||||
f"Expected a torch.device with a specified index or an integer, but got:{device}"
|
||||
)
|
||||
return torch.accelerator.current_device_index()
|
||||
return device_index
|
||||
@@ -0,0 +1,177 @@
|
||||
import gc
|
||||
from typing import Literal
|
||||
from typing_extensions import Self
|
||||
|
||||
import torch
|
||||
from torch._C import _acceleratorGraph
|
||||
|
||||
|
||||
class Graph(_acceleratorGraph):
|
||||
r"""
|
||||
Wrapper around an :ref:`accelerator<accelerators>` graph that supports capture and replay.
|
||||
|
||||
A graph captures a sequence of operations and their dependencies, allowing them to be
|
||||
replayed efficiently with reduced overhead. This class can be used as a context manager
|
||||
to automatically capture operations on the current stream.
|
||||
|
||||
Arguments:
|
||||
keep_graph (bool, optional): If ``False``, the underlying graph is destroyed and the
|
||||
executable graph is instantiated on the GPU at the end of ``capture_end``.
|
||||
If ``True``, the underlying graph is preserved after ``capture_end``. In this case,
|
||||
the executable graph is not instantiated automatically; it must be explicitly created
|
||||
by calling ``instantiate``, or it will be instantiated on the first call to ``replay``.
|
||||
Defaults to ``False``.
|
||||
pool (tuple[int, int], optional): Memory pool identifier for this graph. Multiple graphs
|
||||
can share the same pool by passing the same identifier, which can reduce memory overhead.
|
||||
Defaults to ``None``.
|
||||
capture_error_mode (Literal["default", "global", "thread_local", "relaxed"], optional):
|
||||
Specifies the behavior of graph capture. The exact semantics are backend-specific.
|
||||
``"default"``: backend-defined default capture behavior.
|
||||
``"global"``: potentially unsafe API calls are prohibited. Errors may occur if capture
|
||||
in the current thread affects other threads.
|
||||
``"thread_local"``: potentially unsafe API calls are prohibited. Errors occur only if
|
||||
capture in the current thread affects itself.
|
||||
``"relaxed"``: the current thread is allowed to make potentially unsafe API calls, except
|
||||
for calls that inherently conflict with stream capture.
|
||||
Default: ``"default"``.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> x = torch.zeros([2000], device=0)
|
||||
|
||||
>>> stream = torch.Stream()
|
||||
>>> graph = torch.accelerator.Graph()
|
||||
>>> with stream, graph:
|
||||
... x += 1
|
||||
|
||||
>>> graph.replay()
|
||||
"""
|
||||
|
||||
def __new__(
|
||||
cls,
|
||||
keep_graph: bool = False,
|
||||
*,
|
||||
pool: tuple[int, int] | None = None,
|
||||
capture_error_mode: Literal[
|
||||
"default", "global", "thread_local", "relaxed"
|
||||
] = "default",
|
||||
) -> Self:
|
||||
return super().__new__(cls, keep_graph)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
keep_graph: bool = False,
|
||||
*,
|
||||
pool: tuple[int, int] | None = None,
|
||||
capture_error_mode: Literal[
|
||||
"default", "global", "thread_local", "relaxed"
|
||||
] = "default",
|
||||
) -> None:
|
||||
super().__init__(keep_graph)
|
||||
self.graph_pool = pool
|
||||
self.capture_error_mode = capture_error_mode
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
def capture_begin(self) -> None:
|
||||
r"""
|
||||
Begin graph capture on the current stream.
|
||||
|
||||
All operations on the current stream after this call will be recorded into the graph until
|
||||
``capture_end`` is called, using the memory pool and capture error mode provided at construction time.
|
||||
"""
|
||||
super().capture_begin(
|
||||
pool=self.graph_pool, capture_error_mode=self.capture_error_mode
|
||||
)
|
||||
|
||||
def capture_end(self) -> None:
|
||||
r"""
|
||||
End graph capture on the current stream of the current device.
|
||||
|
||||
After this call, the graph can be replayed via ``replay``.
|
||||
"""
|
||||
super().capture_end()
|
||||
|
||||
def instantiate(self) -> None:
|
||||
r"""
|
||||
Instantiate the underlying graph. Will be called by ``capture_end``
|
||||
if ``keep_graph=False``, or by ``replay`` if ``keep_graph=True`` and
|
||||
``instantiate`` has not already been explicitly called.
|
||||
"""
|
||||
super().instantiate()
|
||||
|
||||
def replay(self) -> None:
|
||||
r"""Replay the work captured by this graph."""
|
||||
super().replay()
|
||||
|
||||
def reset(self) -> None:
|
||||
r"""
|
||||
Delete the graph currently held by this instance.
|
||||
|
||||
After this call, the graph can be recaptured. Set :attr:`graph_pool` or
|
||||
:attr:`capture_error_mode` beforehand to use different settings on the next capture.
|
||||
"""
|
||||
super().reset()
|
||||
|
||||
def pool(self) -> tuple[int, int]:
|
||||
r"""
|
||||
Return an opaque token representing the id of this graph's memory pool.
|
||||
|
||||
This id can optionally be passed to another graph's ``capture_begin``,
|
||||
which hints the other graph may share the same memory pool.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> g1 = torch.accelerator.Graph()
|
||||
>>> g1.capture_begin()
|
||||
>>> # ... operations ...
|
||||
>>> g1.capture_end()
|
||||
|
||||
>>> # Share g1's memory pool with a new graph
|
||||
>>> pool_id = g1.pool()
|
||||
>>> g2 = torch.accelerator.Graph(pool=pool_id)
|
||||
"""
|
||||
return super().pool()
|
||||
|
||||
def enable_debug_mode(self) -> None:
|
||||
r"""Enable debugging mode for ``debug_dump``."""
|
||||
return super().enable_debug_mode()
|
||||
|
||||
def debug_dump(self, path: str) -> None:
|
||||
r"""
|
||||
Dump the captured graph to a file for debugging purposes if the debugging is
|
||||
enabled via ``enable_debug_mode``.
|
||||
|
||||
Arguments:
|
||||
path (str): Path to dump the graph to.
|
||||
|
||||
Example::
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> s = torch.Stream()
|
||||
>>> g = torch.accelerator.Graph()
|
||||
>>> g.enable_debug_mode()
|
||||
|
||||
>>> with s, g:
|
||||
>>> # ... operations ...
|
||||
|
||||
>>> # Dump captured graph to a file "graph_dump.dot"
|
||||
>>> g.debug_dump("graph_dump.dot")
|
||||
"""
|
||||
return super().debug_dump(path)
|
||||
|
||||
def __enter__(self) -> None:
|
||||
torch.accelerator.synchronize()
|
||||
if torch.compiler.config.force_cudagraph_gc:
|
||||
# We previously always ran garbage collection here. While this can help
|
||||
# reclaim accelerator device memory held by dead Python cycles, it is
|
||||
# very expensive, especially when performing multiple graph captures in sequence.
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
torch.accelerator.empty_host_cache()
|
||||
self.capture_begin()
|
||||
|
||||
def __exit__(self, *exc_info: object) -> None:
|
||||
self.capture_end()
|
||||
|
||||
|
||||
__all__ = ["Graph"]
|
||||
@@ -0,0 +1,248 @@
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from ._utils import _device_t, _get_device_index
|
||||
|
||||
|
||||
__all__ = [
|
||||
"empty_cache",
|
||||
"empty_host_cache",
|
||||
"get_memory_info",
|
||||
"max_memory_allocated",
|
||||
"max_memory_reserved",
|
||||
"memory_allocated",
|
||||
"memory_reserved",
|
||||
"memory_stats",
|
||||
"reset_accumulated_memory_stats",
|
||||
"reset_peak_memory_stats",
|
||||
]
|
||||
|
||||
|
||||
def empty_cache() -> None:
|
||||
r"""Release all unoccupied cached memory currently held by the caching
|
||||
allocator so that those can be used in other application.
|
||||
|
||||
.. note:: This function is a no-op if the memory allocator for the current
|
||||
:ref:`accelerator <accelerators>` has not been initialized.
|
||||
"""
|
||||
if not torch._C._accelerator_isAllocatorInitialized():
|
||||
return
|
||||
torch._C._accelerator_emptyCache()
|
||||
|
||||
|
||||
def empty_host_cache() -> None:
|
||||
r"""Release all unoccupied cached host (pinned) memory currently held by the host caching
|
||||
allocator so that it can be used by other applications.
|
||||
|
||||
.. note:: This function is a no-op if the memory allocator for the current
|
||||
:ref:`accelerator <accelerators>` has not been initialized.
|
||||
"""
|
||||
torch._C._accelerator_emptyHostCache()
|
||||
|
||||
|
||||
def memory_stats(device_index: _device_t = None, /) -> OrderedDict[str, Any]:
|
||||
r"""Return a dictionary of accelerator device memory allocator statistics for a given device index.
|
||||
|
||||
The return value of this function is a dictionary of statistics, each of
|
||||
which is a non-negative integer.
|
||||
|
||||
Core statistics:
|
||||
|
||||
- ``"allocated.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
number of allocation requests received by the memory allocator.
|
||||
- ``"allocated_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of allocated memory.
|
||||
- ``"segment.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
number of reserved segments from device memory allocation.
|
||||
- ``"reserved_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of reserved memory.
|
||||
- ``"active.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
number of active memory blocks.
|
||||
- ``"active_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of active memory.
|
||||
- ``"inactive_split.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
number of inactive, non-releasable memory blocks.
|
||||
- ``"inactive_split_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
|
||||
amount of inactive, non-releasable memory.
|
||||
|
||||
For these core statistics, values are broken down as follows.
|
||||
|
||||
Pool type:
|
||||
|
||||
- ``all``: combined statistics across all memory pools.
|
||||
- ``large_pool``: statistics for the large allocation pool
|
||||
(as of June 2025, for size >= 1MB allocations).
|
||||
- ``small_pool``: statistics for the small allocation pool
|
||||
(as of June 2025, for size < 1MB allocations).
|
||||
|
||||
Metric type:
|
||||
|
||||
- ``current``: current value of this metric.
|
||||
- ``peak``: maximum value of this metric.
|
||||
- ``allocated``: historical total increase in this metric.
|
||||
- ``freed``: historical total decrease in this metric.
|
||||
|
||||
In addition to the core statistics, we also provide some simple event
|
||||
counters:
|
||||
|
||||
- ``"num_alloc_retries"``: number of failed device memory allocation calls that
|
||||
result in a cache flush and retry.
|
||||
- ``"num_ooms"``: number of out-of-memory errors thrown.
|
||||
- ``"num_sync_all_streams"``: number of ``synchronize_and_free_events`` calls.
|
||||
- ``"num_device_alloc"``: number of device memory allocation calls.
|
||||
- ``"num_device_free"``: number of device memory free calls.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
Returns:
|
||||
OrderedDict[str, Any]: an ordered dictionary mapping statistic names to their values.
|
||||
"""
|
||||
if not torch._C._accelerator_isAllocatorInitialized():
|
||||
return OrderedDict()
|
||||
device_index = _get_device_index(device_index, optional=True)
|
||||
stats = torch._C._accelerator_getDeviceStats(device_index)
|
||||
flat_stats = []
|
||||
|
||||
def flatten(prefix: str, value: Any) -> None:
|
||||
if isinstance(value, dict):
|
||||
for k, v in value.items():
|
||||
nested_prefix = f"{prefix}.{k}" if prefix else k
|
||||
flatten(nested_prefix, v)
|
||||
else:
|
||||
flat_stats.append((prefix, value))
|
||||
|
||||
flatten("", stats)
|
||||
flat_stats.sort()
|
||||
# pyrefly: ignore [no-matching-overload]
|
||||
return OrderedDict(flat_stats)
|
||||
|
||||
|
||||
def memory_allocated(device_index: _device_t = None, /) -> int:
|
||||
r"""Return the current :ref:`accelerator<accelerators>` device memory occupied by tensors
|
||||
in bytes for a given device index.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
Returns:
|
||||
int: the current memory occupied by live tensors (in bytes) within the current process.
|
||||
"""
|
||||
return memory_stats(device_index).get("allocated_bytes.all.current", 0)
|
||||
|
||||
|
||||
def max_memory_allocated(device_index: _device_t = None, /) -> int:
|
||||
r"""Return the current :ref:`accelerator<accelerators>` maximum device memory occupied by tensors
|
||||
in bytes for a given device index.
|
||||
|
||||
By default, this returns the peak allocated memory since the beginning of
|
||||
this program. :func:`~torch.accelerator.reset_peak_memory_stats` can be used to
|
||||
reset the starting point in tracking this metric.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
Returns:
|
||||
int: the peak memory occupied by live tensors (in bytes) within the current process.
|
||||
"""
|
||||
return memory_stats(device_index).get("allocated_bytes.all.peak", 0)
|
||||
|
||||
|
||||
def memory_reserved(device_index: _device_t = None, /) -> int:
|
||||
r"""Return the current :ref:`accelerator<accelerators>` device memory managed by the caching allocator
|
||||
in bytes for a given device index.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
Returns:
|
||||
int: the current memory reserved by PyTorch (in bytes) within the current process.
|
||||
"""
|
||||
return memory_stats(device_index).get("reserved_bytes.all.current", 0)
|
||||
|
||||
|
||||
def max_memory_reserved(device_index: _device_t = None, /) -> int:
|
||||
r"""Return the current :ref:`accelerator<accelerators>` maximum device memory managed by the caching allocator
|
||||
in bytes for a given device index.
|
||||
|
||||
By default, this returns the peak cached memory since the beginning of this
|
||||
program. :func:`~torch.accelerator.reset_peak_memory_stats` can be used to reset
|
||||
the starting point in tracking this metric.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
Returns:
|
||||
int: the peak memory reserved by PyTorch (in bytes) within the current process.
|
||||
"""
|
||||
return memory_stats(device_index).get("reserved_bytes.all.peak", 0)
|
||||
|
||||
|
||||
def reset_accumulated_memory_stats(device_index: _device_t = None, /) -> None:
|
||||
r"""Reset the "accumulated" (historical) stats tracked by the current :ref:`accelerator<accelerators>`
|
||||
memory allocator for a given device index.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
.. note:: This function is a no-op if the memory allocator for the current
|
||||
:ref:`accelerator <accelerators>` has not been initialized.
|
||||
"""
|
||||
device_index = _get_device_index(device_index, optional=True)
|
||||
return torch._C._accelerator_resetAccumulatedStats(device_index)
|
||||
|
||||
|
||||
def reset_peak_memory_stats(device_index: _device_t = None, /) -> None:
|
||||
r"""Reset the "peak" stats tracked by the current :ref:`accelerator<accelerators>`
|
||||
memory allocator for a given device index.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
.. note:: This function is a no-op if the memory allocator for the current
|
||||
:ref:`accelerator <accelerators>` has not been initialized.
|
||||
"""
|
||||
device_index = _get_device_index(device_index, optional=True)
|
||||
return torch._C._accelerator_resetPeakStats(device_index)
|
||||
|
||||
|
||||
def get_memory_info(device_index: _device_t = None, /) -> tuple[int, int]:
|
||||
r"""Return the current device memory information for a given device index.
|
||||
|
||||
Args:
|
||||
device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
|
||||
If not given, use :func:`torch.accelerator.current_device_index` by default.
|
||||
If a :class:`torch.device` or str is provided, its type must match the current
|
||||
:ref:`accelerator<accelerators>` device type.
|
||||
|
||||
Returns:
|
||||
tuple[int, int]: a tuple of two integers (free_memory, total_memory) in bytes.
|
||||
The first value is the free memory on the device (available across all processes and applications),
|
||||
The second value is the device's total hardware memory capacity.
|
||||
"""
|
||||
device_index = _get_device_index(device_index, optional=True)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return torch._C._accelerator_getMemoryInfo(device_index)
|
||||
Reference in New Issue
Block a user