mlsanityData to 7 Oct 2026

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

arXiv 2604.26294 ↗First posted 29 Apr 2026cs.CL cs.DCcited by 0

On the map: Mixture-of-experts → Expert routing strategies

Abstract (arXiv)

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.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts. Builds on 12 of its 21 references that are in the map. First posted on arXiv 29 Apr 2026.

Builds on

12 of 21 references.

References · 12

  1. 2024The Llama 3 Herd of Modelscited by 14,192
  2. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  3. 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 621
  4. 2023Ring Attention with Blockwise Transformers for Near-Infinite Contextcited by 448
  5. 2022Reducing Activation Recomputation in Large Transformer Modelscited by 300
  6. 2023DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Modelscited by 239
  7. 2018Mesh-TensorFlow: Deep Learning for Supercomputerscited by 168
  8. 2021Sequence Parallelism: Long Sequence Training from System Perspectivecited by 125
  9. 2024USP: A Unified Sequence Parallelism Approach for Long Context Generative AIcited by 58
  10. 2023Striped Attention: Faster Ring Attention for Causal Transformerscited by 52
  11. 2025MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Corecited by 14
  12. 2023Rail-only: A Low-Cost High-Performance Network for Training LLMs with Trillion Parameterscited by 13

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Parallel Track Transformers: Enabling Fast GPU Inference with Reduced Synchronization
  2. 2026Scaling State-Space Models on Multiple GPUs with Tensor Parallelism
  3. 2025Shift Parallelism: Low-Latency, High-Throughput LLM Inference for Dynamic Workloads
  4. 2026Tensor Cache: Eviction-conditioned Associative Memory for Transformers
  5. 2024TPI-LLM: Serving 70B-scale LLMs Efficiently on Low-resource Edge Devices
  6. 2026SwiftFusion: Scalable Sequence Parallelism for Distributed Inference of Diffusion Transformers on GPUs
  7. 2026The Surprising Effectiveness of Shared Memory in Looped Transformers
  8. 2025End-to-End Transformer Acceleration Through Processing-in-Memory Architectures