[GPU] pynvml 모듈로 gpu 사용량 체크하기

GPU 모니터링 툴이 별도로 없을 때 pynvml을 사용하여 GPU 사용량 정보 등을 로그로 남길 수 있는 방법을 정리하고자 한다.

 

1. GPU 정보 확인하기

import pynvml
print("Driver Version:", pynvml.nvmlSystemGetDriverVersion())

>>>> 
Driver Version: 4xx.1xx.xx

 

 

2. GPU 사용량 체크하기

GPU가 여러 장 있을 때 GPU별 utilization, memory 사용량 평균을 다음과 같이 기록해둘 수 있다.

from threading import Thread
import numpy as np
import torch

....


if __name__ == "__main__":
    
    def schedule_gpu_memory_logging():
        def log_gpu_usage():
            if not torch.cuda.is_available():
                return

            from pynvml.smi import nvidia_smi

            
            nvsmi = nvidia_smi.getInstance()
            res_memory_list = nvsmi.DeviceQuery("memory.free, memory.total, memory.used")["gpu"]
            res_util_list = nvsmi.DeviceQuery("utilization.gpu, memory.free, memory.total, memory.used")["gpu"]


            total = np.average([each["fb_memory_usage"]["total"] for each in res_memory_list])
            used = np.average([each["fb_memory_usage"]["used"] for each in res_memory_list])
            free = np.average([each["fb_memory_usage"]["free"] for each in res_memory_list])
            percentage = 100*used/total
                
            utilization = np.average([each["utilization"]["gpu_util"] for each in res_util_list])
            
            print(
                f'GPU Usage. Util: {utilization} Used: {used} Total: {total} ({percentage}% used). Free: {free}'
            )
        
        def log_loop():
            while True:
                log_gpu_usage()
                time.sleep(15)

        t = Thread(target=log_loop, daemon=True)
        t.start()

    schedule_gpu_memory_logging()
    
    main()
    ....

 

logger와 함께 쓰면 다음과 같이 기록된다.

2023-12-09 02:24:54,667 INFO GPU Usage. Util: 100.0 Used: 34550.25 Total: 40536.1875 (85.23310190431698% used). Free: 5985.9375
2023-12-09 02:25:09,701 INFO GPU Usage. Util: 99.9375 Used: 35386.875 Total: 40536.1875 (87.29699851521558% used). Free: 5149.3125
2023-12-09 02:25:24,732 INFO GPU Usage. Util: 100.0 Used: 34897.25 Total: 40536.1875 (86.08912715336143% used). Free: 5638.9375
2023-12-09 02:25:39,764 INFO GPU Usage. Util: 100.0 Used: 36318.5 Total: 40536.1875 (89.59525362369118% used). Free: 4217.6875
2023-12-09 02:25:54,797 INFO GPU Usage. Util: 100.0 Used: 35287.625 Total: 40536.1875 (87.05215555853643% used). Free: 5248.5625

 

 

[참고자료]

• https://github.com/marianokamp/export/blob/main/hf1/HF_LoRA_Memory.ipynb