colossalai.utils.memory_tracer

class colossalai.utils.memory_tracer.AsyncMemoryMonitor(power=10)

An Async Memory Monitor runing during computing. Sampling memory usage of the current GPU at interval of 1/(10**power) sec.

The idea comes from Runtime Memory Tracer of PatrickStar PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management https://arxiv.org/abs/2108.05818

Parameters: power (int) – the power of time interval, defaults to 10

Usage:

```python

async_mem_monitor = AsyncMemoryMonitor() input = torch.randn(2, 20).cuda() OP1 = torch.nn.Linear(20, 30).cuda() OP2 = torch.nn.Linear(30, 40).cuda()

async_mem_monitor.start() output = OP1(input) async_mem_monitor.finish() async_mem_monitor.start() output = OP2(output) async_mem_monitor.finish() async_mem_monitor.save(‘log.pkl’)

```