Yuqi Pan Yupeng Feng Jinghao Zhuang Siyu Ding et al

摘要

主流基于Transformer的大语言模型面临严重的效率瓶颈：训练时的计算量随序列长度呈二次方增长，而推理阶段的内存占用则线性增长，这严重限制了长序列处理能力。此外，在非NVIDIA硬件平台上构建大规模模型也给稳定高效的训练带来了挑战。为应对这些问题，我们提出SpikingBrain，一类受大脑启发的新型模型架构，旨在实现高效的大规模长序列训练与推理。SpikingBrain依托MetaX GPU集群，聚焦三个核心方向：（1）模型架构：采用线性及混合线性注意力机制，并结合自适应脉冲神经元；（2）算法优化：设计了一种高效的基于转换的训练流程，以及专门定制的脉冲编码框架；（3）系统工程：开发了适配MetaX硬件特性的定制化训练框架、算子库及并行策略。基于上述技术，我们构建了两款模型：SpikingBrain-7B（线性大语言模型）与SpikingBrain-76B（混合线性MoE大语言模型）。这两款模型验证了在非NVIDIA平台实现大规模语言模型开发的可行性。SpikingBrain在仅使用约1500亿个token进行持续预训练的情况下，性能可与开源的Transformer基线模型相媲美。我们的模型显著提升了长序列训练效率，并在推理阶段实现了（部分）恒定内存占用与事件驱动的脉冲行为。例如，SpikingBrain-7B在处理400万token长度的序列时，首个token生成时间（Time to First Token）实现了超过100倍的加速。在数百块MetaX C550 GPU上，训练过程可稳定运行数周，其中7B模型的模型FLOPs利用率达到了23.4%。所提出的脉冲机制实现了69.15%的稀疏性，支持低功耗运行。总体而言，本工作展示了受大脑启发的机制在推动下一代高效、可扩展大规模模型设计方面的巨大潜力。