Current AI training infrastructure is dominated by single instruction multiple data (SIMD) and systolic array architectures, such as Graphics Processing Units (GPUs) and Tensor Processing Units (TPUs), that excel at accelerating parallel workloads and dense vector matrix multiplications. Potentially more efficient neural network models utilizing sparsity and recurrence cannot leverage the full power of SIMD processor and are thus at a severe disadvantage compared to today's prominent parallel architectures like Transformers and CNNs, thereby hindering the path towards more sustainable AI. To overcome this limitation, we explore sparse and recurrent model training on a massively parallel multiple instruction multiple data (MIMD) architecture with distributed local memory. We implement a training routine based on backpropagation through time (BPTT) for the brain-inspired class of Spiking Neural Networks (SNNs) that feature binary sparse activations. We observe a massive advantage in using sparse activation tensors with a MIMD processor, the Intelligence Processing Unit (IPU) compared to GPUs. On training workloads, our results demonstrate 5-10x throughput gains compared to A100 GPUs and up to 38x gains for higher levels of activation sparsity, without a significant slowdown in training convergence or reduction in final model performance. Furthermore, our results show highly promising trends for both single and multi IPU configurations as we scale up to larger model sizes. Our work paves the way towards more efficient, non-standard models via AI training hardware beyond GPUs, and competitive large scale SNN models.

翻译：当前AI训练基础设施以单指令多数据（SIMD）与脉动阵列架构为主，例如图形处理单元（GPU）和张量处理单元（TPU），这些架构在加速并行工作负载与密集向量矩阵乘法方面表现优异。然而，利用稀疏性与循环特性的潜在更高效神经网络模型无法充分发挥SIMD处理器的全部性能，因而与Transformer、CNN等当前主流并行架构相比处于显著劣势，阻碍了可持续发展的人工智能发展路径。为突破这一局限，本文探索了在具有分布式本地内存的大规模并行多指令多数据（MIMD）架构上训练稀疏与循环模型。我们针对具有二值稀疏激活特性的脑启发式脉冲神经网络（SNN）类模型，实现了基于时间反向传播（BPTT）的训练流程。实验表明，与GPU相比，在MIMD处理器（智能处理单元IPU）上使用稀疏激活张量具有显著优势。在训练任务中，我们的结果显示：相比A100 GPU，吞吐量提升5至10倍；当激活稀疏度更高时，吞吐量提升可达38倍，且训练收敛速度与最终模型性能未出现显著下降。此外，当我们扩展至更大模型规模时，单IPU与多IPU配置均展现出极具前景的发展趋势。本工作为超越GPU的AI训练硬件支撑下更高效的非标准化模型（包括具有竞争力的规模化SNN模型）奠定了基础。