Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

Ravi Ghadia; Maksim Abraham; Sergei Vorobyov; Max Ryabinin

cs.LGFeb 24, 2026

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

Authors: Ravi Ghadia, Maksim Abraham, Sergei Vorobyov, Max Ryabinin

Organizations: Together AI

Abstract

Efficiently processing long sequences with Transformer models usually requires splitting the computations across accelerators via context parallelism. The dominant approaches in this family of methods, such as Ring Attention or DeepSpeed Ulysses, enable scaling over the context dimension but do not focus on memory efficiency, which limits the sequence lengths they can support. More advanced techniques, such as Fully Pipelined Distributed Transformer or activation offloading, can further extend the possible context length at the cost of training throughput. In this paper, we present UPipe, a simple yet effective context parallelism technique that performs fine-grained chunking at the attention head level. This technique significantly reduces the activation memory usage of self-attention, breaking the activation memory barrier and unlocking much longer context lengths. Our approach reduces intermediate tensor memory usage in the attention layer by as much as 87.5 $\%$ for 32B Transformers, while matching previous context parallelism techniques in terms of training speed. UPipe can support the context length of 5M tokens when training Llama3-8B on a single 8 $\times$ H100 node, improving upon prior methods by over 25 $\%$ .

Figures & tables

Figure 1: Comparison of context parallelism approaches on long-sequence training for Llama3-8B using 8

\times

H100s. UPipe provides maximum efficiency, resulting in a longer maximum context length (5M tokens) while retaining throughput.

Figure 2: Memory usage breakdown when training Llama 3-8B with a sequence length of 3M tokens across 8 H100 GPUs. FSDP denotes Fully Sharded Data Parallel, AC denotes Activation Checkpointing, AO denotes AC with offloading, OOM denotes Out of Memory.

Panel of Table 1 — Table 1: Theoretical peak memory usage breakdown across different stages of the forward pass for a Transformer model. All floating point tensors use BFloat16 precision by default, except cross-entropy loss, which uses fp32 precision.

Method	Before attn block	During inp_all_to_all	During attn kernel	During out_all_to_all
Ulysses	$L$	$L+(\gamma+1)$	$L+(\gamma+1)$	$L+2$
Ulysses + offloading	$1$	$\gamma+2$	$\gamma+2$	$3$
FPDT	$\frac{1}{\pi}$	$\frac{\gamma+2}{\pi}$	$\frac{2\gamma+1}{\pi}$	$\frac{2}{\pi}$
Untied Ulysses	$1$	$2+\frac{\gamma}{\nu}$	$2+\frac{\gamma}{\nu}$	$2+\frac{1}{\nu}$

Table 2: Peak activation memory within the forward attention block under GQA, expressed in units of one local hidden-state shard (

\frac{S}{C}\cdot d_{model}

bf16 elements, i.e.,

2\cdot\frac{S}{C}\cdot d_{model}

bytes).

\pi

represents the number of chunks in FPDT,

\nu

represents the number of chunks in UPipe. UPipe consumes

\nu

times less intermediate (QKV + all-to-all) activation memory than Ulysses with activation offloading. FPDT has lower memory usage due to an arbitrary chunk size, but suffers from performance degradation.

Figure 3: Illustration of (a) DeepSpeed-Ulysses and (b) UPipe designs. UPipe processes attention in a headwise untied manner, so that in each stage, attention is performed only on a subset of heads. This allows memory reuse across different stages, significantly reducing the peak memory usage due to attention activations. HBM (High-Bandwidth Memory) usage illustrates memory utilization of the intermediate buffers and omits other components for brevity.

Figure 4: Illustration of UPipe’s GQA scheduling algorithm. We communicate as many unique key/value heads as possible, along with the corresponding queries in stage 0. In the subsequent stages, we only communicate the next queries of the corresponding groups, reusing the key/value tensors from stage 0 until stage $G$ , where $G$ is the group size.

	Method	128K	256K	512K	1M	2M	3M	4M	5M
Llama 3-8B	Native PyTorch	1373.87	845.99	474.30	249.85	OOM	–	–	–
	Ring	2064.90	1387.67	841.05	458.51	237.99	159.96	OOM	–
	Ulysses	2320.47	1503.80	878.63	475.33	246.05	162.41	OOM	–
	FPDT	1171.68	884.75	621.20	382.42	219.53	153.48	119.76	–
	UPipe	2281.05	1487.29	867.17	472.53	246.07	166.32	125.56	98.25
Qwen3-32B	Native PyTorch	127.03	112.20	91.39	OOM	–	–	–	–

Table 3: Throughput comparison (tokens/second/GPU) for Llama 3-8B (8

\times

H100s) and Qwen3-32B (16

\times

H100s) across varying sequence lengths. On a single 8

\times

H100 node, our method scales Llama 3-8B to a 5M-token sequence length, reaching a 25

\%

improvement over the previous state-of-the-art. OOM : Out of Memory. Note : FPDT execution fails at lengths

>

4M.

Figure 5: Llama 3-8B: Comparison of peak GPU memory usage and throughput (normalized w.r.t. USP-Hybrid) between UPipe and USP-Hybrid at different sequence lengths on 16

\times

H100s.

Figure 6: Ablation analysis of UPipe’s head-chunk size

U

when training Llama 3-8B with a context size of 512K on

C=

4 GPUs.

	Method	xx 128K	256K	512K	1M
TPS	Ulysses	xx 2417.79	1516.51	880.27	476.45
	UPipe	xx 2377.98	1515.59	878.73	470.34
	UPipe-SAO	xx 2445.10	1512.69	872.57	471.07
GiB	Ulysses	xx 62.94	63.01	63.15	63.43
	UPipe	xx 53.24	53.31	53.45	53.73
	UPipe-SAO	xx 62.39	62.47	62.61	62.89

Table 4: Comparison of throughput (TPS) and memory usage (GiB) across Ulysses, UPipe, and UPipe with SAO (Strided Activation Offloading). SAO narrows the performance gap between Ulysses and UPipe due to less frequent activation offloading to CPU.

Method	Throughput	Memory
Ulysses	864.16	27.34
Ulysses+SAC	1066.01	27.85
UPipe	850.97	25.24
UPipe+SAC	1056.77	25.75

Table 5: Throughput (in tokens/second) and memory usage (GiB) of Ulysses and UPipe with and without Selective Activation Checkpointing at a sequence length of 512K for Llama 3-8B.

	Batch size	1	2	4	8	10
TPS	Ulysses	2458.28	2485.78	2543.43	2564.38	2464.10
TPS	UPipe	2405.30	2432.11	2498.80	2518.56	2526.77
GiB	Ulysses	22.53	27.25	38.13	59.99	71.13
GiB	UPipe	22.53	27.03	36.56	56.79	67.11

Table 6: Comparison of throughput (TPS) and memory usage (GiB) for Ulysses and UPipe across different batch sizes.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Before Bwd Attn	During out_all_to_all	During Bwd Attn Kernel	During inp_all_to_all
Ulysses	$L+2$	$L+3$	$L+2\gamma+2$	$L+\gamma+1$
Ulysses + offloading	$3$	$4$	$2\gamma+3$	$\gamma+2$
FPDT	$\frac{1}{\pi}$	$\frac{3}{\pi}$	$\frac{2\gamma+4}{\pi}$	$\frac{\gamma+2}{\pi}$
Untied Ulysses	$3$	$3+\frac{1}{\nu}$	$\max\!\left(3+\frac{2\gamma}{\nu},\;2+\frac{2\gamma}{\nu}+\frac{2(\nu-1)}{\nu}\right)$	$2+\frac{\gamma}{\nu}+\frac{2(\nu-1)}{\nu}$

Appendix

Table 7: Peak activation memory during the backward pass of the attention stage under GQA, expressed in units of one local hidden-state shard (

\frac{S}{C}\cdot d_{model}

bf16 elements, i.e.,

2\cdot\frac{S}{C}\cdot d_{model}

bytes).

\pi

represents the number of chunks in FPDT,

\nu

represents the number of chunks in UPipe.

Model	Method	128K	256K	512K	1M	2M	3M	4M	5M
Llama 3-8B	Native PyTorch	25.32	31.40	43.55	67.86	OOM	–	–	–
	Ring	21.32	23.40	27.58	35.86	52.49	69.11	OOM	–
	Ulysses	21.26	23.02	26.80	34.35	49.49	64.55	OOM	–
	FPDT	21.73	22.50	24.03	27.09	35.17	43.35	51.42	–
	UPipe	21.10	22.30	24.70	29.90	40.50	51.10	61.70	72.30
Qwen3-32B	Native PyTorch	45.81	53.69	69.47	OOM	–	–	–	–

Appendix

Table 8: Memory comparison (GiB) for Llama 3-8B (8

\times

H100s) and Qwen3-32B (16

\times

H100s) across varying sequence lengths. OOM denotes Out of Memory. FPDT execution fails at lengths

>

4M. The lowest memory usage at a given sequence length is in bold.

Method	Stage	128K	256K	512K	1M	2M	3M
Ulysses	All-to-All	0.40	0.90	1.68	4.93	16.30	42.21
	FA3-Fwd	1.58	6.35	25.71	103.49	421.67	995.92
	FA3-Bwd	2.40	9.13	36.74	146.86	588.73	1324.71
	Other	3.03	5.33	10.08	19.78	41.30	56.31
	Total	7.40	21.72	74.21	275.06	1068.00	2419.14
UPipe	All-to-All	0.46	1.10	2.43	5.52	17.12	34.34

Appendix

Table 9: Runtime comparison (in seconds) for Llama 3-8B on 8

\times

H100 between DeepSpeed-Ulysses and UPipe across sequence lengths. FA3-Fwd/Bwd: Total Flash Attention-3 forward and backward kernel time. Total refers to the total time for a single training step.

Method	100	200	300	400	500	600	700	800	900	1000
Ulysses	6.83	6.58	5.98	5.62	5.22	4.88	4.94	5.11	4.68	4.47
UPipe	6.85	6.54	5.98	5.62	5.22	4.90	4.95	5.11	4.69	4.47
%Diff	0.29%	-0.61%	0.00%	0.00%	0.00%	0.41%	0.20%	0.00%	0.21%	0.00%

Appendix

Table 10: Loss Convergence: Comparison of training loss between UPipe and Ulysses when training a Llama 3-8B model on C4 dataset with sequence length of 128K tokens over 1000 steps.

Memory-Efficient Attention

See all 24 papers

Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference
Apr 29, 2026Vasu Shyam, Anna Golubeva, Quentin AnthonyTransformer Inference Transformer
Zyphra
Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules
Oct 13, 2026Sarah Pan, Ali BehrouzLong-Context Language Model Inference Memory-Augmented Language Models
Massachusetts Institute of Technology · Cornell University
Training-Inference Consistent Segmented Execution for Long-Context LLMs
May 12, 2026Xianpeng Shang, Jiang Li, Zehua Duo +2Self-Attention Long-Context Language Modeling
College of Computer Science, Inner Mongolia University, Hohhot 010021, China · National & Local Joint Engineering Research Center of Intelligent Information Processing Technology for Mongolian, Hohhot 010021, China · Inner Mongolia Key Laboratory of Multilingual Artificial Intelligence Technology, Hohhot 010021, China +1

Apr 29, 2026J/K move · Enter open · S save

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

Vasu Shyam, Anna Golubeva, Quentin Anthony

Zyphra

Transformer Inference Transformer Tensor Parallelism Memory-Efficient Inference

Open paper AskPDF 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.

Abstract

Memory-Efficient Attention

Explore similar work

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

Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules

Training-Inference Consistent Segmented Execution for Long-Context LLMs