通过稀疏化前馈网络，显著降低大模型训练与推理的显存占用。
Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference
- 仅激活每个令牌最相关的Top-K通道，基于SwiGLU门控机制实现稀疏计算。
- 训练时减少60%以上激活显存，推理时提升吞吐量30%以上。
- 适合资源受限场景下的大模型高效训练与部署，尤其适配GPU显存有限的设备。
大规模语言模型在多种人工智能任务中表现卓越，其性能随模型规模和训练数据增加而提升。然而，这种扩展范式带来巨大内存开销，尤其在使用FlashAttention时，前馈网络（FFN）的激活内存成为主要瓶颈。本文通过细致的内存分析发现，FFN激活是激活内存的主要来源。为此，提出新型FFN架构Mixture-of-Channels（MoC），根据SwiGLU原生门控机制，仅对每个令牌激活前K个最相关通道。该方法在预训练中大幅减少激活内存，并通过部分权重加载至GPU SRAM降低内存访问，在推理中提升效率。大量实验表明，MoC在保持竞争力性能的同时，显著节省内存并提升吞吐量。
原文摘要 · Abstract (English)
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 (FFNs)-has become the critical bottleneck, especially when FlashAttention is implemented. In this work, we conduct a detailed memory profiling of LLMs and identify FFN activations as the predominant source to activation memory overhead. Motivated by this, we introduce Mixture-of-Channels (MoC), a novel FFN architecture that selectively activates only the Top-K most relevant channels per token determined by SwiGLU's native gating mechanism. MoC substantially reduces activation memory during pre-training and improves inference efficiency by reducing memory access through partial weight loading into GPU SRAM. Extensive experiments validate that MoC delivers significant memory savings and throughput gains while maintaining competitive model performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。