本文将介绍一个任务,即部署几个开源的模型,并将本地经过全量微调的模型与开源模型进行效果对比。
部署的开源模型包括:星火13B,Baichuan2-13B, ChatGLM6B等。
其他两个模型基于transformers架构封装,因此推理服务启动还是十分丝滑,但星火13B是基于Megatron-DeepSpeed框架实现,启动推理服务的过程中发现启动13B的显卡占用71G-78G,有些反直觉。
此文将整理开源星火13B的显存及内存排查并优化的整理过程,至于哪家开源模型效果好,不在此文的讨论范围内。
直观上来说,13B的模型,数据类型为bf16,显卡占用大概在26G左右,但星火13B直接占用70G+,不可思议。排查原因的过程中,少不了源码的调试与分析。在排查的过程中,启动推理服务的文件run_iFlytekSpark_text_generation.py中,model_provider方法是初始化模型并加载模型文件的方法。
其中,加载权重文件可以看到,加载state_dict时,直接将权重文件加载到显卡中,而非加载至CPU,然后再执行to方法,转移到GPU。因此该处是一个潜在的优化点。
再打入iFlytekSparkModel内部,词表Embedding层,线性转换层,等初始化weight时,也是直接将weight分配在GPU上运行。
优化方案:
1. 模型初始化时,模型的Embedding,线性层的权重weight均直接加载至GPU,因此可以优化为先将这些weight加载至CPU。改进的方式也很简单,从上面的源码层面,可以看到,当增加参数” use_cpu_initialization”,将使用CPU进行初始化权重,因此只需要在启动推理服务的脚本中增加” --use-cpu-initialization”参数即可。
2. 加载模型文件时,直接加载至GPU,然后run_iFlytekSpark_text_generation.py中的get_model方法中,当模型加载完成后,会进行分配至GPU以及FP16的转换的操作。因此,优化的方式也很简单,可以优化为先加载至CPU,再运行get_model中的默认分配至GPU,加载完后,再使用垃圾回收机制清除CPU占用的内存即可。
效果对比:
(1) 优化前的显卡占用: 71.5G
(2) 优化前的内存占用: 虚拟内存占用94.5G
(3) 优化后的显卡占用: 26G
(4) 优化后的内存占用: 43.1G
总结:
本文主要是针对开源星火13B的显存及内存占用过大的一个代码优化。核心思想是使用CPU预加载模型,再转换至GPU。后期如有遇到此类问题,可以借鉴之。