DCP: Addressing Input Dynamism In Long-Context Training via Dynamic Context Parallelism

TL;DR

DCP tackles the inefficiencies of static context parallelism in long-context training by introducing dynamic, blockwise data/computation partitioning guided by hypergraph partitioning. It models input dynamism (varying sequence lengths and diverse attention masks) as a per-iteration planning problem, producing per-device execution plans that adapt to current workloads. The framework includes a block-centric executor, a planning-overlapping data loader, and a hypergraph-based scheduler, achieving up to 1.16x–1.46x end-to-end speedups and larger gains on sparse attention scenarios, while preserving training accuracy. This approach enables more memory- and compute-balanced, communication-efficient long-context training that scales with data dynamism and mask patterns.

Abstract

Context parallelism has emerged as a key technique to support long-context training, a growing trend in generative AI for modern large models. However, existing context parallel methods rely on static parallelization configurations that overlook the dynamic nature of training data, specifically, the variability in sequence lengths and token relationships (i.e., attention patterns) across samples. As a result, these methods often suffer from unnecessary communication overhead and imbalanced computation. In this paper, we present DCP, a dynamic context parallel training framework that introduces fine-grained blockwise partitioning of both data and computation. By enabling flexible mapping of data and computation blocks to devices, DCP can adapt to varying sequence characteristics, effectively reducing communication and improving memory and computation balance. Micro-benchmarks demonstrate that DCP accelerates attention by 1.19x~2.45x under causal masks and 2.15x~3.77x under sparse attention patterns. Additionally, we observe up to 0.94x~1.16x end-to-end training speed-up for causal masks, and 1.00x~1.46x for sparse masks.

Figures (22)

• Figure 1: CP communication overhead when training a 8B GPT model on an Amazon EC2 p4d.24xlarge cluster (400Gbps interconnect between nodes) with 4-way tensor and 16-way context parallelism, using the LongAlign bai2024longalign dataset. Overlap: overlapping CP communication and attention computation. Communication overhead fraction (vs. total iteration time) is shown above each bar.
• Figure 2: Sequence length distribution of LongAlign bai2024longalign and LongDataCollection togethercomputer2024longdatacollections datasets, capped at 131072.
• Figure 3: Four parallelizable dimensions of an attention operator. Each block represents the computation of $\mathbf{\hat{O}_{bhqk}}$ (Listing \ref{['lst:parallel_attn_algo']} line 5) with corresponding Q and KV blocks. Figure shows attention of two sequences (B0 and B1) with lengths 4 and 3 tokens, respectively, using causal mask.
• Figure 4: Special data and computation placement for causal mask.
• Figure 5: Communication and computation required under different parallelization configurations. Color distinguishes different sequences. Round-cornered blocks: input token blocks (4 blocks in each sequence), where each block in the blue sequence contains twice as many tokens as green and orange sequences. Square blocks $[a, b]$: attention between token block $a$ and token block $b$. Arrow: communication of KV for corresponding blocks.