开源模型效果对比及优化：星火13B显存及内存排查

本文将介绍一个任务，即部署几个开源的模型，并将本地经过全量微调的模型与开源模型进行效果对比。

部署的开源模型包括：星火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。后期如有遇到此类问题，可以借鉴之。