大语言模型推理中的显存优化 有哪些

大语言模型推理中的显存优化 有哪些

显存优化背景

在大语言模型推理时，显存是显著瓶颈。以开源的BLOOM 176B模型为例，在8张A100计算卡上，通常对话设置下仅能进行批量为10左右的推理。为缓解显存压力，常采用显存池、显存卸载/加载等技术。

Offloading/Checkpoint原理

显存卸载（Offloading）是应对大量中间结果引发显存问题的技术。当显存容量有限时，把暂时用不到的中间结果乃至模型参数转移到主存甚至磁盘。在个人工作站等显存和计算卡数量受限的推理场景中适用，存储卸载和加载过程能与计算过程重叠，保障整体吞吐量。为减少参数存储卸载/加载次数，可在一次加载参数后进行较多批量的推理。不过，因计算能力有限，大语言模型会有较大延时，适用于时间敏感度低的推理任务。

举例

假设要处理一批包含100篇长文本的情感分析任务，使用大语言模型进行推理。模型在处理文本过程中会产生大量中间结果，如每一层神经网络计算后的输出等。

由于工作站显存有限，无法同时存储所有中间结果和模型参数。这时就可