Repository: pytorch
Issue: [RFC] A device-agnostic Python device memory related API design for stream-based accelerators
Labels: triaged,module: python frontend

# Motivation
This RFC aims to propose a design for a series of generic memory-related APIs tailored for stream-based accelerators to help users simplify the runtime code written for different devices (including out-of-tree devices, a registration mechanism is provided for them).

# Background
Currently, PyTorch only provides device-specific memory-related APIs usually used in other components or workloads, like torch/ao(https://github.com/pytorch/ao/pull/753/) and huggingface([reference](https://github.com/huggingface/transformers/blob/fd3238b4b0e1af4fae4a293cbfd1251ead40cd29/src/transformers/trainer_utils.py#L532)). And the user can call this API like this.
```python
if device = "cuda":
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()
    max_memory_reserved = torch.cuda.max_memory_reserved()
elif device= "xpu":
    torch.xpu.reset_peak_memory_stats()
    torch.xpu.empty_cache()
    max_memory_reserved = torch.xpu.max_memory_reserved()
elif ...
```
PyTorch needs to provide a generic memory-related API design to help users write device-agnostic code.

# Design
Our design is following [[RFC] A device-agnostic Python runtime API design for stream-based accelerators](https://github.com/pytorch/pytorch/issues/128403) and generalize `torch.xxx.foo` to `torch.acc.foo`. So we can simplify this above example:
```python
torch.reset_peak_memory_stats(device_type)
torch.empty_cache(device_type)
max_memory_reserved = torch.max_memory_reserved(device_type)
```
We can give a list including the most used memory-related APIs.

<div align="center">
<table>
<tr>
<td> Device-specific memory APIs torch.xxx.foo</td> <td> Device-agnostic memory APIs torch.foo</td>
</tr>
<tr>
<td>

```python
torch.xxx.empty_cache
```

</td>
<td>

```python
torch.acc.empty_cache
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.reset_peak_memory_stats
```

</td>
<td>

```python
torch.acc.reset_peak_memory_stats
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.reset_accumulated_memory_stats
```

</td>
<td>

```python
torch.acc.reset_accumulated_memory_stats
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.memory_stats
```

</td>
<td>

```python
torch.acc.memory_stats
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.memory_allocated
```

</td>
<td>

```python
torch.acc.memory_allocated
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.max_memory_allocated
```

</td>
<td>

```python
torch.acc.max_memory_allocated
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.memory_reserved
```

</td>
<td>

```python
torch.acc.memory_reserved
```

</td>
</tr>

<tr>
<td>

```python
torch.xxx.max_memory_reserved
```

</td>
<td>

```python
torch.acc.max_memory_reserved
```

</td>
</tr>

</table>
</div>

**Our goal** is torch.acc.foo can cover the most common runtime scenarios and usages. And using **if**/**else** statement and torch.xxx.foo as an additional supplement in other situations.

### Additional context
We will implement this design on top of [
[RFC] Intel GPU Runtime Upstreaming for Allocator](https://github.com/pytorch/pytorch/issues/116322).
After we unify a device allocator abstract, given as `DeviceAllocator`(temp name), it will define memory API interface like `emptyCache`, `resetPeakStats`, `resetAccumulatedStats` and `getDeviceStats`. And each backend can implement its own device allocator and leverage the allocator registration mechanism to register it. This is also a nice solution for the out-of-tree backend.
At this moment, we can use `GetAllocator(device_type)->emptyCache()` as a unified API in the C++ side for each accelerator. Finally, we implement python API `torch.acc.empty_cache` on top of the C++ unified APi `GetAllocator(device_type)->emptyCache()`, illustrated as below.
![image](https://github.com/user-attachments/assets/104efcc3-f627-4756-bc1c-b24cf5a9d1e2)

cc @albanD