Mixture-of-Channels: Exploiting Sparse FFNs for Efficient LLMs Pre-Training and Inference
This addresses the critical memory overhead issue in LLM training and inference, especially when using FlashAttention, offering an efficient solution for scaling models.
The paper tackles the activation memory bottleneck in large language models, particularly from feed-forward networks, by introducing Mixture-of-Channels (MoC), which selectively activates only the Top-K channels per token, resulting in substantial memory savings and throughput gains while maintaining competitive performance.
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.