97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
"""
|
||
一些显存调试的工具函数
|
||
"""
|
||
import torch
|
||
|
||
|
||
def get_model_memory_usage(model):
|
||
total_param_memory = 0
|
||
param_info = {}
|
||
|
||
for name, param in model.named_parameters():
|
||
# 计算单个参数的显存（单位：MB）
|
||
param_size = param.data.element_size() * param.data.nelement() / 1024 ** 2
|
||
total_param_memory += param_size
|
||
param_info[name] = {
|
||
"size_mb": f"{param_size:.2f} MB",
|
||
"shape": param.shape,
|
||
"dtype": param.dtype,
|
||
}
|
||
|
||
param_info["total_param_memory"] = f"{total_param_memory:.2f} MB"
|
||
return param_info
|
||
|
||
|
||
def get_optimizer_memory_usage(optimizer):
|
||
total_state_memory = 0
|
||
state_info = {}
|
||
|
||
for group in optimizer.param_groups:
|
||
for p in group["params"]:
|
||
if p in optimizer.state:
|
||
state = optimizer.state[p]
|
||
for key, value in state.items():
|
||
if isinstance(value, torch.Tensor):
|
||
# 计算优化器状态张量的显存
|
||
state_size = value.element_size() * value.nelement() / 1024 ** 2
|
||
total_state_memory += state_size
|
||
state_info[f"Param_{id(p)}_{key}"] = {
|
||
"size_mb": f"{state_size:.2f} MB",
|
||
"shape": value.shape,
|
||
"dtype": value.dtype,
|
||
}
|
||
|
||
state_info["total_optimizer_state_memory"] = f"{total_state_memory:.2f} MB"
|
||
return state_info
|
||
|
||
|
||
def log_memory_snapshot(model, optimizer, step):
|
||
print(f"--- Memory Snapshot (Step {step}) ---")
|
||
|
||
# 模型参数显存
|
||
model_mem = get_model_memory_usage(model)
|
||
print("Model Parameters:")
|
||
for name, info in model_mem.items():
|
||
print(f" {name}: {info}")
|
||
|
||
# 优化器状态显存
|
||
optimizer_mem = get_optimizer_memory_usage(optimizer)
|
||
print("\nOptimizer State:")
|
||
for name, info in optimizer_mem.items():
|
||
print(f" {name}: {info}")
|
||
|
||
# 总显存占用（模型 + 优化器）
|
||
total_memory = float(model_mem["total_param_memory"].split()[0]) + \
|
||
float(optimizer_mem["total_optimizer_state_memory"].split()[0])
|
||
print(f"\nTotal Memory: {total_memory:.2f} MB")
|
||
|
||
|
||
class ActivationMemoryTracker:
|
||
def __init__(self):
|
||
self.activation_info = {}
|
||
self.hooks = []
|
||
|
||
def track_activations(self, model):
|
||
for name, module in model.named_modules():
|
||
hook = module.register_forward_hook(self._hook_fn(name))
|
||
self.hooks.append(hook)
|
||
|
||
def _hook_fn(self, name):
|
||
def hook(_, __, output):
|
||
# 计算激活值的显存（单位：MB）
|
||
if isinstance(output, torch.Tensor):
|
||
mem = output.element_size() * output.nelement() / 1024 ** 2
|
||
self.activation_info[name] = {
|
||
"size_mb": f"{mem:.2f} MB",
|
||
"size": mem,
|
||
"shape": output.shape,
|
||
"dtype": output.dtype,
|
||
}
|
||
return hook
|
||
|
||
def clear(self):
|
||
for hook in self.hooks:
|
||
hook.remove()
|
||
self.activation_info = {}
|
||
self.hooks = []
|