第一种情况
如果这个报错后面跟了想要占用多少显存但是不够这样的字眼,如下:
解决办法就很简单了:
- 改小batchsize,batchsize砍半可以差不多省掉一半的显存
- 推理阶段加上with torch.no_grad(),这个可以将修饰的代码段不要梯度,可以省掉很多显存
- 改小input的shape,例如(224,224)->(112,112)这样可以省掉一半的显存
- 换小的网络结构
- 用多卡训练,torch可以用model = nn.DataParallel(model)启用多卡训练,终端用CUDA_VISIBLE_DEVICES=0,1 python3 train.py启动即可,这样会将batchsize等份的分给n张卡,这里的示例是2张卡,这样相当于减小了单卡的batchsize,就不会OOM了。
- 开启FP16,就是浮点数截断,可以省一大部分显存
- 改代码参考hugging face的gradient checkpoint方法,前向传播不计算梯度,反向传播的时候才计算梯度,用时间换显存。参考:https://github.com/haotian-liu/LLaVA/blob/80540fb4bf4dad118d87d42bd2fb55e6f3b96f16/llava/train/train.py#L816
- 改代码用deepspeed启动,部分tensor将存储在内存来减少显存使用。参考:https://github.com/haotian-liu/LLaVA/blob/80540fb4bf4dad118d87d42bd2fb55e6f3b96f16/scripts/v1_5/finetune.sh#L3
- transformer结构的模型加入Lora或者qLora,冻结大部分网络参数仅训练额外Lora的参数。参考:https://github.com/haotian-liu/LLaVA/blob/80540fb4bf4dad118d87d42bd2fb55e6f3b96f16/scripts/v1_5/finetune_lora.sh#L4
- 超级省显存的教程:https://huggingface.co/blog/ram-efficient-pytorch-fsdp,可以训70B的LLaMa2.可以参考优化下代码。
- 钞能力
第二种情况
直接报错OOM没有说申请多少显存而不够的字眼。如下:
这个情况比较特殊,有多种原因:
- 原因一:linux下某个用户存在了显存泄露,如果是自己的账号有显存泄露,执行fuser -v /dev/nvidia*然后将提示的进程kill掉即可,如果是其他用户显存泄露,需要管理员权限,执行以下命令,执行前最好跟所有用户通知下,不然再跑的任务会被杀了。但是出现这样显存泄露的情况,所有的用户都会用不了显卡,所以应该不会有任务在跑(猜测):
sudo fuser -v /dev/nvidia* |awk '{for(i=1;i<=NF;i++)print "kill -9 " $i;}' | sudo sh
- 原因二:网上看到的,据说模型加载的参数和自己pytorch的版本不匹配就会报错这个,例如你的pretrain使用torch1.1.0训的,你用torch1.2.0的代码加载这个参数就有可能报错。这个没经历过,仅仅看到记录下。