目录
一.引言
nvidia-smi 命令可以查看 LLM 训练与推理中的实时显存占用与 gpu-util,为了避免一直使用命令行监控显卡变化同时为了把指标数据做本地收集和后续处理,这里新增两个处理任务,分别周期性的收集 GPU 运行日志并根据运行日志进行周期的指标可视化。
二.常见 GPU 监控指令
最常用的指令是直接在命令行:
watch -n 5 nvidia-smi
5 代表每几 s 执行一次对应命令,nvidia-smi 是官方提供的 GPU 列表查看命令。通过上述命令可以实现最基本的周期监控,但是需要实时盯着面板且无法直接获取对应的显存、Util 信息。
三.周期回收 GPU 日志
1.功能分类
鉴于 watch -n 命令的局限性,我们通过 python 脚本定期回收 GPU 日志,主要使用下面两个类:
- pynvml 负责获取 GPU 指标
- sched 负责周期调度获取指标
1.1 初始化变量
这里分别初始化了 pynvml 显卡监控类、scheduler 周期调度类,并定义了存储全部日志的 all_log,用于记录训练开始和结束的 Boolean 以及日志轮次的 round_count 参数,最后的 end_count 和 end_limit 负责在训练、推理任务结束后达到一定限制优雅退出监控程序。
1.2 添加统计信息
首先获取当前机器的 GPU 卡数量,随后遍历每张卡的 id 获取当前的显存使用情况与 GPU 利用率。这里有两个 if 判断,第一个负责判断 GPU 是否开始工作,第二个则是负责在开始后记录 GPU 的空闲次数。
1.3 保存与退出
gpu_info 负责记录本次检测的全部 GPU 卡信息,并追加至 all_log 中,随后执行一次 dump 操作将信息存储至 json。最后的 if 负责判断空置多久后停止监控程序,如果是常驻的 7x24h 监控,则无需添加该 if 条件,监控可以做到不间断。
2.完整函数
只需定义 interval 参数规定多久监控一次即可。
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import json
import pynvml
import time
import sched
# 初始化 pynvml 库
pynvml.nvmlInit()
# 创建 sched 对象
s = sched.scheduler(time.time, time.sleep)
# 定义轮次计数器
round_count = 0
# 追加轮次日志
all_log = []
# 标记开始状态
is_start_working = False
# 退出次数累计
end_count = 0
end_limit = 10
def get_gpu_info(interval):
global round_count
global end_count
global end_limit
global is_start_working
round_count += 1
# 打开文件并读取原有数据
file_path = "gpu_info.json"
# 获取显卡数量
device_count = pynvml.nvmlDeviceGetCount()
used = []
util = []
# 遍历每个显卡
for i in range(device_count):
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
# 获取显存使用情况
meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
used.append(f"{meminfo.used / 1024 / 1024 / 1024:.2f}")
# 获取 GPU 利用率
utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
util.append(f"{utilization.gpu}")
if utilization.gpu > 10:
is_start_working = True
if utilization.gpu == 0 and is_start_working:
end_count += 1
gpu_info = {"epoch": round_count, "used": ",".join(used), "util": ",".join(util)}
all_log.append(gpu_info)
# 将显卡信息转换为 JSON 格式并存储到文件中
with open(file_path, 'w', encoding='utf8') as f:
json.dump(all_log, f, ensure_ascii=False, indent=4)
# 延迟 5 秒钟后再次执行该函数
if end_count <= end_limit:
s.enter(interval, 1, get_gpu_info, argument=(interval,))
# 第一次执行函数
interval = 10
get_gpu_info(interval)
# 开始任务调度
s.run()
下面是一次训练后的监控日志样例,由于是单机四卡环境,所以每个 epoch 有四条记录逗号隔开:
四.GPU 指标可视化
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import json
import matplotlib.pyplot as plt
color_dict = {0: 'blue', 1: 'green', 2: 'red', 3: 'cyan'}
style_dict = {0: '-', 1: '--', 2: '-', 3: "--"}
def load_and_plot(file_path):
metric_map = {}
epoch = []
info = json.load(open(file_path, "r"))
for js in info:
epoch.append(js["epoch"])
used_batch = js["used"].split(",")
util_batch = js["util"].split(",")
for index, gpu_info in enumerate(zip(used_batch, util_batch)):
used = gpu_info[0]
util = gpu_info[1]
# 初始值
if index not in metric_map:
metric_map[index] = {"used": [], "util": []}
# 填充
metric_map[index]["used"].append(float(used))
metric_map[index]["util"].append(float(util))
gpu_ids = list(metric_map.keys())
gpu_ids.sort()
for id in gpu_ids:
used = metric_map[id]["used"]
util = metric_map[id]["util"]
plt.plot(epoch, used, label=f"GPU{id}")
plt.plot(epoch, util, c=color_dict[id], linestyle=style_dict[id], label=f"GPU{id}")
plt.title("GPU Util")
plt.legend()
plt.show()
file_path = "./gpu_info.json"
load_and_plot(file_path)
这里直接读取 json 信息获取对应 GPU id 的 Memory Used 与 GPU Util 指标即可。
- Memory Used
- GPU Util
这里训练过程比较稳定,所以显存和 Util 的变化与波动都不大,对于一些推理场景,可能会出现显存的阶跃,这时候定期的监控图会更有参考意义。
五.总结
上面给出了 GPU 指标定期监控与可视化的脚本,同学们可以把两个 py 穿插在训练、推理等 GPU 运行期间的任务,nohup 启动即可,在显卡停止工作时自动停止,记录本次 GPU 工作情况,灰常的方便。