CUDA Memory management in PyTorch

TLDR;

• PyTorch cuda api에서는 memory allocator가 미리 넉넉하게 메모리를 선점한 후, 자체적으로 메모리 관리 (캐싱, 할당, 반환...)를 한다.
• nvidia-smi에는 allocator가 미리 선점해 놓은 메모리도 사용량으로 표기한다.
• 그래서 GPU 메모리를 얼마나 사용하고 있는지 정확히 확인하려면, torch.cuda.memory_stats()등을 사용해야 한다.
  

Intro

Demucs 논문에,

One of the main drawbacks of the Demucs model (...중략...) is its large model size, more than 1014MB,
(...중략...) quantization will reduce the model size down to 120MB without any loss of SDR.

라는 내용이 있다.

그래서 내가 구현한 Demucs 모델이 정말 1,014 MiB인지 확인을 해봤다.

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.142.00   Driver Version: 450.142.00   CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   35C    P0    33W /  70W |   2500MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     13890      C   python                           2497MiB |
+-----------------------------------------------------------------------------+

그런데 2,500 MiB나 차지하는 것이었다...

내가 잘못 구현했나 싶어서 간단한 테스트를 해보았다.

import torch

m = torch.nn.Linear(10, 10)
m.cuda()

# model size in MiB
model_size = sum([p.numel() * p.element_size() for p in m.parameters()]) / 1024 / 1024

print(model_size) # 0.00041961669921875

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.142.00   Driver Version: 450.142.00   CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   67C    P0    41W /  70W |   1072MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     28553      C   python                           1069MiB |
+-----------------------------------------------------------------------------+

0.0004 MiB짜리 모델을 cuda에 올리니 1,072 MiB나 차지했다...

Memory Management in Pytorch

PyTorch CUDA semantics - Memory management를 정리해보면 다음과 같다.

Memory management

• Pytorch는 메모리 할당을 빠르게 하기 위해서 caching memory allocator를 사용합니다.
• 이 allocator가 차지하고 있으나, 실제로 사용되지 않는 메모리가 있을 수 있습니다.
  이런 메모리는 nvidia-smi에서는 현재 사용중인 것으로 표기됩니다.
• 그러니까 pytorch에 의해 실제로 사용되는 메모리를 추적하고싶으면, memory_allocated()나 max_memory_allocated()처럼 우리가 만들어놓은 api를 쓰세요.
• 이외에도 memory_reserved(), max_memory_reserved(), empty_cache() 같은 api도 있으니 참고하세요.
• 참고로 empty_cache()는 allocator가 차지하고 있으나, 실제로 안쓰는 메모리를 반환하므로 다른 GPU application이 반환된 메모리를 사용할 수 있습니다.
• memory_stats()는 자세한 cuda memory 관리 상황을 dictionary로 잘 정리해서 반환합니다.

memory_allocated()로 다시 확인해보았다.

import torch

from demucs_clone.modules import Demucs


m = Demucs()
m.cuda()

# display in MiB
print(torch.cuda.memory_allocated() / 1024 /1024) # 1013.73779296875

1,013.7 MiB로, 논문에서 말한 모델 크기와 같음을 확인했다. 다행히도 제대로 구현한 듯 하다.