colossalai.utils.memory_tracer

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

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')
```