Comparing Ulysses and Megatron Sequence Parallelism

27 Jul, 2026

Introduction

I recently built a small (mainly pedagogical) distributed training engine to understand how different forms of model parallelism work in practice. This post compares sequence parallelism as used in Megatron-LM and DeepSpeed. Although both systems keep the transformer residual stream partitioned as [B, S/P, H] (where P is group size), the mechanisms they use around attention are quite different. Megatron couples sequence parallelism to tensor-parallel QKV, output-projection, and MLP weights, whereas Ulysses (DeepSpeed) redistributes projected QKV activations between sequence and head layouts but makes no changes to the MLP block.

Note that this isn't a complete comparison of the two systems. For instance, I don't speak to how communication can overlap with computation, which can have material effects on performance. The goal of the blog is really to follow the rank-local tensors and collectives used by each sequence-parallel attention implementation.

A simple attention skeleton:

class Attention(nn.Module):
    def __init__(self, ...):
        self.q_proj = Linear(hidden_size, n_head * head_dim)
        self.k_proj = Linear(hidden_size, n_kv_head * head_dim)
        self.v_proj = Linear(hidden_size, n_kv_head * head_dim)
        self.output_proj = Linear(n_head * head_dim, hidden_size)
    
    def forward(self, x):
        B, S, H = x.size()
        # qkv projections
        q = q_proj(x).view(B, S, n_head, head_dim)
        k = k_proj(x).view(B, S, n_kv_head, head_dim)
        v = v_proj(x).view(B, S, n_kv_head, head_dim)
        q, k = apply_rope(q), apply_rope(k)
        q, k, v = transpose(q, k, v)
        # attention
        y = F.scaled_dot_product_attn(q, k, v)
        y = tranpose_y(y)
        # output projection
        y = self.output_proj(y)
        return y

Megatron Attention

Megatron makes the observation that attention heads can be computed in parallel and therefore shards the heads produced by the Q, K, and V projections across tensor-parallel ranks. As seen here, linear_qkv is column parallel, while the attention output projection is row parallel.

Importantly, at the beginning of the forward function, each rank owns a sequence shard of shape [B, S / group_size, H]. Although the QKV projection itself is sequence-independent, the attention computation for a given head requires that head’s queries, keys, and values for the entire sequence. Megatron therefore all-gathers the sequence shards before the column-parallel QKV projection, giving every rank an input of shape [B, S, H]. The input data gets replicated across ranks, but each rank owns the rows of the QKV weights corresponding to its local attention heads.

Assuming ordinary multi-head attention, the local qkv projection (Megatron fuses their q, k, and v projection) consequently produces data shaped [B, S, 3H / group_size], which is reshaped into Q, K, and V tensors containing the full sequence but only n_head / group_size heads. Each rank can then compute attention independently over its local heads, producing an activation of shape [B, S, H / group_size]. The row-parallel output projection consumes this output so that each rank creates a partial contribution of shape [B, S, H], which must be summed across ranks. Megatron combines the summing across ranks and restoring of sequence-parallel in a single collective here.

The backward necessarily needs to reverse this. The row-parallel output projection all-gathers its sequence-sharded output gradient before computing its local gradients. For the QKV projection, every rank computes a partial gradient with respect to the full hidden-state input, and a reduce-scatter, similar to their forward, both sums those contributions and returns [B, S / group_size, H] to the preceding sequence-parallel operation. Megatron also needs an additional all-gather of the saved sequence-sharded input during the QKV backward pass because computing the local QKV weight gradient requires the input for all sequence positions.

Megatron Attention costs

Now, lets look into how much this costs in terms of memory and bandwidth. For both Megatron (and later) DeepSpeed we assume:

TP/SP group size P = 8 GPUs
Sequence length S  = 32,768
Microbatch size B  = 1
Hidden size H      = 4,096
Attention heads    = 32
Head dimension D   = 128
Activation dtype d = BF16 = 2 bytes/element

The tensor [B, S, H] occupies 256MB, which divided amongst 8 ranks is 32MB per rank. Before the QKV projection, Megatron all-gathers those sequence shards giving each rank a full, temporary 256 MB of input. Importantly, these activations do get materialized per rank while executing the column-parallel projection. After the final output project each rank holds a partial output of shape [B, S, H], which is reduce-scattered back to [B, S / P, H].

In terms of communication, for a standard ring all-gather of an N-byte global tensor over P ranks, each rank will send and receive N * (P - 1) / P bytes, which equates to 224 MB sent and recieved per rank or 1.75GB cumulatively across ranks. The forward reduce-scatter has the same ideal ring traffic. So per rank, 448MB are sent and 448 are received in the forward for attention. The backward as discussed contains an extra all-gather, adding an additional 224MB sent and recieved per rank, equating to 672MB. So across one attention layer's block we can generally say that 5 × ((P−1)/P) × N bytes are sent per rank where N = B × S × H × d across the entire attention forward and backward.

In terms of memory, long-lived residual activations come out to about O(B * S * H / P) per rank though the gathered QKV input and partial output-projecion require temporary buffers of size O(B * S * H). Furthermore, multiple communication and GEMM buffers during the backward pass would push these storage requirements still higher.

DeepSpeed Attention

Ulysses has the same attention-time layout as Megatron, where at the transformer-block boundary each rank owns a sequence shard of shape [B, S/P, H]. Each rank then applies the complete QKV projections to its local tokens, producing Q, K, and V tensors of shape [B, S/P, num_heads, head_dim]. DistributedAttention receives these already-projected tensors.

Before computing attention, Ulysses performs an all-to-all that exchanges the sequence partition for a head partition. Each rank sends different groups of heads to their corresponding ranks and receives those same heads for every sequence shard. Consequently, each rank's Q, K, and V tensors change from [B, S/P, num_heads, head_dim] to [B, S, num_heads/P, head_dim]. The rank can then compute its assigned heads independently over the full sequence. Importantly, unlike Megatron, the activations here get sharded and not the weights.

After attention, a second all-to-all reverses this transformation where [B, S, num_heads/P, head_dim] becomes [B, S/P, num_heads, head_dim] (equivalently [B, S/P, H]). Crucially, this allows for each rank to again own all the hidden features for its local tokens, so it can apply the output projection matrix locally without another activation collective.

The core difference is that Ulysses shards activations by sequence outside attention and by attention heads inside attention, while keeping the projection weights replicated. Megatron also shards activations by sequence outside attention and by attention heads inside attention, but achieves this by sharding the QKV and output-projection weights.

DeepSpeed Attention costs

Using the same configuration as above, a full [B, S, H] activation occupies 256MB, while each rank's sequence shard [B, S / P, H] occupies 32MB. The first all-to-all changes the layout of each tensor from [B, S / P, H] to [B, S, H / P], but it does not change the number of elements held by a rank. After attention, the activation tensor is also 32MB per rank, and the reverse all-to-all returns it to [B, S / P, H] before the output projection.

For an all-to-all that redistributes an N byte global tensor over P ranks, each rank begins with only N / P bytes. It divides that local shard among the P destinations, retaining one piece and sending the other P - 1. Each rank therefore sends and receives N * (P - 1) / P^2 bytes. For N = 256MB and P = 8, that is 28MB sent and 28MB received per all-to-all.

DeepSpeed's implementation performs one all-to-all for each of Q, K, and V, followed by one for the attention context. These four collectives result in 112MB sent and 112MB received per rank during the forward pass. Autograd applies the inverse all-to-all for each tensor during backward, which is another 112MB sent and received. Across the complete attention forward and backward, each rank therefore sends and receives 224MB of activation data.

However, because Ulysses does not shard the QKV or output-projection weights, every sequence-parallel rank holds the same logical parameters but applies them to a different S/P token shard. Each rank therefore computes parameter gradients using only its local tokens, so DeepSpeed must all-reduce those gradients across the sequence-parallel group. As DeepSpeed currently defaults sequence-parallel gradient communication to FP32. The four attention projection matrices contain 4H^2 parameters, so the gradient payload is G = 4 * H^2 * 4 bytes = 4 * 4096^2 * 4 bytes = 256MB. Given the same N * (P - 1) / P formula above, synchronizing these attention-weight gradients adds 448MB sent and 448MB received per rank, which can generalize to 2G(P - 1) / P. Unlike Ulysses, Megatron’s TP ranks own disjoint projection-weight shards, so it doesn't need to all-reduce their gradients within the TP group.

The total communication is thus (assuming one eight-rank group with no additional data parallelism or ZeRO):

Megatron: 5N(P - 1) / P
DeepSpeed: 8N(P - 1) / P^2 + 2G(P - 1) / P

For the configuration above, DeepSpeed communicates 224MB of activation data and 448MB of parameter-gradient data, totaling 672MB sent and received per rank. Megatron communicates 1,120MB per rank, so it communicates approximately 1.67x as much as DeepSpeed. Put differently, Ulysses communicates 40% less data under these assumptions.

In terms of activation memory, Megatron materializes temporary [B,S,H] tensors of 256MB, whereas Ulysses keeps the corresponding tensors sharded at [B,S/P,H], or 32MB. This is an 8x reduction for those particular buffers, not for total memory. Without ZeRO, Ulysses also stores complete attention weights and gradients on every rank, while Megatron stores TP shards.

If ZeRO-3 is enabled, DeepSpeed shards the parameters, gradients, and optimizer states across the combined sequence-parallel (and data-parallel) group and all-gathers each projection’s complete weights when they are needed for computation. During the backward pass, those weights may need to be gathered again if they were released after forward, while each rank’s local weight-gradient contributions are summed and partitioned using a reduce-scatter. ZeRO-3 therefore reduces persistent model-state memory by roughly the group size, but adds parameter all-gathers and gradient reduce-scatters on top of Ulysses’s activation all-to-alls.

In the configuration above, one all-gather of the 128MB BF16 attention weights transfers 112MB per rank, while reduce-scattering their 256MB FP32 gradients transfers 224MB. Depending on whether the weights are gathered once or again during backward, ZeRO-3 adds approximately 336–448MB of model-state communication per rank to Ulysses’s 224MB of attention-activation communication.

Compared to not using ZeRO-3, when using ZeRO-3 each rank stores and updates only 1/P of the parameters, gradients, and optimizer states. The gradient therefore needs only to reduce-scatter, which both sums the local contributions and leaves each rank with the gradient shard it owns, reducing the gradient communication to G(P - 1) / P (compared to the reduce-scatter and all-gather needing 2G(P - 1) / P without ZeRO-3). However, because each rank needs the complete weights when executing the projections, ZeRO-3 needs a parameter all-gather costing kW(P - 1) / P, where W is the complete parameter size and k is the number of times the weights are gathered (k being 2 if parameters are released after the forward). Combining this model-state communication with Ulysses’s attention-activation all-to-alls gives:

Ulysses without ZeRO:
8N(P - 1) / P^2 + 2G(P - 1) / P

Ulysses with ZeRO-3:
8N(P - 1) / P^2 + (G + kW)(P - 1) / P

With ZeRO-3, and plugging in the numbers like we did above, we see that Megatron communicates approximately 1.67-2x as much (depending on the value of k).

The MLP addendum

DeepSpeed-Ulysses does not introduce any MLP-specific activation collectives or tensor-parallel weight sharding, which is why this post mainly focuses on attention. Each rank simply applies the complete MLP to its local [B,S/P,H] token shard. Megatron’s sequence parallelism, though, is coupled to tensor parallelism throughout the transformer block, including the MLP.

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.up_proj = ColumnParallelLinear(hidden_size, intermediate_size)
        self.down_proj = RowParallelLinear(intermediate_size, hidden_size)

    def forward(self, x):
        x = self.up_proj(x)
        x = F.gelu(x)
        x = self.down_proj(x)
        return x

Megatron's MLP implementation is functionally similar to the above example. Locally, at the MLP boundary, each rank initially holds [B,S/P,H]. The column-parallel up_proj partitions the intermediate output features across ranks, while the row-parallel down_proj partitions its intermediate input features across ranks.

The forwards and backwards flow, as it relates to collectives, is similar to attention. First an all-gather so each rank holds a full [B,S,H] then a reduce-scatter to sum each rank's partial result.

Benchmark comparison

The benchmark is somewhat contrived in that its not a pretrained model, has no embedding layer, vocabulary projection, tokenized data, or model-specific features. Instead, it just repeats the following ordinary transformer block eight times to build a 3.6 billion parameter model.

# Hidden size              = 6,144
# Attention heads          = 48
# Head dimension           = 128
# MLP intermediate size    = 24,576
# Activation dtype         = BF16

x = x + attention(norm1(x))
x = x + mlp(norm2(x))

On the Megatron side, the eight GPUs form one tensor-parallel and sequence-parallel group. The attention and MLP projections use Megatron Core's native ColumnParallelLinear and RowParallelLinear modules, while the optimizer maintains FP32 master parameters and Adam moments for each rank's local parameter shards. On the DeepSpeed side, the projections are ordinary complete PyTorch linear layers, attention uses DistributedAttention, and ZeRO-3 is applied to the entire model. Both communication-overlap options were disabled, as were activation checkpointing, Transformer Engine, fused LayerNorm and activation kernels, pipeline parallelism, and any additional data parallelism.

In the table below, values written as M / D are Megatron / DeepSpeed. Communication is the analytically estimated number of GiB sent per rank in one direction over a complete forward, backward, and optimizer step as talked about above. This was tested on an 8xH100 over NVLink. At short sequence lengths, Megatron wins decisively, with parity being reached around 64K, then DeepSpeed is consistently faster by about 2-3%.

The 256K sequence length is interesting in that Megatron is modeled as sending 210 GiB per rank in each direction over the training step, compared with a 38.72 GiB for DeepSpeed, which is an approximately 82% reduction, and yet the complete DeepSpeed step is only 2.5% faster. In terms of memory usage, ZeRO-3 must temporarily gather complete projection weights and maintain communication buffers, and so DeepSpeed uses more peak allocated memory at every measured sequence length, although the difference narrows from 6.39 GiB at 8K to 2.30 GiB at 256K.

I also tested a hybrid that uses Ulysses for attention, FSDP2 to shard only the complete attention weights, and Megatron's tensor-parallel strategy for the MLP. FSDP gathers the QKV and output-projection weights when attention executes and reduce-scatters their gradients, but doesn't touch the tensor-parallel MLP weights. The MLP therefore retains Megatron's sequence all-gather and reduce-scatter.

In the table below, H / D means Hybrid / DeepSpeed with full-model ZeRO-3. These results use the same eight-layer model and eight-H100 setup as above.

Like Megatron, the hybrid strategy removes much of full-model ZeRO-3's fixed cost at short sequence lengths. But that advantage shrinks with sequence length because the tensor-parallel MLP continues to communicate activations proportional to S, while DeepSpeed's complete MLP remains local to each sequence shard. Full Ulysses with ZeRO-3 catches up around 64K. Its advantage is still only 1.7% at 128K, but grows to approximately 3.6% at 256K. The hybrid does, however, use less peak memory throughout the measured range. It is therefore a useful middle point in that it's substantially faster and smaller than full-model ZeRO-3 at short sequences, close to Megatron's performance at longer sequences, and lower-memory than either implementation in this experiment.

As a final way of interpreting these results, I fit one regression across all 36 Megatron, DeepSpeed, and hybrid measurements. With S_K = sequence_length / 1024, the fitted model was time_ms = 5.40 + 2.80 * peak_memory_GiB + 1.49 * communication_GiB + 4.94 * S_K + 0.14157 * S_K^2. It has an R^2 of 0.9998 and a mean absolute error of 26.4 ms. Longer sequences will simultaneously increase computation, activation memory, and communication, so the regression cannot cleanly separate their effects, which is to say that peak memory, communication, and sequence length are all highly correlated, though I do think that quadratic sequence length gives a useful means with which to understand the trend. For instance, at 256K, the S_K^2 term alone contributes 9.3 seconds of the 11 second step meaning that attention computation dominates. Moreover, communication matters at shorter sequences, when the effects of attention aren't nearly as severe.

Conclusion

Megatron and Ulysses both keep activations sequence-sharded at transformer-block boundaries, but redistribute them differently inside attention. Megatron all-gathers the sequence while tensor-parallel QKV and output projections shard the weights across ranks. Ulysses uses all-to-all collectives to transpose sequence-sharded activations into full-sequence, head-sharded activations, without requiring tensor-parallel attention weights. These choices affect communication volume and peak memory and strongly influence runtime at shorter sequences, but their relative importance diminishes at longer contexts, where the quadratic cost of attention computation dominates the training step.