EE-LLM：利用 3D 并行性进行大规模提前退出大语言模型的训练与推理
机器学习
2024-06-18 v3 人工智能
分布式、并行与集群计算
摘要
我们提出了 EE-LLM，一个用于大规模训练和推理提前退出大语言模型的框架。虽然近期工作已初步证明了提前退出在加速 LLM 推理方面的有效性，但 EE-LLM 通过支持大规模 3D 并行训练和推理，为扩展提前退出 LLM 迈出了基础性的一步。EE-LLM 基于 Megatron-LM 构建，实现了多种针对提前退出量身定制的算法创新和性能优化，包括一种支持在流水线并行下对提前退出训练目标进行反向传播的轻量级方法、利用原始流水线调度中空闲资源进行与提前退出层相关计算的技术，以及两种与自回归生成的 KV 缓存兼容的提前退出推理方法。我们的分析和实证研究表明，与标准 LLM 训练相比，EE-LLM 在计算开销可忽略不计的情况下实现了出色的训练效率，并在不牺牲输出质量的前提下实现了显著的推理加速。为了促进进一步的研究和采用，我们在 https://github.com/pan-x-c/EE-LLM 发布了 EE-LLM。
引用
@article{arxiv.2312.04916,
title = {EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism},
author = {Yanxi Chen and Xuchen Pan and Yaliang Li and Bolin Ding and Jingren Zhou},
journal= {arXiv preprint arXiv:2312.04916},
year = {2024}
}
备注
ICML 2024 camera-ready version