[Paper] 可扩展的尖峰神经网络构建，使用多达数千块 GPU

发布: 1个月前 (2025年12月10日 GMT+8 18:27)

7 min read

原文: arXiv

Source: arXiv - 2512.09502v1

概览

本文提出了一种在现代 GPU 集群上构建和运行大规模脉冲神经网络（SNN）模拟的新方法，能够扩展到数千个 GPU。通过重新设计网络构建阶段并利用基于 MPI 的通信，作者使得在下一代 exascale 机器上以所需性能模拟皮层规模模型（数十亿突触）成为可能。

网络划分 – 将完整的 SNN 划分为本地子网络，每个 MPI 进程（即每个 GPU）对应一个子网络。每个进程接收一个随机种子和全局连通规则的描述（如连接概率、距离依赖型分布）。
本地构建 – 使用种子，每个 rank 独立生成其前突触和后突触伙伴列表。算法将连接存储为压缩稀疏行（CSR）格式，天然适配 GPU 内存。
脉冲交换准备 – 为每个远程目标 rank 分配一个 脉冲发送缓冲区。作者预先计算路由表，指明哪些输出脉冲需要打包发送到哪些目标。
通信层 – 评估了两种基于 MPI 的方法：
- 点对点：每个 rank 仅对实际需要通信的 rank 发起非阻塞发送/接收（稀疏通信）。
- 集合式：当网络足够密集、几乎所有 rank 每个时间步都交换脉冲时，使用 MPI_Alltoallv。
仿真循环 – 构建阶段结束后，常规的 GPU 核函数推进神经元状态、产生脉冲、将其打包进预计算缓冲区并触发 MPI 交换。所有入站脉冲解包完毕后进入下一个时间步。

大规模脑模型 – 研究者现在可以在现有 GPU 集群上模拟具有真实突触数量的皮层柱或全脑片段，将从数周到数天的洞察时间大幅缩短。
神经科学在环 AI – 高效运行大规模 SNN 为将深度学习与生物学上合理的脉冲动力学相结合的混合 AI 系统打开了大门。
Exascale 准备度 – 构建和通信模式已针对即将到来的 exascale 架构（如 NVLink、高速互连）进行设计，使代码对国家实验室和云服务商的未来硬件具备前瞻性。
工具链集成 – 由于实现基于广泛使用的 NEST 模拟器，开发者可以在不重写底层 GPU 代码的情况下插入自定义神经元模型、可塑性规则或传感器接口。
性能感知设计 – 论文的基准方法为其他 HPC 开发者提供了评估点对点与集合式通信在不规则工作负载下的模板。

静态连通性的假设 – 当前流水线仅在网络构建一次后使用；结构可塑性等动态重连需要重新构建或增量更新，本文未涉及。
GPU 内存受限 – 对于极密网络，即使使用 CSR 压缩也可能超出单 GPU 内存；作者建议将来采用 out‑of‑core 技术。
硬件特异性 – 基准聚焦于 NVIDIA GPU 与 InfiniBand；在 AMD GPU 或新兴互连（如 Slingshot）上的表现仍待验证。
超出 2 000 GPU 的可扩展性 – 虽然算法在理论上已具备 exascale 可行性，但在 >2 000 GPU（或真正的 exascale 系统）上的实测仍留待后续工作。