Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

TL;DR

The paper introduces tensor and sequence parallelism (TSP), a unified parallel execution strategy that reduces memory usage in transformer training by combining weight and token sharding on a single device axis.

Contribution

TSP is a novel parallelism scheme that unifies tensor and sequence parallelism, reducing memory overhead and enabling efficient training of long-context models.

Findings

TSP reduces memory usage compared to traditional TP and SP.

TSP demonstrates competitive performance in benchmarks.

Theoretical analysis shows communication trade-offs of TSP.

Abstract

We present tensor and sequence parallelism (TSP), a parallel execution strategy that folds tensor parallelism and sequence parallelism onto a single device axis. In conventional multi-dimensional parallelism layouts, tensor parallelism (TP) shards model weights while sequence parallelism (SP) shards tokens, reducing per-device parameter or activation memory, respectively. Traditionally, each scheme is assigned its own mesh dimension. TSP instead assigns each rank both a weight shard and a sequence shard, reducing both parameter and activation memory along the same device axis. We implement this design with two runtime schedules. For attention, ranks iterate over broadcast parameter shards and reconstruct context through a sequence-wise key/value exchange. For gated MLPs, weight shards circulate in a ring while partial outputs accumulate locally. By sharding both weights and activations across the same devices, TSP trades additional communication volume for reduced memory overhead. We provide a theoretical communication and memory analysis, describe our implementation of TSP attention and gated MLP blocks, and benchmark TSP against TP, SP, and TP+SP. These results position TSP as a hardware-aware alternative for long-context and memory-constrained model training, and as a viable axis of parallelism in concert with existing parallelism schemes such as pipeline and expert parallelism for dense and mixture-of-expert models.