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

• URL: http://arxiv.org/abs/2604.26294v1
• Date: Wed, 29 Apr 2026 04:44:58 GMT
• Title: Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference
• Abstract summary: tensor and sequence parallelism (TSP) is a parallel execution strategy that folds tensor parallelism and sequence parallelism onto a single device axis.<n>We implement this design with two runtime schedules. For attention, ranks over broadcast parameter shards and context through a sequence-wise key/value exchange.<n>We provide a theoretical communication and memory analysis, describe our implementation of attention and gated blocks, and benchmark TSP against TP, SP, and TP+SP.
• Score: 7.139526063136717
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Demystifying Pipeline Parallelism: First Theory for PipeDream [53.657104889705856]
  In this paper we introduce Randomized PipeDream (PD) as a stale block-SGD abstraction that yields clean non convergence for a PD-style method.<n>We prove that the delay induced by steady-state PD grows as $S2 - S/2 + O(1)$ for $S$ stages, so the stale-read contribution scales as $(2 S4)$, equivalently $(S4/K)$ in the tuned-rate form.
  arXiv  Detail & Related papers  (2026-06-02T11:14:57Z)
• Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior [107.2098567818173]
  Latent Recurrent Transformer (LRT) is a lightweight augmentation of autoregressive transformers.<n>LRT reuses a high-level source-layer hidden state from the previous token as recurrent memory for the next token.
  arXiv  Detail & Related papers  (2026-05-26T10:10:26Z)
• Parallel Recursive LSTM [26.883902799977964]
  We introduce the Parallel Recursive LSTM, a hierarchical recurrent architecture that replaces left-to-right recurrence with nonlinear state composition over a balanced computation tree.<n>PR-LSTM achieves strong sequence-length generalization on formal-language benchmarks, solving more tasks than standard RNN, LSTM, and Transformer baselines, while avoiding the quadratic scaling of attention.
  arXiv  Detail & Related papers  (2026-05-16T18:28:59Z)
• PRISM: Parallel Residual Iterative Sequence Model [52.26239951489612]
  We propose PRISM (Parallel Residual Iterative Sequence Model) to resolve this tension.<n>PRISM introduces a solver-inspired inductive bias that captures key structural properties of multi-step refinement in a parallelizable form.<n>We prove that this formulation achieves Rank-$L$ accumulation, structurally expanding the update manifold beyond the single-step Rank-$1$ bottleneck.
  arXiv  Detail & Related papers  (2026-02-11T12:39:41Z)
• Parallel Track Transformers: Enabling Fast GPU Inference with Reduced Synchronization [19.97521786735984]
  Parallel Track (PT) Transformer is a novel architectural paradigm that restructures to minimize cross-device dependencies.<n>We report consistent improvements in serving efficiency, including up to 15-30% reduced time to first token, 2-12% reduced time per output token, and up to 31.90% increased throughput in both settings.
  arXiv  Detail & Related papers  (2026-02-07T01:42:20Z)
• db-SP: Accelerating Sparse Attention for Visual Generative Models with Dual-Balanced Sequence Parallelism [14.406306253079515]
  Scaling Diffusion Transformer (DiT) inference via sequence parallelism is critical for reducing latency in visual generation.<n>We formalize a sparse imbalance ratio to quantify the imbalance, and propose db-SP, a sparsity-aware sequence parallelism technique.<n>We show that db-SP delivers an end-to-end speedup of 1.25x and an attention-specific speedup of 1.40x over state-of-the-art sequence parallel methods.
  arXiv  Detail & Related papers  (2025-11-28T11:55:46Z)
• GSPN-2: Efficient Parallel Sequence Modeling [101.33780567131716]
  Generalized Spatial Propagation Network (GSPN) addresses this by replacing quadratic self-attention with a line-scan propagation scheme.<n>GSPN-2 establishes a new efficiency frontier for modeling global spatial context in vision applications.
  arXiv  Detail & Related papers  (2025-11-28T07:26:45Z)
• PT$^2$-LLM: Post-Training Ternarization for Large Language Models [52.4629647715623]
  Large Language Models (LLMs) have shown impressive capabilities across diverse tasks, but their large memory and compute demands hinder deployment.<n>We propose PT$2$-LLM, a post-training ternarization framework tailored for LLMs.<n>At its core is an Asymmetric Ternary Quantizer equipped with a two-stage refinement pipeline.
  arXiv  Detail & Related papers  (2025-09-27T03:01:48Z)
• HelixPipe: Efficient Distributed Training of Long Sequence Transformers with Attention Parallel Pipeline Parallelism [14.067070576474086]
  As transformer sequence lengths grow, existing pipeline parallelisms incur suboptimal performance due to the quadratic attention computation and the substantial memory overhead.<n>We propose HelixPipe, a novel pipeline parallelism for long sequence transformer training.<n>It introduces attention parallel partition, which schedules attention computations of different micro batches across different pipeline stages in parallel, reducing pipeline bubbles.<n>It employs a two-fold first-in-last-out micro batch schedule to balance memory usage and overlap communication with fragmentation.
  arXiv  Detail & Related papers  (2025-07-01T03:11:18Z)
• xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism [5.704297874096985]
  Diffusion models are pivotal for generating high-quality images and videos. This paper introduces xDiT, a comprehensive parallel inference engine for DiTs. Notably, we are the first to demonstrate DiTs scalability on Ethernet-connected GPU clusters.
  arXiv  Detail & Related papers  (2024-11-04T01:40:38Z)
• USP: A Unified Sequence Parallelism Approach for Long Context Generative AI [1.973144426163543]
  Sequence parallelism (SP) is becoming key to unlocking the long-context capabilities of generative AI models. This paper investigates the state-of-the-art SP approaches, i.e. DeepSpeed-Ulysses and Ring-Attention, and proposes a unified SP approach. We achieved 47% MFU on two 8xA800 nodes using SP for the LLAMA3-8B model training using sequence length 208K.
  arXiv  Detail & Related papers  (2024-05-13T13:08:02Z)
• Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning [54.99749970495241]
  Alpa automates model-parallel training of large deep learning (DL) models. Alpa generates execution plans that unify data, operator, and pipeline parallelism. Unlike specialized systems, Alpa also generalizes to models with heterogeneous architectures and models without manually-designed plans.
  arXiv  Detail & Related papers  (2022-01-28T10:13:35Z)
• Parallel Training of Deep Networks with Local Updates [84.30918922367442]
  Local parallelism is a framework which parallelizes training of individual layers in deep networks by replacing global backpropagation with truncated layer-wise backpropagation. We show results in both vision and language domains across a diverse set of architectures, and find that local parallelism is particularly effective in the high-compute regime.
  arXiv  Detail & Related papers  (2020-12-07T16:38:45Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.