摘要: Transformer 的训练开销随序列长度呈平方级增长,推理显存占用随序列长度线性增加,这造成了基于 Transformer 的大模型资源消耗巨大、长序列处理能力受限。为应对这些问题,本项目借鉴大脑结构功能与信息处理机制,开发 7B 和 76B 规模的高效类脑脉冲大模型,整个训练与推理全流程在国产算力(沐曦科技曦云 C550)集群上进行。本报告涉及的核心技术点为:(1)在模型层面,基于脉冲神经元构建了具有新型线性或混合线性复杂度的基础模型架构;(2)在算法层面,构建了与现有大模型兼容的通用模型转换技术和高效训练范式,并配套开发了专用的脉冲化编码框架;(3)在工程层面,在国产 GPU 集群开发了兼容的大模型训练框架、Triton/CUDA 算子库、模型并行策略以及集群通信原语。测试评估亮点包括:(1)两款模型的长序列训练效率显著提升,能以极低的数据量实现与众多开源 Transformer 模型相媲美的通用语言建模性能(约为主流大模型的 2%);(2)推理阶段模型的低计算/存储复杂度,结合脉冲事件驱动特性,1M 长度下的 TTFT(生成第一个 Token 所需时间)加速可达 26.5倍, 4M 长度上加速超过 100 倍,长序列处理上展现出数量级的效率和速度提升;(3)面向国产算力集群的训练框架算子加速和通信适配,能保持百卡规模训练的数周稳定运行,7B 模型训练 MFU 超过 23.4%;(4)将压缩到 1B 的类脑模型部署到 CPU 手机端推理框架上,在64k-128k-256k 长度下较 Llama3.2 的 1B 模型解码速度分别提升 4.04×-7.52×-15.39×;(5)细粒度动态阈值脉冲化策略结合粗粒度的 MoE 方案,使得网络整体稀疏性超过 69.15%,为低功耗类脑大模型运行提供有力支撑。本次尝试为国产算力平台上的高效类脑脉冲大模型的研发扩展提供了有益探索和实践,也将启发下一代类脑芯片的设计。
1. 引言
Transformer中,“基本单元”是乘加-非线性构成的,模型的复杂性主要来自于扩大参数规模、扩大数据量规模,本文将其称为“外生复杂性”。但是基于参数的Scaling加上Transformer的平方复杂度给目前的训练、推理造成了比较大的影响,难以继续提升模型规模、上下文长度(尤其是训练阶段)。
本文认为,可以模仿人脑找到一条基于“内生复杂性”的道路:
人脑是目前唯一已知的通用智能系统,人脑包含约 860 亿神经元和约 1000 万亿突触数量、具有丰富的神经元种类、不同神经元又具有丰富的内部结构,但功耗仅 20W 左右。鉴此,我们相信还有另一条路径,称为“基于内生复杂性的通用智能模型”:即找到一条融合神经元丰富动力学特性构建具有生物合理性和计算高效性的神经网络新路径,其将充分利用生物神经网络在神经元和神经环路上的结构和功能特性。
本文结合一个新提出的复杂神经元(前段时间发了Science子刊),加上一个Linear Attention机制,基于转换的方法构建了新的类脑LLM。
2. 背景
- 类脑计算
- 脉冲神经网络
- 神经元内生复杂动力学
- 线性基础模型架构
- 长序列应用前景
3. 技术要点
主要的Contribution:
- 生物启发的稀疏脉冲事件驱动计算 ,将原本的稠密GeMM替换成基于脉冲的稀疏MM;进一步适配MoE;
- 国产集群的大规模分布式训练 ,在MetaX集群上进行训练,完成了框架适配、算子兼容、通信兼容;
- 高效且通用的模型转换工具
- 注意力模块,将现有的Transformer映射到一个稀疏的SWA或者Linear Attention;
- FFN模块,将现有的FFN转换为MoE;
4. 模型架构
4.1. 核心组件
注意力机制
标准的softmax注意力:
Sliding Window Attention:
其中这样注意力变成固定大小,降低到了Linear开销。
另一种Linear Attention直接去掉Softmax:
形成了某种RNN形式。
上述两种机制中,SWQ可以高效建模局部信息,Linear Attention则可以压缩长距离信息。位了坚固多种注意力机制,提出混合注意力机制。
混合注意力的一种实现是层间串行:
或层内并行混合:
本文的SpikingBrain-7B采用层间混合,76B采用层内混合。
混合专家模型MoE
在实践中,为了在已有的 Dense 模型上的基础上高效扩展为 MoE 模型,可以采用上采 样(Upcycling)技术,在不损失原有性能的情况下高效扩展模型规模。其核心步骤 如下:(1)将 Dense 模型中前馈网络的权重参数复制到所有 MoE 专家中,从而保证上采 样后的模型在初始状态下与原模型保持一致;(2)为确保上采样后的 MoE 输出与原 Dense 模型的输出尺度一致,需要对专家输出进行适当缩放。
脉冲神经元
LIF:
存在以下局限性:
- 时间动态特性导致训练复杂度升高、稳定性下降;
- 固定threshold存在神经元沉默和过度激活的情况;
为了解决以上问题,提出了自适应阈值脉冲神经元,让阈值随膜电位发生变化;同时移除衰减因子并使用soft-reset机制。
4.2. 模型整体结构
SpikingBrain-7B
LinearAttention层和4B大小的SWA层1:1堆叠交错,FFN不采用MoE,采用原模型的SwiGLU。将Linear Attention进一步修改为Gated版本:
SpikingBrain-76B-A12B
LA和SWA 1:1层内混合,同时1:6配置标准注意力机制1。并行混合时两个分支都通过RMSNorm保证数值稳定性。采用128个Sink Tokens,在每个输入embedding前插入128个embbeding vector,让所有token的注意力都能关注到它们,并且在它们内部去除casual mask、能够相互关注。
基于Flash Attention实现了上述功能的Kernel。
FFN部分采用MoE,top-1 * 16个Expert。最终保留了7层DenseFFN。
与类脑机制的联系
我们的架构选择与生物大脑中观察到的机理紧密契合。(1) 线性注意力模块展现出与人类记忆类似的建模特性,它依赖于压缩和持续更新的“记忆状态”。在每个时间步,它只从当前记忆中提取信息,呈现出类马尔可夫性质的行为。从生物学角度看,其状态化的时间递归可以被视为具有多分支形态的树突动力学的一种简化抽象。(2) MoE(专家混合)组件体现了模块化稀疏激活和功能分化的原理,这与神经回路中分布式和专业化的信息处理方式非常相似。我们的脉冲编码方案则借鉴了生物系统中事件驱动和自适应稀疏的神经元激活机制。通过将网络层面的稀疏性(MoE)与神经元层面的脉冲稀疏性相结合,我们的方法实现了按需分配的计算,并提供了一种强大的双尺度效率策略。总而言之,这些研究结果为设计既高效又具有生物学合理性的大模型架构提供了一条有前景的途径。
5. 训练范式
5.1. 通用性:注意力图对应
标准注意力:
SWA是它的强局部先验的稀疏化:
Linear Attention是对它的低秩近似:
的rank最大为。
因此,基于 Attention map 的对应关系,我们可以直接以预训练 transformer 检查点的 QKV 投影参数作为初始点,利用已经学到的 QK 相似度信息,通过少量数据的训练,使得 attention 迁移到局部、低秩的 special cases。
在训练过程中,遵循以下要点:
- 对Linear Attention的 施加非负激活 ,如ReLU/Sigmoid函数,因为Softmax具有非负的特性,转换后的应该具有类似的性质;
- 保持转换阶段重新初始化的参数是低秩的 ,如RMSNorm和Gate等。核心是因为转换后finetune希望减少训练开销,如果将大部分参数全部随机化重新开始学习过于困难,并且希望原有的模型可以主导优化的过程。因此选择尽量reuse所有的注意力和FFN部分的权重2。
- 将长序列扩展放到转换阶段完成3。
- 确保性能的转换过程下的充分训练 。进行了全参数训练,不冻结backbone参数。
5.2. 高效性:基于转换的训练方案
转换流程
持续预训练过程包含三个阶段,在完成模型转换的同时逐步扩展上下文窗口长度。第一阶段使用序列长度为 8k、总规模为 100B tokens 的数据进行训练,旨在将注意力图重新调整至局部/低秩的 special case,并使模型损失达到收敛。随后,第二和第三阶段分别逐步将序列长度扩展至 32k 和 128k,并分别使用 20B 至 30B tokens 的数据进行训练。整个转换流程共使用约 150B tokens 的训练数据,以从头训练所需约 10T tokens 的数据量为参照,持续预训练过程仅使用了约 2% 的数据量,在训练资源、预算有限的场景下,高效地完成了模型转换。三个阶段的训练数据均使用 Matrix 数据集,其中长序列数据通过最基础的分 domain 的 packing 策略产生。Rope base 维持基座模型的 1M 大小。 SFT 分三阶段训练,每一个阶段使用不同领域的数据,分阶段逐步提升模型在通用知
识、对话、推理方面的能力。第一阶段,主要侧重于增强模型的基础语言理解和各领域的专
业知识。该阶段使用 Infinity Instruct_fond 作为训练数据集,其中包含科学知识、代码
解析以及数学问题求解各类基础知识。训练数据包含 50w samples ,在 8k 序列长度训练。第
二阶段,训练专注于对话能力与指令遵循专项提升。训练数据 (Infinity Instruct_chat)
包含多轮对话,任务导向和知识问答等多种对话类型。第二阶段训练数据量与训练序列长
度与第一阶段保持一致。第三阶段提升模型在简单推理任务上的能力,这一阶段采用通过
DeepSeek-R1 蒸馏得到的高质量推理数据集,其中包含大量带有详细思维链标
注的推理过程数据,涵盖数学证明、逻辑推理、案例分析等多种需要多步推理的问题类型。
为了确保在不同语言环境下的推理能力,该阶段数据中英文比例控制为 1:1,在序列长度 8k
上共 15w samples 进行训练。
也就是除了训练之外还蒸馏了DS-R1.
MoE上采样
将原来的Dense模型复制成份,随机初始化Router。保证有一个专家始终激活,确保稳定性能。
注意到复制了多份之后权重的scale发生了变化,按照:
则
5.3. 模型的脉冲化编码
5.3.1. 脉冲化编码方案与自适应阈值设计
我们的激活值的脉冲化采用解耦的两步策略,分为优化阶段的基于自适应阈值的整数脉冲发放,以及推理阶段在异步硬件上进行的脉冲化编码。这种方法使得基于整数的表达形式能够在 GPU 上支持计算高效的优化过程,而展开后的脉冲表达形式在与专用硬件结合时,则能提供事件驱动的、高能效的推理。
- 第一步:自适应阈值脉冲化——单步生成整数脉冲计数,维持适当发放活性
该阶段通过简化的自适应阈值神经元建模获得整数脉冲计数,核心是通过自适应阈值设计,让神经元在统计意义上始终保持“不过度兴奋、不过度静息”的适中活性,避免固定阈值导致的脉冲冗余或信息丢失。特别地,我们定义:
其中,代表大模型投影层输入的连续浮点激活值(等价于传统 SNN 中多时间步输入的累积和), 为与膜电位相关的自适应阈值, 为调节神经元发放率的超参数, 为合并生成的整数脉冲计数。
神经元活动:
在优化阶段可以合并为:
这种单步模式在GPU上更加稳定且高效4
自适应阈值及超参数对神经元发放活动的影响可归纳为以下几点:
- 调节脉冲发放水平避免过发放/完全沉默
- 可以调整超参适应不同推理设备所需要的脉冲数量,实现精度和能效的tradeoff5
- 对outlier更适应
- 第二步:脉冲编码——虚拟时间步展开稀疏序列
在推理执行计算的阶段,需将第一步生成的整数脉冲计数 重新展开为时间维度的稀疏脉冲序列,以适配事件驱动计算。其过程表达为:
展开的过程设计了三种方法:
- {0, 1}编码,最简单;
- {-1, 0, 1}编码,允许发放抑制性脉冲,神经元的模式更像对称量化函数,对outlier更友好,理论稀疏度更高;
- 二进制编码,将数值转换为二进制位上对应的bit6,优势是可以极大程度压缩scale。
5.3.2. 硬件适配与落地潜力
但需明确的是:GPU 的同步计算架构无法充分发挥脉冲信号“事件驱动、稀疏异步”的核心优势——现有 GPU 仍需按固定的高频时钟周期处理数据,无法像生物神经系统那样“无脉冲时待机、有脉冲时触发计算”。因此,要完全释放本方案的低能耗潜力,需适配专门的异步硬件架构(如基于异步电路设计的类脑芯片、脉冲处理器)处理矩阵运算:这类硬件可响应稀疏脉冲事件,无需高频时钟同步,无脉冲时电路处于静息低功耗状态,有脉冲时执行加法运算,最大化能效优势,为低功耗类脑大模型的边缘端部署(如工业控制、移动设备)提供可能的落地路径。通过该设计,我们为下一代高能效类脑计算硬件的研发提供了参考技术路径,推动大模型从“算力依赖型”向“能效优化型”升级。
6. MetaX国产集群训练
6.1. 分布式训练适配
- 冷热专家优化 ,在本地备份hot expert减少通信量
- 自适应重计算 ,训练初期的热点专家需要处理大量 token, 当 token 数量超过给定阈值,通过激活重计算技术降低显存开销。当训练后期专家 token 分配趋于均衡,不会触发重计算。
- 多粒度重计算 ,针对显存压力大的专家,将重计算分为三个层级,轻量级重计算仅计算激活函数,router 等,中度重计算会计算全连接层以及共享专家,全量重计算对整个 MoE Layer 进行重计算。通过多粒度重计算,达到计算与存储较好的平衡。
- 长度对齐 ,训练过程中每个专家接收到的 token 数量不同可能引起 GEMM 效率的降低,通过 dropping 与 padding 的方式,对齐需处理 token 的长度,提升整体的 GEMM效率。
6.2. 算子适配
- Triton:JIT编译优化→网格搜索与架构匹配→固化缓存结构→生成机器码
- CUDA→MACA:适用层匹配(CUDA转MACA这样不用改代码)→优化分析→缓存与架构匹配→替换为MACA自有算子库
6.3. 并行技术和训练拓扑
- DP,用了ZeRO优化
- PP,p2p,1F1B
- EP,all2all,分组GEMM
- Sequence Parallel,沿seq分割,DeepSpeed Ulysses
7. 结果展示
7.1 下游任务评测
7.2 长序列处理速度评测
7.3. CPU侧推理
7.4. 国产集群训练指标
我们从训练效率与系统稳定性两个方面衡量 MetaX GPU 集群的训练性能指标。在训练效率指标方面,包括模型 FLOPs 利用率(MFU)和每秒每 GPU 处理 Tokens (TGS),MetaX GPU 集群在整个模型训练过程中表现出较好的性能。其中,SpikingBrain-7B 模型的 TGS 达到 1558,MFU 达到 23.4% (8-way DP, 4-way PP, PP micro-batch size 2, globalbatch size 512),体现了较高的计算效率和有效的资源利用。
在系统稳定性方面,通过系统监测集群运行,集群表现出显著的可靠性和鲁棒性。训练工作持续进行了超过两周的时间,中途过程没有中断,展现了国内硬件和软件生态系统的稳定性和成熟度。
7.5 脉冲化结果分析
具体来说,我们对模型中所有线性投影层的激活输入应用脉冲化编码,同时使用对称 INT8 量化其相应的权重。使用 128 个文本样本的小集合对量化参数进行校准调优。如表 6 所示,在常识推理、MMLU 和 CMMLU 基准测试中的评估表明,该方案下的平均性能下降在 SpikingBrain-7B和-76B 模型中都被限制在约 2% 以内,这证实了方案在精度保留上的有效性
这个话的意思似乎是说,上面的所有测试其实都不是在插入了Neuron的情况下做,而是完全在LinearAttention模型上做的?
7.5.2. 脉冲分布特性
7.5.3. 脉冲可视化
8. 结论
本项目全面展示了我们在沐曦国产 GPU 集群上进行高效类脑大模型训练的创新实践。我们通过一系列关键技术,包括非 Transformer 的新型(混合)线性架构、稀疏 MoE 开发、轻量级通用转换训练方案,以及动态阈值的稀疏脉冲激活编码,成功在数百卡规模的国产集群上验证并践行了高效的类脑脉冲大模型开发流程。
此次研究成功发布了两款模型:线性类脑脉冲大模型 SpikingBrain-7B 和 MoE 混合线性类脑脉冲大模型 SpikingBrain-76B-A12B。这两款模型的核心优势在于:(1)训练效率:模型具备线性或近线性的复杂度,显著提升了长序列训练效率,并能以极低的数据量(仅不到 2%)实现与众多开源 Transformer 模型相媲美的性能。(2)推理性能:模型在推理阶段具有事件驱动的脉冲特性,并展现出常数(7B)或部分层常数(76B)级别的复杂度与存储开销,在长序列处理上实现了数量级的效率和速度飞跃(例如,1M 长度 TTFT 加速超过20 倍)。这一突破性实践不仅为国产算力平台上的高效大模型研发提供了宝贵经验,更对未来大模型的规模化部署与应用开辟了新路径。