这是为数不多深入比较使用消费级 GPU(RTX 3090、4090)和服务器显卡(A800)进行大模型预训练、微调和推理的论文。
大型语言模型 (LLM) 在学界和业界都取得了巨大的进展。但训练和部署 LLM 非常昂贵,需要大量的算计资源和内存,因此研究人员开发了许多用于加速 LLM 预训练、微调和推理的开源框架和方法。然而,不同硬件和软件堆栈的运行时机能可能存在很大差异,这使得选择最佳配置变得困难。
最近,一篇题为《Dissecting the Runtime Performance of the Training, Fine-tuning, and Inference of Large Language Models》的新论文从宏观和微观的角度详细分析了 LLM 训练、微调、推理的运行时机能。
论文地址:https://arxiv.org/pdf/2311.03687.pdf
具体来说,该研究首先在三个 8-GPU 上对不同规模(7B、13B 和 70B 参数)的 LLM,面向预训练、微调、服务进行端到端的机能基准测试,涉及具有或不具有单独优化技术的平台,包括 ZeRO、量化、重新算计、FlashAttention。然后,该研究进一步提供了子模块的详细运行时分析,包括 LLM 中的算计和通信运算符。
方法介绍
该研究的基准测试采用自上而下的方法,涵盖 Llama2 在三个 8-GPU 硬件平台上的端到端步骤光阴机能、模块级光阴机能和运算符光阴机能,如图 3 所示。
三个硬件平台分别为 RTX4090、RTX3090 和 A800,具体规格参数以下表 1 所示。
在软件方面,该研究比较了 DeepSpeed 和 Megatron-LM 在预训练和微调方面的端到端步骤光阴。为了评估优化技术,该研究使用 DeepSpeed 逐一启用以下优化:ZeRO-2、ZeRO-3、offloading、激活重算计、量化和 FlashAttention,以衡量机能改进以及光阴和内存消耗方面的下降。
在 LLM 服务方面,存在三个高度优化的系统,vLLM、LightLLM 和 TGI,该研究在三个测试平台上比较了它们的机能(延迟和吞吐量)。
为了保证结果的准确性和可重复性,该研究算计了 LLM 常用数据集 alpaca 的指令、输入和输出的平均长度,即每个样本 350 个 token,并随机生成字符串以达到 350 的序列长度。
在推理服务中,为了综合利用算计资源并评估框架的鲁棒性和效率,所有请求都以突发模式调度。实验数据集由 1000 个合成句子组成,每个句子包含 512 个输入token。该研究在同一 GPU 平台上的所有实验中始终保持「最大生成 token 长度」参数,以保证结果的一致性和可比性。
端到端机能
该研究通过预训练、微调和推理不同尺寸 Llama2 模型(7B、13B 和 70B)的步骤光阴、吞吐量和内存消耗等指标,来衡量在三个测试平台上的端到端机能。同时评估了三个广泛使用的推理服务系统:TGI、vLLM 和 LightLLM,并重点关注了延迟、吞吐量和内存消耗等指标。
模块级机能
LLM 通常由一系列模块(或层)组成,这些模块可能具有独特的算计和通信特性。例如,构成 Llama2 模型的关键模块是 Embedding、LlamaDecoderLayer、Linear、SiLUActivation 和 LlamaRMSNorm。
预训练结果
在预训练实验环节,研究者首先分析了三个测试平台上不同尺寸模型(7B、13B 和 70B)的预训练机能(迭代光阴或吞吐量、内存消耗),然后进行了模块和操作层面的微基准测试。
端到端机能
研究者首先进行实验来比较 Megatron-LM 和 DeepSpeed 的机能,二者在 A800- 80GB 服务器上预训练 Llama2-7B 时没有使用任何内存优化技术(比如 ZeRO)。
他们使用的序列长度为 350,并为 Megatron-LM 和 DeepSpeed 提供了两组批巨细,从 1 到最大批巨细。结果以下表 II 所示,以训练吞吐量(tokens / 秒)和消费级 GPU 内存(单位 GB)为基准。
结果表明,当批巨细都为 1 时,Megatron-LM 稍快于 DeepSpeed。不过当批巨细达到最大时,DeepSpeed 在训练速度上最快。当批巨细相同时,DeepSpeed 消耗了比基于张量并行的 Megatron-LM 更多的 GPU 内存。即使批巨细很小,这两个系统都消耗了大量的 GPU 内存,导致 RTX4090 或 RTX3090 GPU 服务器的内存溢出。
在训练 Llama2-7B(序列长度为 350,批巨细为 2)时,研究者使用了带有量化的 DeepSpeed 来研究不同硬件平台上的扩展效率。结果以下图 4 所示,A800 几乎是线性扩展,RTX4090 和 RTX3090 的扩展效率略低,分别为 90.8% 和 85.9%。在 RTX3090 平台上,NVLink 连接比没有 NVLink 时的扩展效率提升了 10%。
研究者使用 DeepSpeed 来评估不同内存和算计高效方法下的训练机能。为公平起见,所有评估设置成序列长度为 350,批巨细为 1,默认加载模型权重为 bf16。
对于具有卸载功能的 ZeRO-2 和 ZeRO-3,他们分别将优化器状态和优化器状态 + 模型卸载到 CPU RAM。对于量化,他们使用了具有双重量化的 4bits 配置。此外报告了 NVLink 失效时 RTX3090 的机能(即所有数据通过 PCIe 总线传输)。结果以下表 III 所示。
研究者通过最大化每种方法的批巨细以获得最大吞吐量,进一步利用不同 GPU 服务器的算计能力。结果以下表 IV 所示,表明了增加批巨细可以很容易改进训练过程。因此,具有高带宽和大内存的 GPU 服务器比消费级 GPU 服务器更适合做全参数混合精度训练。
模块级分析
下表 V 展示了单步预训练 Llama2-7B 模型的前向、后向和优化器的整体及算计核心光阴成本。对于后向阶段,由于总光阴包含了非重叠光阴,算计核心光阴远远小于前向阶段和优化器。如果非重叠光阴从后向阶段中移除,该值变成 94.8。
重算计和 FlashAttention 的影响
加速预训练的技术大致可以分为两类:节省内存增加批巨细、加速算计核心。以下图 5 所示,GPU 在前向、后向和优化器阶段有 5-10% 的光阴处于闲置状态。
研究者认为这种闲置光阴是批巨细较小造成的,因此他们测试了所有技术可以使用的最大批巨细。最终利用重算计来增加批巨细,利用 FlashAttention 来加速算计核心分析。
以下表 VII 所示,随着批巨细的增加,前向和后向阶段的光阴大幅增加,GPU 闲置光阴几乎没有。
下表 VIII 显示出,FlashAttention 分别可以将前向和后向的注意力模块加速 34.9% 和 24.7%。
微调结果
在微调环节,研究者主要讨论参数高效微调方法(PEFT),展示 LoRA 和 QLoRA 在各种模型巨细和硬件设置下的微调机能。使用序列长度为 350,批巨细为 1,默认将模型权重加载到 bf16。
结果以下表 IX 所示,使用 LoRA 和 QLoRA 微调 Llama2-13B 的机能趋势与 Llama2-7B 保持一致。与 lama2-7B 相比,微调 Llama2-13B 的吞吐量下降了约 30%。
不过当结合所有优化技术时,即使 RTX4090 和 RTX3090 也可以微调 Llama2-70B,实现 200 tokens / 秒的总吞吐量。
推理结果
端到端机能
下图 6 显示了各种硬件平台和推理框架下吞吐量的全面分析,其中省略了 Llama2-70B 的相关推理数据。其中 TGI 框架展现了卓越的吞吐量,尤其是 RTX3090 和 RTX4090 等具有 24GB 内存的 GPU。此外 LightLLM 在 A800 GPU 平台上的机能显著优于 TGI 和 vLLM,吞吐量几乎翻倍。
这些实验结果表明,TGI 推理框架在 24GB 内存 GPU 平台上具有卓越的机能,而 LightLLM 推理框架在 A800 80GB GPU 平台上表现出最高的吞吐量。这一发现表明 LightLLM 专门针对 A800/A100 系列高机能 GPU 进行了优化。
不同硬件平台和推理框架下的延迟表现以下图 7、8、9、10 所示。
综上所示,A800 平台在吞吐量和延迟方面均显著优于 RTX4090 和 RTX3090 两款消费级平台。并且在两款消费级平台中,RTX3090 比 RTX4090 略有优势。当在消费级平台上运行时,TGI、vLLM 和 LightLLM 三个推理框架在吞吐量方面没有表现出实质性差异。相比之下,TGI 在延迟方面始终优于其他两个。在 A800 GPU 平台上,LightLLM 在吞吐量方面表现最好,其延迟也非常接近 TGI 框架。
更多实验结果参考原论文。