[论文解读] EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism

Yanxi Chen, Xuchen Pan+3 more|arXiv (Cornell University)|Dec 8, 2023

Topic Modeling被引用 4

一句话总结

EE-LLM 提出了一种可扩展的框架，通过使用三维并行（张量并行、流水线并行和数据并行）来训练和推理早期退出的大规模语言模型。它实现了流水线并行阶段间的高效反向传播，利用空闲计算资源进行早期退出计算，并支持与 KV 缓存兼容的推理，从而在 13B 和 7B Llama 2 等大规模模型上实现近乎零的训练开销，且不损失模型质量的情况下获得显著的推理加速。

ABSTRACT

We present EE-LLM, a framework for large-scale training and inference of early-exit large language models (LLMs). While recent works have shown preliminary evidence for the efficacy of early exiting in accelerating LLM inference, EE-LLM makes a foundational step towards scaling up early-exit LLMs by supporting their training and inference with massive 3D parallelism. Built upon Megatron-LM, EE-LLM implements a variety of algorithmic innovations and performance optimizations tailored to early exiting, including a lightweight method that facilitates backpropagation for the early-exit training objective with pipeline parallelism, techniques of leveraging idle resources in the original pipeline schedule for computation related to early-exit layers, and two approaches of early-exit inference that are compatible with KV caching for autoregressive generation. Our analytical and empirical study shows that EE-LLM achieves great training efficiency with negligible computational overhead compared to standard LLM training, as well as outstanding inference speedup without compromising output quality. To facilitate further research and adoption, we release EE-LLM at https://github.com/pan-x-c/EE-LLM.

研究动机与目标

• 为了在大规模场景下实现早期退出 LLM 的训练与推理，解决现有框架缺乏支持的问题。
• 为了解决使用流水线并行训练早期退出 LLM 的挑战，即多个损失项分布在不同阶段的问题。
• 通过优化梯度累积和资源利用率，将训练过程中的计算开销最小化。
• 支持在早期退出推理中实现高效、自回归生成，且兼容 KV 缓存。
• 发布一个生产就绪的开源框架，以促进社区采纳和进一步研究。

提出的方法

• 提出一种轻量级梯度累积方案，通过调整损失缩放和微批次处理，实现对流水线并行早期退出模型的无偏反向传播。
• 利用流水线调度中的空闲计算资源，在前向传播过程中执行早期退出层的计算，提升硬件利用率。
• 提出两种与 KV 缓存兼容的推理策略：一种是在所有退出点之间共享 KV 缓存，另一种是采用按退出点管理的缓存机制。
• 采用三维并行（张量并行、流水线并行和数据并行）来扩展大规模早期退出模型的训练与推理能力，突破单个设备的容量限制。
• 通过最小化早期退出头的额外计算成本，尤其是大型输出嵌入矩阵，优化训练效率。
• 在梯度估计中使用方差减少技术，结合微批次中早期退出和最终输出的梯度，降低梯度噪声。

实验结果

研究问题

• RQ1 能否通过流水线并行实现大规模早期退出 LLM 的高效训练，尽管传统上这与多退出训练目标存在冲突？
• RQ2 与标准 LLM 训练相比，训练早期退出 LLM 的计算开销是多少，是否可以被最小化？
• RQ3 如何使早期退出推理与 KV 缓存兼容，以支持自回归生成？
• RQ4 在大规模 LLM 中，早期退出对推理加速和输出质量的影响如何？
• RQ5 所提出的框架是否能在不损害模型准确率的前提下，实现高训练效率和高推理加速？

主要发现

• 与标准 LLM 训练相比，EE-LLM 实现了近乎零的训练开销，早期退出头的额外计算成本可忽略不计。
• 该框架支持大规模早期退出 LLM 的训练，包括 13B 和 7B Llama 2 模型，且使用了完整的三维并行。
• EE-LLM 的推理实现了显著的加速——平均提速最高达 2.5 倍，同时输出质量与完整模型生成相当。
• 实验表明，所提出的梯度累积方法将梯度方差降低了最多 30%，提升了训练稳定性。
• 该框架实现了高效的、与 KV 缓存兼容的推理，使早期退出在自回归生成任务中具备可行性。
• 实证结果表明，使用 EE-LLM 训练的早期退出 LLM 即使在简单输入上提前退出，也能在下游任务中保持强劲性能。