依赖项说明
pip install nvitop
pip install timm
pip install peft
后续的显存占用数据截图，均基于nvitop命令实现
1、模型显存占用说明
1.1 理论占用值
在 一文讲明白大模型显存占用（只考虑单卡）与大模型显存占用分析都对模型训练过程中的显存占用进行了分析。但存在一定的出入。为此编写代码进行实测。
但基本上，都是确认训练显存消耗（可估算部分）主要为：模型参数（Model）+ 优化器状态（Optimizer status）+梯度值（Gradient）+激活值（Activation），只是对这个具体为多少存在差异。
大模型显存占用分析观点 ？
占用显存的大头主要分为四部分：模型参数、前向计算过程中产生的中间激活、后向传递计算得到的梯度、优化器状态。这里着重分析参数、梯度和优化器状态的显存占用。训练大模型时通常会采用AdamW优化器，并用混合精度训练来加速训练，基于这个前提分析显存占用。
使用AdamW优化器和混合精度训练来训练参数量为 A的大模型，模型参数、梯度和优化器状态占用的显存大小为 20A。Adam需要一阶参数与二阶参数，所以需要2个float32，8个字节；博主认为但在混合精度训练中，一阶参数是fp16，二阶参数才是float32，应该是6个字节
根据其观点，在不考虑激活值的情况下，模型混合精度训练显存占用为模型参数的20倍。如果是fp32训练显存占用则为24倍。与其他观点相比是多了一个梯度值的参数
一文讲明白大模型显存占用（只考虑单卡） 观点 ？
按照训练运行的逻辑来讲：
Step1:优化器会先备份一份FP32精度的模型权重，初始化好FP32精度的一阶和二阶动量（用于更新权重）。
Step2:开辟一块新的存储空间，将FP32精度的模型权重转换为FP16精度的模型权重。
Step3:运行forward和backward，产生的梯度和激活值都用FP16精度存储。
Step4:优化器利用FP16的梯度和FP32精度的一阶和二阶动量去更新备份的FP32的模型权重。
Step5:重复Step2到Step4训练，直到模型收敛。
我们可以看到训练过程中显存主要被用在四个模块上：
模型权重本身（FP32+FP16）
梯度（FP16）
优化器（FP32），Adam优化器为（FP32+FP32）
激活值（FP16）
根据其观点，在不考虑激活值的情况下，模型混合训练显存占用为模型参数的16倍。如果是fp32训练则为20倍
1.2 显存实测
模型参数量与激活值占用情况。因为考虑了较大的激活值，无法准确衡量显存占用。
import timm
import torch
import ipdb
def print_memory_use(tag):
useed=torch.cuda.memory_reserved() / (1024 ** 3)
torch.cuda.empty_cache()
print(f"{tag}: {useed:.4f}G")
if __name__=="__main__":
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print_memory_use("import时显存占用：")
x = torch.randn(90, 3, 224, 224).to(device)
print_memory_use("torch.randn(90, 3, 224, 224) 显存占用：")
model = timm.create_model('resnet50', pretrained=False, num_classes=102).to(device)
print_memory_use("模型1初始化显存占用：")
model2 = timm.create_model('resnet50', pretrained=False, num_classes=102).to(device)
print_memory_use("模型2初始化显存占用：")
del model2
torch.cuda.empty_cache()
optimizer = torch.optim.Adam(model.parameters())
print_memory_use("adam优化器显存占用：")
from fvcore.nn import flop_count_table, FlopCountAnalysis, ActivationCountAnalysis
print(f'params: {sum(map(lambda x: x.numel(), model.parameters()))/1024/1024:.3f}M')
st=flop_count_table(FlopCountAnalysis(model, x), activations=ActivationCountAnalysis(model, x))
st=str(st).split('\n')[:3]
st="\n".join(st)
print(st)
print_memory_use("FlopCountAnalysis显存占用：")
with torch.no_grad():
model(x)
print_memory_use("torch.no_grad()推理时显存占用：")
out=model(x)
print_memory_use("模型+激活值显存占用：")
loss_func = torch.nn.CrossEntropyLoss()
label=torch.randn(90, 102).to(device)
loss = loss_func(out, label) # 计算loss
print_memory_use("loss计算显存占用：")
loss.backward()
print_memory_use("反向传播显存占用：")
optimizer.step()
print_memory_use("优化时显存占用：")
optimizer.zero_grad() # 梯度置0
print_memory_use("zero_grad显存占用：")
输出结果如下，可以发现模型参数量为23M约对应90M显存，激活值为1G约对应4G显存。
在执行import时，显存占用是0，基于导入torch并不会引起显存变化
初始化x时，显存占用为0.05g
基于模型1与模型2初始化的显存差异，可以确认resnet50的显存占用为0.2441-0.1680=0.0761G=76.1M ，根据参数量22.618M推算理论显存占用应当为90.472M。
同时，基于此，可以推测torch.cuda环境显存占用为0.1680G-0.0761G-0.05G=0.0419G=50M
import时显存占用：: 0.0000G
torch.randn(90, 3, 224, 224) 显存占用：: 0.0508G
模型1初始化显存占用：: 0.1680G
模型2初始化显存占用：: 0.2441G
adam优化器显存占用：: 0.1680G
params: 22.618M
| module | #parameters or shape | #flops | #activations |
|:-----------------------|:-----------------------|:-----------|:---------------|
| model | 23.717M | 0.373T | 1G |
FlopCountAnalysis显存占用：: 7.4160G
torch.no_grad()推理时显存占用：: 1.5469G
模型+激活值显存占用：: 7.4336G
loss计算显存占用：: 7.4043G
反向传播显存占用：: 8.5820G
优化时显存占用：: 0.7520G
zero_grad显存占用：: 0.7344G
将模型修改为resnet101后，输出的信息如下所示。在同等的输入添加下，resnet101的激活值比resnet50多0.461g（约2g的显存），参数多20m（约80m的显存）。
在这里，基于模型1与模型2初始化的显存差异，可以确认resnet101的显存占用为0.3887-0.2207=0.1679G=167.9M ，根据参数量40.731M推算理论显存占用应当为162.924M。这里的差异比resnet要少很多；同时，可以确定resnet101比resnet50要多占用80m的显存
基于此，可以推测torch.cuda环境显存占用为0.2207-0.1679-0.05=0.00280G=2M，这应该是不正确的。
adam优化器，在没有计算时是不占用显存的。
对比两次torch.no_grad()推理时显存占用差异，可以确定torch.cuda环境是在推理或训练时占用1.4G左右的显存。
对比两次模型+激活值显存占用，可以发现差值为3.5G，
import时显存占用：: 0.0000G
torch.randn(90, 3, 224, 224) 显存占用：: 0.0508G
模型1初始化显存占用：: 0.2207G
模型2初始化显存占用：: 0.3887G
adam优化器显存占用：: 0.2207G
params: 40.731M
| module | #parameters or shape | #flops | #activations |
|:-----------------------|:-----------------------|:-----------|:---------------|
| model | 42.709M | 0.709T | 1.461G |
FlopCountAnalysis显存占用：: 11.0273G
torch.no_grad()推理时显存占用：: 1.5996G
模型+激活值显存占用：: 11.0352G
loss计算显存占用：: 11.0332G
反向传播显存占用：: 7.4043G
优化时显存占用：: 1.3809G
zero_grad显存占用：: 1.3301G
#---------------------------------------------------------------------------------------------------
当将模型修改为resnet18时，显存信息如下所示
import时显存占用：: 0.0000G
torch.randn(90, 3, 224, 224) 显存占用：: 0.0508G
模型1初始化显存占用：: 0.1133G
模型2初始化显存占用：: 0.1562G
adam优化器显存占用：: 0.1133G
params: 10.709M
| module | #parameters or shape | #flops | #activations |
|:-----------------------|:-----------------------|:-----------|:---------------|
| model | 11.229M | 0.164T | 0.224G |
FlopCountAnalysis显存占用：: 1.9863G
torch.no_grad()推理时显存占用：: 0.6816G
模型+激活值显存占用：: 1.9961G
loss计算显存占用：: 1.9668G
反向传播显存占用：: 3.6562G
优化时显存占用：: 0.2969G
zero_grad显存占用：: 0.2754G
1.3 无激活值下测试
这里修改代码，将尽可能将激活值设置为0。这里可以发现，fp32下，显存占用大约是模型的24倍左右。
import timm
import torch
import ipdb
def print_memory_use(tag):
useed=torch.cuda.memory_reserved() / (1024 ** 3)
torch.cuda.empty_cache()
print(f"{tag}: {useed:.4f}G")
if __name__=="__main__":
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print_memory_use("import时显存占用：")
x = torch.randn(1, 3, 64, 64).to(device)
print_memory_use("torch.randn(1, 3, 64, 64) 显存占用：")
model = timm.create_model('resnet18', pretrained=False, num_classes=102).to(device)
print_memory_use("模型1初始化显存占用：")
model2 = timm.create_model('resnet18', pretrained=False, num_classes=102).to(device)
print_memory_use("模型2初始化显存占用：")
del model2
torch.cuda.empty_cache()
optimizer = torch.optim.Adam(model.parameters())
print_memory_use("adam优化器显存占用：")
from fvcore.nn import flop_count_table, FlopCountAnalysis, ActivationCountAnalysis
print(f'params: {sum(map(lambda x: x.numel(), model.parameters()))/1024/1024:.3f}M')
st=flop_count_table(FlopCountAnalysis(model, x), activations=ActivationCountAnalysis(model, x))
st=str(st).split('\n')[:3]
st="\n".join(st)
print(st)
print_memory_use("FlopCountAnalysis显存占用：")
with torch.no_grad():
model(x)
print_memory_use("torch.no_grad()推理时显存占用：")
out=model(x)
print_memory_use("模型+激活值显存占用：")
loss_func = torch.nn.CrossEntropyLoss()
label=torch.randn(1, 102).to(device)
loss = loss_func(out, label) # 计算loss
print_memory_use("loss计算显存占用：")
loss.backward()
print_memory_use("反向传播显存占用：")
optimizer.step()
print_memory_use("优化时显存占用：")
optimizer.zero_grad() # 梯度置0
print_memory_use("zero_grad显存占用：")
当设置为resnet18模型时，输入如下
import时显存占用：: 0.0000G
torch.randn(1, 3, 64, 64) 显存占用：: 0.0020G
模型1初始化显存占用：: 0.0625G
模型2初始化显存占用：: 0.1055G
adam优化器显存占用：: 0.0625G
params: 10.709M
| module | #parameters or shape | #flops | #activations |
|:-----------------------|:-----------------------|:----------|:---------------|
| model | 11.229M | 0.149G | 0.203M |
FlopCountAnalysis显存占用：: 0.0859G
torch.no_grad()推理时显存占用：: 0.0840G
模型+激活值显存占用：: 0.0859G
loss计算显存占用：: 0.0645G
反向传播显存占用：: 0.1074G
优化时显存占用：: 0.2520G
zero_grad显存占用：: 0.2109G
当设置为resnet50模型时，输入如下
import时显存占用：: 0.0000G
torch.randn(1, 3, 64, 64) 显存占用：: 0.0020G
模型1初始化显存占用：: 0.1172G
模型2初始化显存占用：: 0.1934G
adam优化器显存占用：: 0.1172G
params: 22.618M
| module | #parameters or shape | #flops | #activations |
|:-----------------------|:-----------------------|:----------|:---------------|
| model | 23.717M | 0.338G | 0.907M |
FlopCountAnalysis显存占用：: 0.1445G
torch.no_grad()推理时显存占用：: 0.1367G
模型+激活值显存占用：: 0.1445G
loss计算显存占用：: 0.1230G
反向传播显存占用：: 0.2168G
优化时显存占用：: 0.4805G
zero_grad显存占用：: 0.4043G
当设置为resnet101模型时，输入如下
import时显存占用：: 0.0000G
torch.randn(1, 3, 64, 64) 显存占用：: 0.0020G
模型1初始化显存占用：: 0.1699G
模型2初始化显存占用：: 0.3379G
adam优化器显存占用：: 0.1699G
params: 40.731M
| module | #parameters or shape | #flops | #activations |
|:-----------------------|:-----------------------|:----------|:---------------|
| model | 42.709M | 0.644G | 1.325M |
FlopCountAnalysis显存占用：: 0.1992G
torch.no_grad()推理时显存占用：: 0.1914G
模型+激活值显存占用：: 0.1992G
loss计算显存占用：: 0.1992G
反向传播显存占用：: 0.3613G
优化时显存占用：: 0.8418G
zero_grad显存占用：: 0.6738G
2、实际占用测试
这里通过对比，发现对于小模型而言，全量训练与冻结部分参数、lora训练差异不大（lora训练反而更消耗资源），影响模型训练显存需求的反而是fp32与混合精度的差异。
2.1 完整代码
基于以下代码实现模型不同情况下的训练
import torch
import torchvision
from torch.autograd import Variable
import torch.utils.data.dataloader as Data
import timm
from torchvision.transforms import transforms
from tqdm import tqdm
from torch.amp import GradScaler
from torch.amp import autocast
from peft import LoraConfig,get_peft_model
from peft.tuners import lora
def freeze_model(model,freeze_rate=0.95):
if True:
#对layer进行冻结---------
layers=[]
for name, module in model.named_modules():
layers.append(name)
freeze_index=int(len(layers)*freeze_rate)
freeze_layers=layers[:freeze_index]
for name, module in model.named_modules():
if name in freeze_layers:
module.requires_grad=False
module.requires_grad_=False
def set_lora_model(model):
# 可以按照个人需求冻结模型
# for name, param in model.named_parameters():
# param.requires_grad = False
target_modules=[]
# 获取要继续lora操作的层
for name,module in model.named_modules():
if isinstance(module,torch.nn.Conv2d):
target_modules.append(name)
# 实例化LoraConfig
config = LoraConfig(
r=4,
lora_alpha=32,
lora_dropout=0.1,
bias="none",
task_type=None,
init_lora_weights="gaussian",
target_modules=target_modules,
)
lora_model = get_peft_model(model, config)
for name,module in model.named_modules():
if isinstance(module,lora.layer.Conv2d):
for name, param in module.base_layer.named_parameters():
param.requires_grad = False
module.in_channels=module.base_layer.in_channels
module.out_channels=module.base_layer.out_channels
module.kernel_size=module.base_layer.kernel_size
module.stride=module.base_layer.stride
module.padding=module.base_layer.padding
module.dilation=module.base_layer.dilation
module.groups=module.base_layer.groups
return lora_model
def train():
train_loss = 0.
train_acc = 0.
n=0
d_len=0
pbar= tqdm(total=len(train_loader),desc='Train: ')
for batch_x, batch_y in train_loader:
batch_x, batch_y = Variable(batch_x).to(device), Variable(batch_y).to(device)
# print(batch_x.shape,batch_y.shape)
optimizer.zero_grad() # 梯度置0
d_len+=batch_x.shape[0]
if train_amp:
#混合精度运算作用域
with autocast(device_type='cuda'):
out = model(batch_x) # 前向传播
loss = loss_func(out, batch_y) # 计算loss
#将梯度进行相应的缩放
scaler.scale(loss).backward() # 返向传播
#设置优化器计步
scaler.step(optimizer)
#更新尺度
scaler.update()
else:
out = model(batch_x) # 前向传播
loss = loss_func(out, batch_y) # 计算loss
loss.backward()
optimizer.step()
# ------计算loss，acc
train_loss += loss.item()
# torch.max(out, 1) 指第一维最大值，返回[最大值，最大值索引]
pred = torch.max(out, 1)[1]
train_correct = (pred == batch_y).sum()
train_acc += train_correct.item()
n += batch_y.shape[0]
pbar.update(1)
pbar.set_postfix({
'loss': '%.4f' % (train_loss / n),
'train acc': '%.3f' % (train_acc / n),
'dlen':d_len
})
pbar.close()
print('Train Loss: {:.6f}, Acc: {:.6f}'.format(train_loss / (len(
train_data)), train_acc / (len(train_data))) ,batch_x.shape)
def eval():
model.eval()
eval_loss = 0.
eval_acc = 0.
n=0
d_len=0
pbar= tqdm(total=len(test_loader),desc='Test: ')
for batch_x, batch_y in test_loader:
# 测试阶段不需要保存梯度信息
with torch.no_grad():
batch_x, batch_y = Variable(batch_x).to(device), Variable(batch_y).to(device)
if train_amp:
with autocast(device_type='cuda'):
out = model(batch_x)
loss = loss_func(out, batch_y)
else:
out = model(batch_x)
loss = loss_func(out, batch_y)
eval_loss += loss.item()
pred = torch.max(out, 1)[1]
num_correct = (pred == batch_y).sum()
eval_acc += num_correct.item()
d_len+=batch_x.shape[0]
n+=1
pbar.update(1)
pbar.set_postfix({
'loss': '%.4f' % (eval_loss / n),
'eval acc': '%.3f' % (eval_acc / d_len),
'dlen':d_len
})
pbar.close()
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(
test_data)), eval_acc / (len(test_data))))
transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
train_data = torchvision.datasets.Flowers102(
'./data', split = "train", transform=transform, download=True
)
test_data = torchvision.datasets.Flowers102(
'./data', split = "test", transform=transform
)
train_loader = Data.DataLoader(dataset=train_data, batch_size=90, shuffle=True)
test_loader = Data.DataLoader(dataset=test_data, batch_size=90)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = timm.create_model('resnet18', pretrained=True, num_classes=102).to(device)
if __name__=="__main__":
train_amp = True #
freeze_95 = True
train_lora = True
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
2.2 全量训练-fp32
设置三个关键参数为False，可以使模型在fp32下进行训练
if __name__=="__main__":
train_amp = False # True
freeze_95 = False
train_lora = False
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
当进行全量训练时的显存占用如下，为5G
同时训练速度为2.3it/s
2.3 冻结部分参数-fp32
仅设置freeze_95=True，可以使模型在冻结95%的参数下，采用fp32进行训练
if __name__=="__main__":
train_amp = False # True
freeze_95 = True
train_lora = False
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
冻结部分参数训练时的显存占用如下，为5G
同时训练速度也为2.3it/s
2.4 lora训练-fp32
基于以下代码进行lora训练模型
if __name__=="__main__":
train_amp = False # True
freeze_95 = False
train_lora = True
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
冻结部分参数训练时的显存占用如下，为5.1G，比正常训练要高
同时模型参数迭代速度下降
2.5 全量训练-fp16
基于以下代码进行模型训练
if __name__=="__main__":
train_amp = True #
freeze_95 = False
train_lora = False
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
可以发现模型显存占用为3G
同时，参数迭代速度为2.6it/s
2.6 冻结部分参数-fp16
基于以下代码进行模型训练
if __name__=="__main__":
train_amp = True #
freeze_95 = True
train_lora = False
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
可以发现模型显存占用为3G
同时，参数迭代速度为2.6it/s
2.7 lora训练-fp16
基于以下代码进行模型训练
if __name__=="__main__":
train_amp = True #
freeze_95 = False
train_lora = True
if freeze_95:
freeze_model(model,freeze_rate=0.95)
if train_lora:
set_lora_model(model)
print(model)
optimizer = torch.optim.Adam(filter(lambda p : p.requires_grad, model.parameters()))
loss_func = torch.nn.CrossEntropyLoss()
scaler = GradScaler()
for epoch in range(10):
print('epoch {}'.format(epoch + 1))
# training-----------------------------
train()
# evaluation--------------------------------
eval()
可以发现模型显存占用为3.2G
同时，参数迭代速度为2.5it/s
扫一扫
314
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
