objwatch.wrappers.torch_memory_wrapper module

class objwatch.wrappers.torch_memory_wrapper.TorchMemoryWrapper[source]

Bases: ABCWrapper

TorchMemoryWrapper extends ABCWrapper to log memory statistics from GPU.

This class is designed to gather memory statistics for the GPU memory allocator and track various memory metrics associated with memory allocation, reservation, and freeing. It utilizes the memory_stats function to fetch these statistics and organizes them for further use.

If you want to capture specific GPU memory metrics, you can configure the mem_types variable before starting the tracking. For example, to capture the current and peak allocated memory across all pools, you can add the following code before instantiation of objwatch wrappers:

`python from objwatch.wrappers import TorchMemoryWrapper TorchMemoryWrapper.mem_types = ["allocation.all.current", "allocation.all.peak"] `

mem_types define which statistics should be captured. The available statistics are:

  • "allocation.{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 cudaMalloc().

  • "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 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 October 2019, for size >= 1MB allocations).

  • small_pool: statistics for the small allocation pool (as of October 2019, 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, torch also provide some simple event counters:

  • "num_alloc_retries": number of failed cudaMalloc 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 GPU allocation calls. This includes both cuMemMap and cudaMalloc.

  • "num_device_free": number of GPU free calls. This includes both cuMemUnmap and cudaFree.

The caching allocator can be configured via ENV to not split blocks larger than a defined size (see Memory Management section of the Gpu Semantics documentation). This helps avoid memory fragmentation but may have a performance penalty. Additional outputs to assist with tuning and evaluating impact:

  • "max_split_size": blocks above this size will not be split.

  • "oversize_allocations.{current,peak,allocated,freed}": number of over-size allocation requests received by the memory allocator.

  • "oversize_segments.{current,peak,allocated,freed}": number of over-size reserved segments from cudaMalloc().

The caching allocator can be configured via ENV to round memory allocations in order to reduce fragmentation. Sometimes the overhead from rounding can be higher than the fragmentation it helps reduce. The following stat can be used to check if rounding adds too much overhead:

  • "requested_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}": memory requested by client code, compare this with allocated_bytes to check if allocation rounding adds too much overhead.

For the latest help and more detailed information, please refer to the official PyTorch documentation at: https://pytorch.org/docs/stable/generated/torch.cuda.memory_stats.html#torch-cuda-memory-stats

__init__()[source]

Initialize the TorchMemoryWrapper with optional memory types to capture stats.

Parameters:

mem_types (List[str]) – A list of memory types to capture.

mem_types: List[str] = ['allocation.all.current', 'allocation.all.peak']
_capture_memory() → dict[source]

Capture the current GPU memory statistics.

Returns:

A dictionary of GPU memory stats.

Return type:

dict

_format_memory(stats: dict) → str[source]

Format the memory statistics into a string.

Parameters:

stats (dict) – The memory stats to format.

Returns:

A formatted string representing the memory stats.

Return type:

str

wrap_call(func_name: str, frame: FrameType) → str[source]

Wrap the function call to log memory stats before the function is executed.

Parameters:

  • func_name (str) – Name of the function being called.

  • frame (FrameType) – The current stack frame.

Returns:

A string representing the formatted memory stats.

Return type:

str

wrap_return(func_name: str, result: Any) → str[source]

Wrap the function return to log memory stats after the function is executed.

Parameters:

  • func_name (str) – Name of the function returning.

  • result (Any) – The result returned by the function.

Returns:

A string representing the formatted memory stats.

Return type:

str

_abc_impl = <_abc._abc_data object>
wrap_upd(old_value: Any, current_value: Any) → Tuple[str, str][source]

Wrap the update of a variable to log memory stats after the update.

Parameters:

  • old_value (Any) – The old value of the variable.

  • current_value (Any) – The new value of the variable.

Returns:

Formatted old and new values with memory stats.

Return type:

Tuple[str, str]