背景
现在大语言模型越来越大,占用的内存越来越多,这导致内存较小的设备无法体验大模型的效果。transformer
提供了将一个大模型分别加载在gpu和cpu上的方法。
加载方法
- 以多模态模型BLIP2为例,将其语言模型放在gpu上,其余部分放在cpu上。配置加载预加载模型的
device_map
.device_map
可以设置为auto
,则根据设备的显存情况,自动加载在gpu或者cpu上。 - 使用BLIP2模型地址
device_map = {'language_model':0,\
'language_projection':'cpu', \
'qformer':'cpu', \
'query_tokens':'cpu', \
'vision_model':'cpu'}
multimodal = Blip2ForConditionalGeneration.from_pretrained("huggingface_opt",torch_dtype=torch.float16,device_map=device_map)
如上所示,将language_model
放在0号显卡上,其余模型放在cpu上。模型的名称在模型的index文件中
测试情况
加载位置 | 显存占用 | 耗时 | 备注 |
---|---|---|---|
GPU | 13G | 0.5s | |
GPU +CPU | 10G | 2s | 语言放在0号显卡上,其余模型放在cpu上 |