[Distributed w/ TorchTitan] Breaking Barriers: Training Long Context LLMs with 1M Sequence Length in PyTorch Using Context Parallel

XWu: <p>with <a href="https://github.com/fegin" rel="noopener nofollow ugc">Chien-Chin Huang</a> (<a class="mention" href="/u/fegin">@fegin</a>), <a href="https://discuss.pytorch.org/u/lessw2020/summary">Less Wright</a> (<a class="mention" href="/u/lessw2020">@lessw2020</a>), <a href="https://discuss.pytorch.org/u/wanchaol/summary">Tianyu Liu</a> (<a class="mention" href="/u/tianyu">@tianyu</a>), <a href="https://github.com/wconstab" rel="noopener nofollow ugc">Will Constable</a> (<a class="mention" href="/u/wconstab">@wconstab</a>), <a href="https://github.com/gnadathur" rel="noopener nofollow ugc">Gokul Nadathur</a> (<a class="mention" href="/u/gnadathur">@gnadathur</a>)</p>
<h1><a name="p-462811-tldr-1" class="anchor" href="#p-462811-tldr-1"></a>TL;DR</h1>
<ul>
<li>
<p>We implemented <a href="https://arxiv.org/abs/2411.01783" rel="noopener nofollow ugc">pass-KV Ring Attention</a> for Context Parallel in PyTorch</p>
</li>
<li>
<p>We integrated it in torchtitan and verified its effectiveness as well as composability with other native techniques in PyTorch such as <a href="https://pytorch.org/docs/main/distributed.fsdp.fully_shard.html" rel="noopener nofollow ugc">FSDP</a> and <a href="https://pytorch.org/docs/main/torch.compiler.html" rel="noopener nofollow ugc">torch.compile</a></p>
<ul>
<li>
<p>Sequence length scaling to 1M on <a href="https://ai.meta.com/blog/meta-llama-3-1/" rel="noopener nofollow ugc">Llama3-8B</a> model using 32 H100 GPUs</p>
</li>
<li>
<p>Matching loss curves with existing parallelisms</p>
</li>
<li>
<p>Compatible with torch.compile. The compile mode provides an average 10% TPS gain on a set of benchmarks</p>
</li>
</ul>
</li>
</ul>
<h1><a name="p-462811-distributed-with-torchtitan-series-2" class="anchor" href="#p-462811-distributed-with-torchtitan-series-2"></a>Distributed with TorchTitan Series</h1>
<p>The GitHub repository <a href="https://github.com/pytorch/torchtitan" rel="noopener nofollow ugc">torchtitan</a> is a proof of concept for large-scale LLM training using native PyTorch, designed to be easy to understand, use, and extend for different training purposes, supporting multi-dimensional parallelisms with modular components. In this series of topics, we introduce the latest PyTorch features for distributed training enabled in Torchtitan.</p>
<ul>
<li>
<p>Topic 1: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323">Enabling Float8 All-Gather in FSDP2</a></p>
</li>
<li>
<p>Topic 2: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487">Introducing Async Tensor Parallelism in PyTorch</a></p>
</li>
<li>
<p>Topic 3: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250">Optimizing Checkpointing Efficiency with PyTorch DCP</a></p>
</li>
<li>
<p>Topic 4: <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420">Training with Zero-Bubble Pipeline Parallelism</a></p>
</li>
<li>
<p>→ <strong>Topic 5:</strong> Breaking Barriers: Training Long Context LLMs with 1M Sequence Length in PyTorch Using Context Parallel</p>
</li>
</ul>
<h1><a name="p-462811-context-parallel-3" class="anchor" href="#p-462811-context-parallel-3"></a>Context Parallel</h1>
<p><a href="https://openreview.net/forum?id=WsRHpHH4s0" rel="noopener nofollow ugc">Context Parallel</a> is an approach used in LLM to reduce peak activation size by sharding the long input sequence across multiple devices. It breaks the constraint on input sequence length resulting from peak memory usage on storing activations in Transformer blocks.</p>
<p>The core of Context Parallel is Ring Attention, a novel parallel implementation of the Attention layer. Ring Attention shuffles the KV shards and calculates the partial attention scores, repeats until all KV shards have been used on each device. This shard shuffling can be done via either multiple p2p send/recv between devices, or one allgather. This sequence length scaling is claimed to be zero-overhead if the shard rotation among devices overlaps well with the attention computation over local QKV shards.</p>
<p>The overview of Ring Attention algorithm is below:</p>
<pre><code class="lang-auto"># Ring Attention algorithm
N: number of devices where QKV is sharded.
Q_i; K_i; V_i: the QKV shard placed on device i

# on device i
for j in 1..N:
    Collect K_j and V_j from other device
    Compute Attention(Q_i, K_j, V_j)
</code></pre>
<h1><a name="p-462811-enabling-context-parallel-4" class="anchor" href="#p-462811-enabling-context-parallel-4"></a>Enabling Context Parallel</h1>
<p>It is actually pretty easy to automatically replace Attention layers with Ring Attention and utilize Context Parallel in users’ model code. We provide a public API <code>context_parallel</code> which takes in several arguments regarding how to parallelize <a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html" rel="noopener nofollow ugc">SDPA</a> and returns a Python context within which the SDPA function will be replaced with an equivalent Ring Attention.</p>
<p>Enabling Context Parallel in model code is quite simple and only takes two steps:</p>
<ol>
<li>Create a Context Parallel context for each training/inference step</li>
<li>Wrap the forward/backward pass using the context</li>
</ol>
<pre><code class="lang-auto">while step &lt; job_config.training.steps:  # training loop
	step += 1

	# Step 1:
	# create a context parallel python context for current training step
	context_parallel_ctx = context_parallel(...)

	# Step 2:
	# wrap training/inference code in the context
	with train_context(context_parallel_ctx):  # enable Context Parallel
		# no changes to model code needed!!
		pred = model(input_ids)
		loss = loss_fn(pred, labels)
		loss.backward()
</code></pre>
<p>Now let’s dive into an end-to-end example of adopting Context Parallel in Long-context LLMs training in PyTorch.</p>
<h2><a name="p-462811-long-context-training-in-torchtitan-5" class="anchor" href="#p-462811-long-context-training-in-torchtitan-5"></a>Long-context Training in Torchtitan</h2>
<p>We enabled Context Parallel in torchtitan to verify the effectiveness and composability of our implementation and showcase how Context Parallel can be easily enabled in user code.</p>
<p>First of all, just like other parallelisms adopted in torchtitan, there’s no change needed for the model code. All PyTorch SDPA function calls will be replaced with Ring Attention automatically.</p>
<pre><code class="lang-auto"># torchtitan/models/llama/model.py
# the model code remains untouched
class Attention(nn.Module):
	def forward(self, x: torch.Tensor, ...):
		...
		output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)  # SDPA
		...

</code></pre>
<p>Now, a number of GPUs need to be allocated to the Context Parallel device mesh dimensions.</p>
<pre><code class="lang-auto"># torchtitan/parallelisms/parallel_dims.py
# create Context Parallel device mesh from parsed config file
world_mesh = init_device_mesh(
	device_type,
	mesh_shape=...,
	mesh_dim_names=(..., "cp", ...),
)
cp_mesh = world_mesh["cp"]

# create "dp_shard_cp" mesh if composing with FSDP/HSDP
if DP:
	world_mesh["dp_shard", "cp"]._flatten(mesh_dim_name="dp_shard_cp")
else:
	world_mesh["cp"]._flatten(mesh_dim_name="dp_shard_cp")  # equivalent
</code></pre>
<p>Leverage the FSDP API to shard module parameters and install hooks that all-gather parameters before forward pass and reduce-scatter to sync gradients.</p>
<pre><code class="lang-auto"># torchtitan/parallelisms/parallelize_llama.py
# parallelize the model through FSDP fully_shard API
dp_mesh_dim_names = []

if HSDP:
	dp_mesh_dim_names.append("dp_replicate")

dp_mesh_dim_names.append("dp_shard_cp")
apply_fsdp(
	model,
	world_mesh[tuple(dp_mesh_dim_names)],
	cpu_offload=job_config.training.enable_cpu_offload,
	...
)
</code></pre>
<p>The last step is creating the python context for Context Parallel every step and putting the model forward and backward in the context so that the input tensor will be automatically sharded and SDPA will be replaced with Ring Attention.</p>
<pre><code class="lang-auto"># train.py
# enable context parallel in model training code
def main(job_config: JobConfig):
	...  # init

	while step &lt; job_config.training.steps:  # training loop
		step += 1
		...  # data loading, etc..
		# create a context parallel python context for current training step
		context_parallel_ctx = context_parallel(
			cp_mesh=world_mesh["cp"],
			cp_buffers=[input_ids, labels, ...],
			cp_seq_dims=[1, 1, ...],  # shard on seq dimension
			cp_no_restore_buffers={input_ids, labels},  # don't restore
			cp_rotate_method=job_config.experimental.context_parallel_rotate_method,  # shard rotation
		)

	with train_context(context_parallel_ctx):  # enable Context Parallel
		pred = model(input_ids)
		loss = loss_fn(pred, labels)
		del pred
		loss.backward()
</code></pre>
<p>With the above changes in torchtitan, we can simply turn on Context Parallel in torchtitan by modifying the “context_parallel_degree” field in the .toml config file. For example, we want to use Ring Attention and allocate 8 GPUs for every Context Parallel group when running the llama3-70B model in torchtian. We can do the following:</p>
<pre><code class="lang-auto"># file: llama3_70b.toml

[experimental]
context_parallel_degree = 8  # use 8 GPUs for each Context Parallel group
</code></pre>
<h2><a name="p-462811-composing-context-parallel-with-data-parallel-6" class="anchor" href="#p-462811-composing-context-parallel-with-data-parallel-6"></a>Composing Context Parallel With Data Parallel</h2>
<p>If users want to compose Context Parallel with Data Parallel (FSDP or HSDP) in PyTorch, there’s one more step: users need to create a joint DeviceMesh of Data Parallel dimension and Context Parallel dimension. Data Parallel dimension refers to the data parallel groups for FSDP and the data parallel groups that would actually shard model weights for HSDP.</p>
<pre><code class="lang-auto"># example: create a device mesh from HSDP and Context Parallel mesh dimensions
from torch.distributed.device_mesh import init_device_mesh 

# organize 8-GPU into 3 device mesh dimensions
world_mesh = init_device_mesh(
	device_type,
	mesh_shape=[[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
	mesh_dim_names=("dp_replicate", "dp_shard", "cp"),
)
hsdp_mesh = world_mesh["dp_replicate", "dp_shard"]
cp_mesh = world_mesh["cp"]

# flatten "dp_shard" and "cp" into the "dp_shard_cp" mesh, over which the
# weight all-gather and grads sync reduce-scatter in FSDP/HSDP is performed
world_mesh["dp_shard", "cp"]._flatten(mesh_dim_name="dp_shard_cp")
</code></pre>
<h1><a name="p-462811-implementation-details-7" class="anchor" href="#p-462811-implementation-details-7"></a>Implementation Details</h1>
<h2><a name="p-462811-ring-attention-8" class="anchor" href="#p-462811-ring-attention-8"></a>Ring Attention</h2>
<p>We implemented Ring Attention in PyTorch and provide <a href="https://github.com/pytorch/pytorch/blob/bb5e439f2d8a46172b8b7d2fdb7609822b9a97b1/torch/distributed/tensor/experimental/_attention.py#L1210-L1275" rel="noopener nofollow ugc">Context Parallel</a> as a Python context where the user-specified tensor will be automatically sharded over the sequence-dim within the Context Parallel group and all ATEN scaled dot product attention op (at this moment we only support _scaled_dot_product_flash_attention and _scaled_dot_product_efficient_attention) will be replaced with our ring attention handler during op dispatching.</p>
<p>The implementation consists of several submodules:</p>
<ul>
<li>Tensor sharding</li>
<li>Attention op dispatching</li>
<li>Shard rotation</li>
<li>Masking and load balancer</li>
<li>SDPA merger</li>
</ul>
<h3><a name="p-462811-tensor-sharding-9" class="anchor" href="#p-462811-tensor-sharding-9"></a>Tensor Sharding</h3>
<p>This module is responsible for sharding the Input tensor to attention layers as well as labels and embeddings if present, within the Context Parallel group. Currently we adopt two sharding strategies: Sequential sharder and Round Robin sharder.</p>
<p>Sequential sharder evenly splits QKV into N contiguous blocks where N is the size of the Context Parallel group and shards one block to each rank.</p>
<p>Round Robin sharder is designed for Causal Masking. When Causal Masking is used, only the lower triangle part of QKV is needed for computation. This sharder evenly splits QKV into <em>2N</em> blocks and shard the currently shortest block and longest block to rank <em>0, 1, …, N-1</em> (i.e. shard <em>0</em> and <em>2N-1</em> go to rank <em>0</em>, shard <em>1</em> and <em>2N-2</em> go to rank <em>1</em>, etc).</p>
<h3><a name="p-462811-attention-op-dispatching-10" class="anchor" href="#p-462811-attention-op-dispatching-10"></a>Attention Op Dispatching</h3>
<p>The replacement of PyTorch scaled dot product attention (SDPA) with Ring Attention takes two steps:</p>
<ol>
<li><a href="https://github.com/pytorch/pytorch/blob/bb5e439f2d8a46172b8b7d2fdb7609822b9a97b1/torch/distributed/tensor/experimental/_attention.py#L904-L910" rel="noopener nofollow ugc">Install hooks</a> to SDPA function to convert input from torch.Tensor to DTensor</li>
<li><a href="https://github.com/pytorch/pytorch/blob/bb5e439f2d8a46172b8b7d2fdb7609822b9a97b1/torch/distributed/tensor/experimental/_attention.py#L919-L920" rel="noopener nofollow ugc">Replace ATEN SDPA</a> ops with Ring Attention</li>
</ol>
<p>With the above change, all SDPA function calls within the Context Parallel Python context are seamlessly replaced with a Ring Attention implementation.</p>
<h3><a name="p-462811-shard-rotation-11" class="anchor" href="#p-462811-shard-rotation-11"></a>Shard Rotation</h3>
<p>Rotating QKV shards within the Context Parallel group between local SDPA calls instead of applying a single SDPA to the whole input batch is the key change in Ring Attention. How well this data communication can be overlapped with local SDPA computation dominates the scaling overhead. If the communication can be well hidden, the scaling will be zero-overhead.</p>
<p>We offer two strategies of rotating the QKV shards within the group: ring rotation using all-to-all collective between each local SDPA, or collect all shards using all-gather before launching the first local SDPA.</p>
<ul>
<li>Ring rotation using all-to-all between local SDPAs: ring Attention starts at a state where each Context Parallel rank holds a shard of QKV tensor. The first rotation happens when the first local SDPA is launched, and each rank passes its local shard to the next rank in a ring style: rank i passes to rank i+1 and the last rank passes to rank 0. Ideally this rotation is completed before the next local SDPA is about to launch. This iteration repeats until a shard has been passed throughout the ring.</li>
</ul>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXetvnwopNJtD7oNDE4ihGeZXaXMDVVZniSgGm_zF1amD7yraHwOFe0e6MSWqh5zBgQrN6E8cVj2K8V4u6e-4Xuh1D75pySvjdWp373OSC94EW5O0xtuMEWmpYB7EzNdWknf9ioIvA?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" width="624" height="128" role="presentation"></strong><br>
<em>Figure 1. Ideal Comp-Comm overlapping in all-to-all based Context Parallel</em></p>
<ul>
<li>All-gather all shards ahead of the first local SDPA: each rank still only holds one QKV shard when Ring Attention starts. However, each rank calls all-gather to collect the shards from other ranks. Then each rank first launches local SDPA on its own local QKV shard, selects the next shard from the all-gathered shards for its next local SDPA, and repeats until all shards have been used.</li>
</ul>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXcd78FMRwkzJJNuAZqlmsSsYVltWcHi5-LLP6-e9ZizAsDxb1beruX0LplprNr9YR_ydbbXQ2heY9lxCyq0DDzI7zN7yJ0wg0M4vxY8WYGwRtlm0lYQOZ3EwpWdQ5X_u2p9yN2QbQ?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" width="624" height="128" role="presentation"></strong><br>
<em>Figure 2. Ideal Comp-Comm overlapping in all-gather based Context Parallel</em></p>
<p>We did a simple <a href="https://pytorch.org/blog/introducing-pytorch-profiler-the-new-and-improved-performance-tool/" rel="noopener nofollow ugc">profiling</a> on the Llama3-8B model on 8 H100 GPUs with Context Parallel degree 8 to gain the trace of both rotation approaches and see which approach performs better. When profiling, we disabled torch.compile and activation checkpointing. And the sequence length and local batch size are 8k and 1 correspondingly.</p>
<p><strong>All-to-All</strong>: The all-to-all collective which is necessary for the next SDPA computation does overlap with the current SDPA. However, in this case the collective cost is too high (~470 us) and cannot be fully overlapped.</p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXf43lMb22X02JUFtrujmS9EIstT9DDmgOGTjKQOm4O9E6gIWF3d319fm9-Vc9XrlNglo2eAhWNqxY8p_kOJAsd2GtCPt6ErXXeQ_QEuMY9FDdpm8PF8wd-GhlyQNibJfjTyaf4dRA?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" width="624" height="91" role="presentation"></strong><br>
<em>Figure 3. CPU thread trace of Context Parallel using all-to-all</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXcD_JHh0cN8BAAOR9QKp6b9U0K5c48c_qc8iYcQJrcOutAk_jd4XJKZLxBlf8Uu1iVrwJfVFPj1bZu--VZWeuwylJqNNLgcykE6zxXWRfKcLbV0i5iLZbxhRVS3dk-ws4rfQNWfxQ?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" width="624" height="57" role="presentation"></strong><br>
<em>Figure 4. CUDA stream trace of Context Parallel using all-to-all</em></p>
<p><strong>All-gather</strong>: The situation is similar to the all-to-all case except there’s only one all-gather needed while the all-to-all approach needs to issue N-1 collectives where N is the Context Parallel degree.</p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXcrO33SBOQTTdHFmHg4Ic1vgojgnBP5yJrodyDGq-lqEMG8Mci-9tpcUXXJ5iWQVx5AotnQG1xQ9X7dl72Bj1cLygnaxTqfyEKn2CGnumlgvN96gTVdZ1U8zCMglM-h9hhlfeZi?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" width="624" height="88" role="presentation"></strong><br>
<em>Figure 5. CPU thread trace of Context Parallel using all-gather</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXdQn1gD9KoLsKtNx8tRAOsDYAeCkxtIj8ENCnubgavwmtNXQ6m1eBfx8aBeOBOAGO9xD0pEmcV92MsT7P9z0qPx0Bq3yQosMitFuHhM4gs5A4kTAHeNnx0YedqW3yQnCy6A3r2AhA?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" width="624" height="45" role="presentation"></strong><br>
<em>Figure 6. CUDA stream trace of Context Parallel using all-gather</em></p>
<p>We choose the All-gather approach as the default rotation approach for the following reasons:</p>
<ol>
<li>The all-gather approach is simpler and easier to tune performance because it has less collectives. This is especially true when Context Parallel spans over multiple hosts.</li>
<li>The all-to-all approach requires full overlap for best performance and the penalty for exposing is higher than the all-gather one.</li>
</ol>
<h3><a name="p-462811-causal-masking-and-load-balancer-12" class="anchor" href="#p-462811-causal-masking-and-load-balancer-12"></a>Causal Masking and Load Balancer</h3>
<p>Just as PyTorch scaled_dot_product_attention, we support Causal Masking in our Ring Attention. If users have <code>is_causal=True</code> in their model code, the replacing Ring Attention will also use Causal Masking.</p>
<p>Since a lower triangular masking will be applied to the attention scores if the causal masking is specified, the local SDPA can be optimized to save computation by applying to a part of the QKV shard only. However, this also leads to the imbalance of computation among local SDPAs, and prohibits a good overlap of data communication and computation.</p>
<p>To remedy this problem, we also implemented a Round Robin Load Balancer which shards the shortest query block and the longest one to rank 0, then the second shortest query block and the second longest one to rank 1, and so on. The length variance of query blocks comes from the triangularity of the causal mask.</p>
<p>If <code>is_causal=False</code>, a Sequential sharder is used instead, which shards one contiguous QKV block to each rank within the Context Parallel group.</p>
<h3><a name="p-462811-sdpa-merger-13" class="anchor" href="#p-462811-sdpa-merger-13"></a>SDPA Merger</h3>
<p>This module is responsible for calculating the output and log_sumexp from the results of local SDPAs (the ATEN op). Note that the merging formula is slightly different for the Causal Masking case.</p>
<h1><a name="p-462811-experimental-methodology-14" class="anchor" href="#p-462811-experimental-methodology-14"></a>Experimental Methodology</h1>
<p>We first used torchtitan’s Llama3-8B model to verify the experimental Context Parallel feature and verify the composability. Then we used torchtitan’s LLama3-70B model to understand the benefit of using torch.compile with Context Parallel. The benchmark was performed on a cluster where each host has 8 H100 Nvidia GPUs.</p>
<h1><a name="p-462811-evaluation-15" class="anchor" href="#p-462811-evaluation-15"></a>Evaluation</h1>
<h2><a name="p-462811-sequence-length-scaling-16" class="anchor" href="#p-462811-sequence-length-scaling-16"></a>Sequence Length Scaling</h2>
<p>To measure the longest possible sequence length under our experiment setup, we fix the local batch size to 1 and modify the input sequence length and Context Parallel degree (i.e. group size).</p>
<p>We want to show the effectiveness of our implementation in several aspects:</p>
<ul>
<li>To measure the longest possible sequence length we can achieve w/o OOM, we turned off torch.compile in this case because it introduces memory usage imbalance across GPUs where some ranks have significantly higher memory footprint than others. We suspect this is related to our causal masking and load balance not working well with torch.compile.</li>
<li>When doubling the Context Parallel degree, the max sequence length without causing OOM doubles.</li>
<li>When seq_length/context_parallel_degree is fixed, the memory usage is stable while MFU (Model FLOPS Utilization) drops because of intra-host communication.</li>
<li>The all-gather approach has better performance in the multi-host case because it issues less communications.</li>
</ul>
<div class="md-table">
<table>
<thead>
<tr>
<th>seq_length</th>
<th>256K</th>
<th>512K</th>
<th>1M</th>
</tr>
</thead>
<tbody>
<tr>
<td>all-to-all</td>
<td>MEM: 81.01GiB<br>TPS: 781<br>MFU: 36.11%</td>
<td>MEM: 77.11GiB<br>TPS: 196<br>MFU: 17.22%</td>
<td>MEM: 75.73GiB<br>TPS: 70<br>MFU: 12.05%</td>
</tr>
<tr>
<td>all-gather</td>
<td>MEM: 82.89GiB<br>TPS: 706<br>MFU: 32.64%</td>
<td>MEM: 75.65GiB<br>TPS: 320<br>MFU: 28.16%</td>
<td>MEM: 66.62GiB<br>TPS: 150<br>MFU: 25.66%</td>
</tr>
<tr>
<td></td>
<td></td>
<td></td>
<td></td>
</tr>
<tr>
<td>Required CP degree</td>
<td>8</td>
<td>16</td>
<td>32</td>
</tr>
</tbody>
</table>
</div><p><em>Table 1. max seq_len scaling vs. Context Parallel degree</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfUM19MQ7TPl8K5a7qq1JvdTb1wKIUI5IuphO8oH_A8ptBot3jXnIPHqBnncIXYKH4rGuYXaTUHrtMmB7GZYdZeVR4zRQEx1JY9xNeuTTpMZ2qOhUqvbVG4CqCOHfGTnu0aSraH?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" title="Chart" width="624" height="385" role="presentation"></strong><br>
<em>Figure 7. peak memory usage of training long context data with Context Parallel</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXerjNQSWZiyLJRl6Ojl5jejVONdaKM2iCeMAQVqsqrDrfNWwibZCFNjFeUlY2QajugwD5-nBWAt_xulpGdXzElKr8L4Hr0h7qgVjiVOW6NUPPS5MogOXWAwSIta_vhDMgyrEfxUQQ?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" title="Chart" width="624" height="385" role="presentation"></strong><br>
<em>Figure 8. MFU of training long context data with Context Parallel</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfoLXWunDd5gocLo0xWLn-GlDGjK3Qe7t7e7s7bSQTN4y6sR929ouR3LsqQYM72hSzEcG3Tl6Eg93rp7TEh69fx1U3MAfGMK1zOuiX5afY3jxHVLYVTtSKZWNBQnynSUxVkpDuZRA?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" title="Chart" width="624" height="385" role="presentation"></strong><br>
<em>Figure 9. TPS of training long context data with Context Parallel</em></p>
<h2><a name="p-462811-composability-with-other-parallelisms-17" class="anchor" href="#p-462811-composability-with-other-parallelisms-17"></a>Composability With Other Parallelisms</h2>
<p>So far, we already have many types of parallelisms in PyTorch that can be adopted in transformer model training such as FSDP/HSDP, Tensor Parallel, and Pipeline Parallel. It’s critical for our Context Parallel to work seamlessly with those techniques, therefore we tested the composability with them by launching test training jobs and observing the loss converging trend. Below is a collection of verified combinations of parallelism on Llama3 8B in torchtitan. The verification requires the loss converging curve shows a close match with the one from the non-CP training (e.g. if the combination is PP+FSDP+TP+CP, then the object of comparison is PP+FSDP+TP).</p>
<p>Each training job runs for 600 warm-up steps and 3,000 total steps. All verifications are performed with and without torch.compile. We verified Context Parallel works in combination with the following parallelisms:</p>
<ul>
<li>FSDP/HSDP</li>
<li><a href="https://pytorch.org/tutorials/intermediate/TP_tutorial.html" rel="noopener nofollow ugc">Tensor Parallel</a> (TP)</li>
<li>FSDP/HSDP+TP</li>
<li><a href="https://pytorch.org/tutorials/intermediate/pipelining_tutorial.html" rel="noopener nofollow ugc">Pipeline Parallel</a> (PP) + FSDP/HSDP</li>
<li>PP + FSDP/HSDP + TP</li>
</ul>
<p>It’s worth note that for Pipeline Parallel we chose “Interleaved 1F1B” schedule instead of “Interleaved ZeroBubble” because Zero Bubble currently doesn’t work with torch.compile.</p>
<h2><a name="p-462811-compatibility-with-torchcompile-18" class="anchor" href="#p-462811-compatibility-with-torchcompile-18"></a>Compatibility With Torch.compile</h2>
<p>The experience of using torch.compile w/ Context Parallel in torchtitan is smooth, and the user only needs to set “compile = true” in .toml file to turn on compiler mode.</p>
<p>To understand the impact of using torch.compile, we ran Llama3-70B model on 64 H100 GPUs for 20 steps, with and without torch.compile, and compared their peak memory usage and tokens per second (TPS). The Data Parallel degree is 8 and Context Parallel degree is also 8. We also vary the input sequence length from 8K to 128K to understand the impact over a broad spectrum of sequence size.</p>
<p>The result shows that when the sequence length is small, the peak memory usage is slightly higher and the TPS is significantly higher with the compile mode. When the sequence length is large, the peak memory usage and TPS are both lower with the compile mode. On average, the compile mode provides a 10% TPS gain with 2% memory usage regression.</p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfBZ5TVIqUB4lqCAxzGrHWZyw4N7ZI-cfIDLT8xn_rUtR5gg7gnl9coFj0N1Cn3u-wEO8L51tG1L1LzuQ2ByZy_9POmpXlIgtm3wnAg40Tmpzpl4ZwcsIBX0hwqjlssmm-cgQimyg?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" title="Chart" width="624" height="385" role="presentation"></strong><br>
<em>Figure 10. peak memory usage of training long context data with torch.compile on 64 GPUs (FSDP+CP)</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfqX3MKnT9tUCOv3WAWItyV7w1QL2cwCL2JCPPSWUC0fh2-dMilkTKyO0UQ1R2EhiXQfthsW74N__wUjEoyjx49emIx2rliI0PIBi-63r72-yMApf_BTZuD6U0bSGUlSGEloq35?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" title="Chart" width="624" height="385" role="presentation"></strong><br>
<em>Figure 11. MFU of training long context data with torch.compile on 64 GPUs (FSDP+CP)</em></p>
<p><strong><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXfhGQhw1LFA2WvDkW0cFAMIamKpqUoyhWHSHv9qtuLNJECDNQrjcU9WQqfw46MP7vhN6QzpwbojoczZaK3uJdrQ7DqjxK-fjCACsCMGlUA-aDBSvl9O3wueYeM9zcQHaEc1MXkzcw?key=r7fEqpJHwYG2-L_17KTpA49p" alt="" title="Chart" width="624" height="385" role="presentation"></strong><br>
<em>Figure 12. TPS of training long context data with torch.compile on 64 GPUs (FSDP+CP)</em></p>
<p>One of the reasons is sometimes the compiler decides to not overlap the SDPA computation and shard rotation. In this case, the compiler believes the shard rotation and the first SDPA call are data-dependent and decides to enforce the collective and computation being strictly separate.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/3/c3efb2bbdc52d1c0a9a45d0fa6e4ea0dd2f8fbc4.png" data-download-href="https://discuss.pytorch.org/uploads/default/c3efb2bbdc52d1c0a9a45d0fa6e4ea0dd2f8fbc4" title="allgather_compile_data_dependent"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3efb2bbdc52d1c0a9a45d0fa6e4ea0dd2f8fbc4_2_690x65.png" alt="allgather_compile_data_dependent" data-base62-sha1="rXkEMgx5jtQTVI7FXdDjJ7S4Q7y" width="690" height="65" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3efb2bbdc52d1c0a9a45d0fa6e4ea0dd2f8fbc4_2_690x65.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3efb2bbdc52d1c0a9a45d0fa6e4ea0dd2f8fbc4_2_1035x97.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/3/c3efb2bbdc52d1c0a9a45d0fa6e4ea0dd2f8fbc4_2_1380x130.png 2x" data-dominant-color="D4BCB7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">allgather_compile_data_dependent</span><span class="informations">4738×452 187 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<em>Figure 13. CPU thread trace showing all-gather is correctly placed between the first and second SDPA call</em></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/c/6ca338b85eed31190647cd6c64c109127db4e39b.png" data-download-href="https://discuss.pytorch.org/uploads/default/6ca338b85eed31190647cd6c64c109127db4e39b" title="allgather_compile_data_dependent_cuda"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6ca338b85eed31190647cd6c64c109127db4e39b_2_690x45.png" alt="allgather_compile_data_dependent_cuda" data-base62-sha1="fv3gzRIyE1nrJmotutqZQ8DyTR9" width="690" height="45" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6ca338b85eed31190647cd6c64c109127db4e39b_2_690x45.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6ca338b85eed31190647cd6c64c109127db4e39b_2_1035x67.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/c/6ca338b85eed31190647cd6c64c109127db4e39b_2_1380x90.png 2x" data-dominant-color="D0C0C7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">allgather_compile_data_dependent_cuda</span><span class="informations">4574×300 52.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<em>Figure 14. CUDA stream showing all-gather is exposed</em></p>
<p>We plan to address this issue in future updates.</p>
<h1><a name="p-462811-conclusion-19" class="anchor" href="#p-462811-conclusion-19"></a>Conclusion</h1>
<p>In conclusion, we introduced Context Parallel to PyTorch through the use of Ring Attention that allows users to train long-context LLMs without OOM. With Context Parallel, we successfully train Llama3-8B in torchtitan with 1 million sequence length. We also tested its composability with other PyTorch’s native features such as torch.compile, FSDP/HSDP, Tensor Parallel, and Pipeline Parallel, and Context Parallel works out-of-box with those features. Looking forward, we will continue to improve the Context Parallel on compatibility and performance. We welcome any contributions and suggestions from the community on this work.</p>

githubsgi: <p>Is there a stand alone  sample code available to try cp ?</p>

nlpfollower: <p>Try torchtitan cp integration tests.</p>

sen-ppl: <p>Will there be support for custom attn_mask and flex attention?</p>

fduwjj: <p>I think <a class="mention" href="/u/xwu">@XWu</a> is having plans for this already?</p>

sen-ppl: <p>Hello, I was trying out the example at <a href="https://pytorch.org/tutorials/prototype/context_parallel.html" class="inline-onebox" rel="noopener nofollow ugc">Introduction to Context Parallel — PyTorch Tutorials 2.7.0+cu126 documentation</a>,  when i remove the detach from cp_qkv = [t.detach().clone() for t in qkv]<br>
I got the following error</p>
<pre><code class="lang-auto">.venv/lib/python3.12/site-packages/torch/distributed/tensor/experimental/_attention.py", line 1348, in context_parallel
[rank1]:     buffer.resize_(chunk.shape)
[rank1]: RuntimeError: cannot resize variables that require grad
</code></pre>
<p>May I ask why it only works with tensors that don’t require grad? It would be much easier if we can wrap only the SDPA with CP so that we don’t have to deal with other parts of module where the logic depends on full sequence.</p>

fegin: <p><a class="mention" href="/u/sen-ppl">@sen-ppl</a> The context parallel parallelize the input/activation (and hence the QKV). Any <code>buffers</code> that are sequence related, such as buffer for positional embedding should also need to be sharded. This will allow to actually save the memory because the whole activations are sharded. Essentially, you can the design of PTD’s CP as a variation of data parallelization. If you only want to wrap the SDPA only, some collectives (depending on what parallelism that is used before CP) are required, which currently we don’t support yet.</p>
<p>As for FlexAttention, <a class="mention" href="/u/xwu">@XWu</a> has a prototype. But it will still require to parallelize the entire model with CP, not just FlexAttention.</p>

Rohit_Kumar_Singh: <aside class="quote no-group" data-username="XWu" data-post="1" data-topic="215082">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/x/b487fb/48.png" class="avatar"> XWu:</div>
<blockquote>
<pre><code class="lang-auto">with train_context(context_parallel_ctx):  # enable Context Parallel
		pred = model(input_ids)
		loss = loss_fn(pred, labels)
		del pred
		loss.backward()
</code></pre>
</blockquote>
</aside>
<p><a class="mention" href="/u/xwu">@XWu</a> Can we use <code>loss.backward()</code> outside <code>context_parallel_ctx</code>? As I understand <code>loss.backward()</code> starts backpropagation of each cp rank separately. So, it should not affect if backward is invoked inside or outside the context, as loss remains shareded even outside <code>context_parallel_ctx</code>.</p>

Rohit_Kumar_Singh: <p>Found my answer here: <a href="https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/experimental/_attention.py#L1221" class="inline-onebox" rel="noopener nofollow ugc">pytorch/torch/distributed/tensor/experimental/_attention.py at main · pytorch/pytorch · GitHub</a> old (Non ring-attention based) <code>F.scaled_dot_product_attention</code>is being restored before exiting cp context.</p>

githubsgi: <p>Wondering what are the reasons for requiring fused attention kernel for  Context Parallel ?  Ring  Attention in TorchTitan appears to be a wrapper around a fused attention kernel.  Would appreciate any insight,</p>

oahz48: <p>During the training, the context parallel shards the input into partial sequence across the CP group, that means after backward, model parameters grads diverge within the same CP group, do we need to manually reduce the parameter grad?</p>

XWu: <p><a class="mention" href="/u/oahz48">@oahz48</a> This is a good question. The param grad reduce-scatter is automatically handled because we use shard model params over both DP and CP meshes: <a href="https://github.com/pytorch/torchtitan/blob/3b9df60ed1c04e5ea86cfc339881e5dbfb18670a/torchtitan/models/llama3/infra/parallelize.py#L95-L110" class="inline-onebox" rel="noopener nofollow ugc">torchtitan/torchtitan/models/llama3/infra/parallelize.py at 3b9df60ed1c04e5ea86cfc339881e5dbfb18670a · pytorch/torchtitan · GitHub</a></p>
<p>FSDP: <a href="https://docs.pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.ShardingStrategy" class="inline-onebox" rel="noopener nofollow ugc">FullyShardedDataParallel — PyTorch 2.8 documentation</a></p>