AppendicesAppendix F of 7

Appendix F: Context Parallelism and MoE Parallel Folding

  • 3 min read

Long sequences create an activation-memory and attention-computation problem that Expert Parallelism (EP) does not address. Context Parallelism (CP) divides the sequence dimension among ranks that process the same training examples. With a CP degree \(C\), each rank enters a Transformer block with approximately \(S/C\) of the sequence positions. CP therefore reduces the local size of sequence-shaped activations, but it does not itself partition parameters or create additional data-parallel examples. Routed-expert parameter placement remains the job of EP.

Olmo-core uses the Ulysses form of CP for the MoE path (Jacobs et al. 2023). Before attention, each rank holds queries, keys, and values in a context-sharded layout. An all-to-all operation changes that layout as follows for batch size \(B\), sequence length \(S\), head count \(H\), and head width \(d_h\):

\[ [B, S/C, H, d_h] \quad\longrightarrow\quad [B, S, H/C, d_h], \tag{96} \]

so each rank evaluates the full sequence for a subset of attention heads, which requires \(C\) to divide \(H\). A second all-to-all returns the output to the context-sharded layout. The remaining operations in the Transformer block continue on local token rows. In particular, the MoE router receives the local \(S/C\) rows directly; the model does not gather the full sequence before routing.

CP ranks are not independent data-parallel replicas because they contribute different sequence slices of the same examples. This distinction matters when CP is combined with EP. Let \(D\) and \(C\) be the dense Data Parallelism (DP) and CP degrees within one pipeline stage, and let \(D_E\) and \(M_E\) be the expert-replica and expert-sharding degrees. Olmo-core forms two views of the same per-stage rank set:

\[ \mathcal{G}_{\mathrm{stage}} = \mathcal{A}_{\mathrm{DP}} \times \mathcal{A}_{\mathrm{CP}} = \mathcal{A}_{\mathrm{EP\text{-}DP}} \times \mathcal{A}_{\mathrm{EP\text{-}MP}}, \qquad DC = D_E M_E. \tag{97} \]

The dense view keeps CP visible to attention and input preparation. Dense gradients reduce over the combined DP\(\times\)CP group. The MoE view instead uses EP-MP to partition routed experts and EP-DP to synchronize copies of the same expert shard. Dispatch and combine operate over the folded expert groups, so routed tokens from all DP and CP coordinates can reach the ranks that store their selected experts.

This construction is a focused application of MoE parallel folding, which allows dense attention and sparse expert layers to use different parallel mappings on the same devices (D. Liu et al. 2025). For example, four ranks can form a dense mesh with \(D=2\) and \(C=2\), then be viewed by the MoE layers as \(D_E=1\) and \(M_E=4\). Enabling CP\({}=2\) and EP-MP\({}=4\) therefore still uses four ranks within the stage, rather than allocating an independent four-rank expert group for each CP coordinate. More generally, parallel folding prevents CP and EP from being treated as orthogonal multipliers when their layer-local work can share the same ranks.

Pipeline Parallelism (PP) adds an outer partition over layers. Within each PP stage, inputs and labels are sequence-sharded before the pipeline schedule starts, while the original sequence length is retained for positional encoding. Later stages receive already-sharded activations and do not apply CP a second time. Equation 97 then holds independently inside every pipeline stage. Within this layout, CP assigns sequence rows, EP assigns routed experts, and PP assigns layers to ranks.

CP lowers the per-rank footprint of many sequence-shaped activations and makes longer contexts feasible, while Ulysses adds attention-layout all-to-alls. Parallel folding does not remove that communication; it prevents CP from unnecessarily multiplying the device count required by EP and gives attention and routed experts the process groups that match their respective computation.

Search the full report

Search across sections, appendices, and references.

    Enter a term to find passages. Esc closes search.