SpikingBrain技术报告：基于尖峰大脑灵感的大型模型

主流的基于Transformer的大型语言模型面临着重大的效率瓶颈：训练计算的复杂度与序列长度呈平方关系，而推理内存线性增长，限制了长上下文处理。在非NVIDIA平台上构建大型模型也给稳定和高效的训练带来了挑战。为了解决这个问题，我们提出了SpikingBrain，这是一系列旨在高效进行长上下文训练和推理的灵感模型。SpikingBrain利用MetaX GPU集群，并关注三个方面：（1）模型架构：线性和混合线性注意力架构与自适应尖峰神经元；（2）算法优化：高效的基于转换的训练管道和专用的尖峰编码框架；（3）系统工程：为MetaX硬件量身定制的训练框架、操作符库和并行策略。通过这些技术，我们开发了两个模型：SpikingBrain-7B，一个线性大型语言模型，以及SpikingBrain-76B，一个混合线性MoE大型语言模型。这些模型展示了在非NVIDIA平台上开展大规模LLM开发的可行性。SpikingBrain的性能可与开源Transformer基准相媲美，同时仅使用约150B的token进行持续预训练。我们的模型显著提高了长序列训练的效率，并以（部分）恒定内存和事件驱动的尖峰行为提供推理。例如，SpikingBrain-7B在4M-token序列的首次token时间上获得了超过100倍的加速。训练在数百个MetaX C550 GPU上保持稳定，7B模型的FLOPs利用率达到了23.4%。提出的尖峰方案达到了69.15%的稀疏度，支持低功耗操作。总体而言，这项工作展示了灵感于大脑的机制在推动下一代高效可扩展的大模型设计中的潜力。