💬 NLP

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

This paper introduces Tensor and Sequence Parallelism (TSP), a novel execution strategy that folds weight and token sharding onto a single device axis to simultaneously reduce parameter and activation memory overhead, offering a hardware-efficient alternative for training and inference of long-context and memory-constrained transformer models.

Original authors: Vasu Shyam, Anna Golubeva, Quentin Anthony

Published 2026-04-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vasu Shyam, Anna Golubeva, Quentin Anthony

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive jigsaw puzzle with a group of friends, but you only have a very small table (your computer's memory) to work on. The puzzle is so big that no single person can hold all the pieces at once.

This paper introduces a new way for a team of computers (GPUs) to work together to train giant AI models, which are essentially these massive puzzles. The authors call their new strategy TSP (Tensor and Sequence Parallelism).

Here is the breakdown using simple analogies:

The Problem: Two Old Ways to Share the Work

To solve the puzzle, the team usually uses one of two old methods, but both have flaws:

The "Weight Split" Method (Tensor Parallelism):
Imagine the puzzle pieces are the "rules" of the game (the model's weights). In this method, you cut the rulebook in half. Person A holds the first half of the rules, and Person B holds the second half.
- The Good: You save space on the table because you aren't storing the whole rulebook twice.
- The Bad: If the puzzle has a long story (a long sequence of words), everyone still has to hold the entire story in their hands to play. If the story is huge, your hands (memory) get full, and you crash.
The "Story Split" Method (Sequence Parallelism):
Imagine the puzzle pieces are the story itself. In this method, Person A holds the first half of the story, and Person B holds the second half.
- The Good: You save space on the table because you aren't holding the whole story at once.
- The Bad: Everyone still has to memorize the entire rulebook. If the rulebook is huge, your brain (memory) gets full, and you crash.

The Old Hybrid: Usually, teams try to do both by using two separate groups of friends. One group splits the rules, and a different group splits the story. But this is inefficient because it uses up all your friends just to split the work, leaving no one left to help with other tasks (like Data Parallelism).

The Solution: The "Folded" Method (TSP)

The authors say: "Why use two separate groups? Let's fold the work onto a single axis."

In TSP, every single person in the group does both things at the same time:

They hold a slice of the rulebook (weights).
They hold a slice of the story (sequence).

The Analogy:
Imagine you are at a dinner party.

Old Way: You have a table where one person passes around the menu (weights) while everyone reads the whole book. Another table has people passing around the book (story) while everyone memorizes the whole menu.
TSP Way: Everyone at the table gets a small piece of the menu and a small piece of the story.

How They Make It Work (The Magic Tricks)

Since everyone has a tiny piece of the menu and a tiny piece of the story, they have to talk to each other a lot to finish the puzzle. The paper describes two clever ways they do this without getting overwhelmed:

For the "Story" parts (Attention):
Imagine the group needs to know the whole story to understand a specific sentence. Instead of everyone shouting the whole story at once, they take turns. One person broadcasts their piece of the menu to everyone. Then, everyone calculates their part of the story, and they quickly swap their story pieces (keys and values) to reconstruct the full context. It's like a relay race where they pass the baton (data) while running.
For the "Rules" parts (MLP):
Imagine the group needs to apply different rules to their story pieces. Instead of stopping to shout the rules, they pass the rulebook pages around in a circle (a ring). Person A does their math with Page 1, then passes Page 1 to Person B while Person B passes Page 2 to Person C. While the pages are moving, everyone is busy doing math. This keeps the "traffic" moving while the "work" gets done.

Why Is This Better?

The paper claims TSP is a "hardware-aware" solution, meaning it is designed specifically for how modern computer chips talk to each other.

Memory Savings: Because everyone holds a piece of the rules and a piece of the story, the memory required on each computer drops significantly. This allows the team to handle much longer stories (longer context) without running out of memory.
Speed: Even though they are passing more data back and forth (which sounds slower), they do it in a way that overlaps with their thinking. The "passing" happens while they are "thinking," so the total time doesn't get much longer.
Fitting in the Room: In a computer cluster, the fastest connection is usually between chips on the same machine (like people sitting at the same table). The slower connection is between different machines (people in different rooms).
- Old methods often forced the team to split across different rooms, slowing them down.
- TSP allows the whole "split" team to fit on a single machine (one table), keeping them in the fast lane.

The Results

The authors tested this on a massive cluster of 1,024 powerful GPUs (MI300X).

Memory: TSP used the least amount of memory at every test, especially when the stories were very long.
Speed: TSP was just as fast, or faster, than the old methods.
Scalability: As they added more computers to the group, TSP continued to perform well, whereas the old methods started to struggle with memory limits.

In short: TSP is a smarter way to organize a team of computers. Instead of splitting the "rules" and the "story" into separate groups, it combines them so every computer holds a little bit of both. This saves space, allows for longer stories, and keeps the team working efficiently on the same fast network.

1. Problem Statement

Training and inference of large language models (LLMs) face two primary memory bottlenecks:

Parameter Memory: The storage required for model weights, gradients, and optimizer states.
Activation Memory: The storage required for intermediate activations, which scales quadratically with sequence length ( $S^2$ ) in attention mechanisms.

Existing parallelism strategies address these issues in isolation or with trade-offs:

Tensor Parallelism (TP): Shards model weights across devices, reducing parameter memory but leaving activation memory fully replicated. It becomes inefficient at long sequence lengths due to high activation memory and communication volume.
Sequence Parallelism (SP): Shards the input sequence across devices, reducing activation memory but leaving model weights fully replicated. It is ineffective when parameter memory is the bottleneck.
Conventional TP+SP: Combines both by assigning them to orthogonal axes in a device mesh (e.g., a $T \times \Sigma$ grid). While effective, this consumes a large number of devices for model parallelism ( $T \times \Sigma$ ), leaving fewer devices for data parallelism. Furthermore, it often forces parallel groups to span slower inter-node links, underutilizing high-bandwidth intra-node interconnects (like NVLink or AMD Infinity Fabric).

The Core Challenge: How to simultaneously reduce both parameter and activation memory on a single device axis to maximize memory efficiency and free up devices for data parallelism, without incurring prohibitive communication costs.

2. Methodology: Tensor and Sequence Parallelism (TSP)

The authors propose TSP, a strategy that "folds" tensor and sequence parallelism onto a single device axis. Instead of assigning TP and SP to orthogonal dimensions, TSP assigns every rank (GPU) both a shard of the model weights and a shard of the sequence.

Key Design Components:

Unified Sharding: Each device holds $1/D$ of the model parameters and $1/D$ of the sequence tokens. This reduces both parameter-proportional memory and activation memory by a factor of $D$ on a single axis.
Runtime Schedules: To handle the fact that weights and activations are sharded differently, TSP employs specific communication schedules for different layers:
- Attention Layers: Uses a broadcast schedule.
  - Ranks iterate over weight shards (packed $W_Q, W_K, W_V, W_O$ ).
  - In each iteration, one rank broadcasts its weight shard to all others.
  - All ranks compute local projections on their sequence shards.
  - An All-Gather operation reconstructs the full Key/Value context across the sequence dimension to perform FlashAttention.
  - A Zigzag partitioning scheme is used to balance the causal attention workload (ensuring later tokens don't incur disproportionately more work).
- Gated MLP Layers: Uses a ring communication schedule.
  - Weight shards ( $W_1, W_2, W_3$ ) circulate point-to-point around the ring.
  - Each rank accumulates partial outputs locally.
  - This eliminates the need for an All-Reduce operation (standard in TP), as the sequence remains local and only weights move.
Hardware Awareness: The design is optimized for topologies like AMD's MI300X (Infinity Fabric), where collective bandwidth scales with the number of participating GPUs. TSP ensures that communication primitives (broadcast and ring) engage the full group or are overlapped with compute to maximize efficiency.

3. Key Contributions

Novel Parallelism Strategy: Introduction of TSP, which collapses two orthogonal parallelism axes into one, allowing simultaneous reduction of parameter and activation memory.
Theoretical Analysis: A rigorous mathematical model comparing TSP against DP, TP, SP, and TP+SP regarding:
- Memory: TSP is the only scheme that reduces both weight-proportional and activation memory by $1/D$ without requiring $D^2$ devices.
- Communication: While TSP introduces weight-movement communication (which TP+SP avoids), the authors prove that for long sequences or moderate batch sizes, the activation-proportional communication dominates, making TSP competitive.
- FLOPs: TSP maintains the same per-device computational efficiency as TP and SP.
Implementation & Overlap: Detailed implementation of TSP attention and MLP blocks with explicit communication/computation overlap strategies (pipelining weight transfers behind GEMM computations).
Topology Optimization: Demonstration that TSP allows the entire model replica to fit within a single high-bandwidth node (intra-node), avoiding the performance penalties of spanning slower inter-node links often required by 2D TP+SP meshes.

4. Results

The authors evaluated TSP on a cluster of up to 1,024 AMD MI300X GPUs (128 nodes) with a 7B parameter dense transformer model.

Memory Efficiency:
- TSP consistently achieved the lowest peak memory across all tested sequence lengths (16K to 128K tokens).
- At short contexts, TSP matched TP (dominated by parameter memory).
- At long contexts, TSP significantly outperformed TP and TP+SP (dominated by activation memory), preventing Out-Of-Memory (OOM) errors where other strategies failed.
Throughput:
- Despite higher theoretical communication volume (due to weight movement), TSP maintained competitive throughput with TP+SP.
- The extra communication was successfully hidden behind computation via pipelining.
- Throughput advantages grew with increasing parallelism degrees ( $D$ ) and sequence lengths.
Scalability:
- TSP allowed for larger micro-batches due to reduced memory footprint, further improving throughput.
- It enabled fitting model replicas entirely within a single node (8 GPUs), freeing up ranks for Data Parallelism (DP) in larger clusters.

5. Significance

Hardware-Aware Scaling: TSP provides a practical solution for modern HPC clusters where intra-node bandwidth is high but inter-node bandwidth is a bottleneck. It maximizes the utility of expensive intra-node interconnects.
Long-Context Viability: As models move toward longer context windows (100k+ tokens), activation memory becomes the primary bottleneck. TSP is uniquely positioned to handle this without the massive device overhead of 2D parallelism.
Flexibility: TSP can be composed orthogonally with other parallelism schemes (Pipeline Parallelism, Expert Parallelism, Data Parallelism), offering a flexible "axis" for system designers to optimize for specific hardware topologies.
Cost Efficiency: By reducing the number of devices needed for model parallelism, TSP allows more devices to be used for data parallelism, potentially reducing the total cluster size required for a given training job or increasing the effective batch size.

In conclusion, TSP represents a significant evolution in distributed training strategies, trading a manageable increase in communication volume for substantial gains in memory efficiency and hardware utilization, making it a critical tool for training next-generation long-context LLMs.