Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference
Tong Wu, Yutong He, Bin Wang, Kun Yuan
大型语言模型(LLM)在各种人工智能任务中取得了显着的成功,这是由扩展法律推动的,这些法律将模型大小和训练数据与性能改进相关联。 然而,这种扩展范式会承受大量的内存开销,给训练和推理都带来了重大挑战。 虽然现有的研究主要解决了参数和优化器状态内存的减少,但激活内存 - 特别是来自前馈网络(FFN)的激活内存已成为关键瓶颈,特别是当实现FlashAttention时。 在这项工作中,我们对LLM进行详细的内存分析,并将FFN激活确定为激活内存开销的主要来源。 受此激励,我们引入了Mixture-of-Channels(MoC),这是一种新颖的FFN架构,仅选择性地激活SwiGLU本地门接机制确定的每个令牌的Top-K最相关的通道。 MoC在预训练期间大大减少了激活内存,并通过通过部分重量加载到GPU SRAM中来减少内存访问,从而提高了推理效率。 广泛的实验验证了 MoC 在保持竞争模型性能的同时,可显著节省内存和吞吐量。
Large language models (LLMs) have demonstrated remarkable success across diverse artificial intelligence tasks, driven by scaling laws that correlate model size and training data with performance improvements. However, this scaling paradigm incurs substantial memory overhead, creating significant challenges for both training and inference. While existing research has primarily addressed parameter and optimizer state memory reduction, activation memory-particularly from feed-forward networks (FFN...
