训练清华ChatGLM-6B时报错, 原因是显存不够
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 96.00 MiB (GPU 0; 23.70 GiB total capacity; 4.37 GiB already allocated; 64.81 MiB free; 4.37 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
尝试将
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
改为
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().quantize(4).cuda()
仍然报错
RuntimeError: CUBLAS error: CUBLAS_STATUS_NOT_INITIALIZED
排错流程如下
查看服务器显存占用情况
watch -n 0.1 nvidia-smi
发现gpu:0显存被PID:19409程序大量占用, 报错应该是默认在gpu:0训练导致显存不足, 接着查看gpu:0上程序所属用户(如果不是师兄的我就kill了)
top
一看是root的, 惹不起还躲不起嘛, 换张卡跑, 顺嘴一提, 权限内的程序可以kill -9 {pid}掉释放显存
kill -9 19409
发现gpu:1空闲, 指定gpu:1上训练模型, 有多种方法,
(1) 可以在py代码开头(一定要在开头)加
import osos.environ['CUDA_VISIBLE_DEVICES']='1'
这样即可指定在gpu:1上训练, 实际上是只设置gpu:1可见, 而屏蔽其他gpu卡
(2) 可以在代码运行前shell或bash脚本中加
CUDA_VISIBLE_DEVICES=1 python xxx.py
这样即可指定在gpu:1上训练, 实际上是只设置gpu:1可见, 而屏蔽其他gpu卡
(3)在程序中使用set_device()
import torchtorch.cuda.set_device(id)
设置完成后查看显存占用情况可以看到, gpu:1显存占用马上上升了, 不影响其他gpu卡的显存
可以看到清华的ChatGLM-6B约占12G显存(其他卡显存增加是写文章的时候其他小伙伴在跑)
来源地址:https://blog.csdn.net/Hello_World1023/article/details/130355998