frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "[MTIA] (3/n) Implement PyTorch APIs to query/reset device peak memory usage (#143347)"

Browse Source 
This reverts commit efe21ee59dfdd6642cc693e69e07aa9d8be13eb9.

Reverted https://github.com/pytorch/pytorch/pull/143347 on behalf of https://github.com/huydhn due to D67118173 has been backed out internally ([comment](https://github.com/pytorch/pytorch/pull/143347#issuecomment-2557983266))
This commit is contained in:
PyTorch MergeBot
2024-12-21 04:04:16 +00:00
parent dabc9566c4
commit c7d7eff798
3 changed files with 12 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/source/mtia.memory.rst
View File

@ -11,4 +11,3 @@ The MTIA backend is implemented out of the tree, only interfaces are be defined
:nosignatures:
memory_stats
max_memory_allocated

1
docs/source/mtia.rst
View File

@ -19,7 +19,6 @@ The MTIA backend is implemented out of the tree, only interfaces are be defined
is_available
is_initialized
memory_stats
max_memory_allocated
get_device_capability
empty_cache
record_memory_history

25
torch/mtia/memory.py
View File

@ -10,19 +10,6 @@ from . import _device_t, is_initialized
from ._utils import _get_device_index
def max_memory_allocated(device: Optional[_device_t] = None) -> int:
r"""Return the maximum memory allocated in bytes for a given device.
Args:
device (torch.device, str, or int, optional) selected device. Returns
statistics for the current device, given by current_device(),
if device is None (default).
"""
if not is_initialized():
return 0
return memory_stats(device).get("dram", 0).get("peak_bytes", 0)
def memory_stats(device: Optional[_device_t] = None) -> Dict[str, Any]:
r"""Return a dictionary of MTIA memory allocator statistics for a given device.
@ -36,6 +23,18 @@ def memory_stats(device: Optional[_device_t] = None) -> Dict[str, Any]:
return torch._C._mtia_memoryStats(_get_device_index(device, optional=True))
def max_memory_allocated(device: Optional[_device_t] = None) -> int:
r"""Return the maximum memory allocated in bytes for a given device.
Args:
device (torch.device or int, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.mtia.current_device`,
if :attr:`device` is ``None`` (default).
"""
return memory_stats(device=device).get("allocated_bytes.all.peak", 0)
__all__ = [
"memory_stats",
"max_memory_allocated",