TPipe: Efficient Spiking Transformer Training with Time Parallelism and Asynchronous Pipeline

摘要

Spiking Transformers, which couple Spiking Neural Networks (SNNs) with Transformer architectures, promise low-energy inference and strong accuracy. However, their training is hindered by the intrinsic temporal dynamics of SNNs, which cause severe memory and synchronization bottlenecks. Through an in-depth analysis of the SpikeFormer architecture and training workload, we uncover local dependency in its training process, which enables time parallelism. We present TPipe, the first unified framework for efficient SpikeFormer training that integrates time, data, and pipeline parallelism. We formalize the Automated SpikeFormer Parallelism (ASP) problem to jointly optimize parallelism and placement strategies. We further formalize the SpikeFormer Scheduling (SS) problem, which eliminates pipeline bubbles via fine-grained, sub–micro-batch asynchronous scheduling. We prove that TPipe achieves bounded speed-up over synchronous baselines while preserving training correctness. Experiments on multiple SpikeFormer models and datasets show that TPipe attains 1.1–3.8× speed-up over default pipeline baselines while maintaining accuracy.