FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism

Yujie Wang, Shiju Wang, Shenhan Zhu, Fangcheng Fu, Xinyi Liu, Xuefeng Xiao and 4 others

Training a language model on long documents means splitting each sequence across many GPUs, but real training data is a mix of a few very long sequences and a flood of short ones, and every existing system forces all of them onto the same large, communication-heavy split sized for the longest — wasting time on the short majority; this work from Peking University and ByteDance lets each training step form parallel groups of different sizes at once, routing short sequences to small fast groups and long ones to big groups, choosing the split by solving an optimization problem, for up to 1.98 times faster training.

The argument in one glance

Problem

Long-context training splits sequences across GPUs, and it is costly

Training a language model on long documents needs enormous GPU memory, so sequence parallelism splits each input sequence across multiple devices to fit. But splitting forces the devices to communicate to exchange intermediate results, and the more devices a sequence is spread across, the heavier that communication becomes — dragging down training speed. So long sequences need a large split to fit in memory, but a large split is slow.
What people did

Pack sequences and use one fixed split for the whole batch

Existing systems assume all sequences are the same length. They use sequence packing — concatenating varied-length sequences into uniform-length inputs, adjusting attention masks to keep them separate — and then apply a single, static split size (sequence-parallelism degree) to every input, hand-tuned for the workload.
Why it fell short

A fixed split forces short sequences onto the large group

Real training corpora are not uniform: their sequence lengths follow a long-tail distribution, a few very long sequences among a majority of short ones. To fit the longest sequences in memory the split must be sized large, and because the split is fixed, every short sequence is forced onto that same oversized group, paying its heavy communication cost even though a small group would handle it far faster. The dominant short-sequence majority ends up bottlenecked by the rare long ones.
This paper's idea

Form heterogeneous parallel groups per step, chosen by optimization

FlexSP lets each training step form multiple sequence-parallel groups of different sizes at once — large groups to hold the long sequences within memory, small groups to run the short ones with far less communication — and carefully assigns each sequence to a group to balance the load. It picks the groups and the assignment by modeling compute, communication, and memory costs as a mixed-integer linear program, shrunk by adaptive length-bucketing so a solver cracks it in seconds, for up to 1.98 times faster training.

Why this problem is worth solving

The core trade-off is that the very thing that lets a long sequence fit is what makes it slow.

To train on long context, a sequence too big for one GPU’s memory is split along its length across several GPUs — sequence parallelism. Each GPU holds a shard, and to compute attention, which needs the whole sequence, the GPUs must exchange their shards through collective communication (an All-to-All exchange). The size of the group a sequence is split across — its parallelism degree — sets both how long a sequence it can hold and how much communication it pays. A larger group can fit a longer sequence, but it also means more devices exchanging data, and crucially, once a group grows beyond a single machine, that exchange must cross the network between machines, which is far slower than the high-bandwidth links inside one machine. On a typical cluster the intra-machine interconnect (NVLink) is many times faster than the inter-machine one (InfiniBand), so a group small enough to fit inside one machine communicates cheaply, while a group spanning several machines spends a large fraction of its time just moving data. The authors measure this directly: for a batch of 8K-token sequences, using a group larger than one machine pushes communication past 31 percent of the time, while a single-machine group cuts it to under 8 percent — for identical computation. So the group size is a lever with memory on one end and communication cost on the other.

What turns that trade-off into an opportunity is the shape of real data.

If every training sequence were the same length, one group size would be the right answer for all of them and there would be nothing to optimize. But real training corpora are nothing like uniform. The authors measure the sequence-length distributions of popular datasets — GitHub, CommonCrawl, Wikipedia — and find a pronounced long tail: a small number of very long sequences, and a large majority of short ones. In Wikipedia, over 96 percent of sequences are below 8K tokens, with only a tiny sliver above 32K. This matters because the memory needs of the two ends are wildly different: a 32K sequence fits in a group of 8 devices, but a 128K sequence needs at least 32. When a system commits to one fixed group size, it must size that group for the longest sequences it might see, and then the short-sequence majority — which would run happily in a small, fast, single-machine group — is dragged onto the huge, slow, multi-machine group anyway. Because the short sequences are the overwhelming majority, this misallocation wastes most of the training time. The long tail is precisely why a one-size-fits-all split is so wasteful, and why matching the split to each sequence’s actual length can accelerate almost the entire workload.

What people did before, and where it broke

Homogeneous sequence parallelism with packing (Megatron-LM, DeepSpeed). Assume uniform lengths, pack sequences into equal-length inputs, and apply one fixed group size to all. Where it broke. The group must be sized for the longest sequences, so the short-sequence majority is forced onto that oversized, communication-heavy group and runs far slower than it needs to.
Best-fit sequence packing. Pack varied-length sequences to best fill device memory, eliminating the padding waste of naive approaches. Where it broke. Packing makes inputs uniform in length, but it still hands every packed input to the same static split, so it inherits the same one-size-fits-all inefficiency for short content.
Per-batch adaptive strategy. Change the parallelism strategy from one data batch to the next based on that batch. Where it broke. Within any single batch the strategy is still homogeneous, so a batch that mixes long and short sequences (as real batches do) still forces the short ones onto the batch’s large group.
The gap. No system matched the parallelism to the heterogeneity that exists within a single step: large groups for the rare long sequences to satisfy memory, small groups for the common short ones to run fast, all at once. Forming heterogeneous groups per step and assigning sequences to them optimally is what FlexSP provides.

How FlexSP works

From one fixed split to per-step heterogeneous groups

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Match the split to the workload’s heavy tail, spending big groups only on the rare long sequences. The corpus is dominated by short sequences that a small, fast group handles beautifully, with a thin tail of long sequences that genuinely need large, slow groups; FlexSP gives each what it needs in the same step rather than forcing the majority to pay for the minority. Splitting a workload by its heavy-tailed structure and handling the dominant common case cheaply while spending the expensive machinery only on the rare hard case is the same strategy the low-bit language-model serving work covered elsewhere on this site uses, keeping the vast majority of values in the cheap fast format and reserving the costly high-precision path for the rare outliers.

Turn the scheduling choice into an optimization a solver can settle. Deciding how many groups of each size to form and which sequence goes where is a tangle of memory limits and communication costs, so FlexSP writes it down precisely — as a mixed-integer linear program over compute, communication, and memory — and lets a solver find the optimum, taming the size with adaptive bucketing. Recasting a hard combinatorial systems decision as a formal math-programming problem, complexity-reduced so a standard solver cracks it, is exactly the move the fermion-to-qubit encoding work covered elsewhere on this site makes, expressing a knotty constrained-optimization as a solver problem and shrinking it so the solver finishes quickly.

Hide the planning cost by overlapping it with the training. Re-planning every step would be self-defeating if it stalled the GPUs, so FlexSP runs its CPU solver disaggregated from the GPU training and overlapped with it, pipelining the plan for the next step behind the current one so the optimization costs nothing on the critical path. Overlapping unavoidable auxiliary work with the main computation so it disappears from the timeline is the same discipline the communication-overlapping work covered elsewhere on this site follows, running the expensive side-work concurrently with the computation that was going to happen anyway so it never shows up as delay.

Following a training step through FlexSP

Planning one step over a mixed batch

Tracing one step shows bucketing, solving, and overlapped planning in action.

Step 1 of 5

A batch arrives with wildly different lengths

A training step brings a batch of sequences of very different lengths — a few 100K-token documents among many 8K ones — to be trained across a 64-GPU cluster. A fixed split would force all of them into one large group.

The long tail means most sequences are short but a few are enormous.

Why it works

Small groups communicate far less. Splitting a sequence across more devices adds communication that dominates once it crosses to the slow between-machine network → but a short sequence fits in a small group inside one machine’s fast interconnect → so routing short sequences to small groups slashes their communication cost.
The short majority is what to accelerate. Real corpora are long-tailed, so short sequences are the overwhelming majority → but a fixed split sized for the rare long sequences forces the short majority onto it → so adapting the split to length speeds up almost every sequence.
Big and small groups coexist in one step. Long sequences still need a large group to fit in memory → but heterogeneous groups let those coexist with small groups in the same step → so memory limits are met without penalizing the short sequences.
Balanced assignment stops groups idling. Sending each sequence to the smallest group that fits would overload the small groups with the dominant short sequences → but modeling the problem and balancing the assignment spreads the load → so no group sits idle waiting for an overloaded one.
The planning is effectively free. Re-planning every step could cost more than it saves → but the choice is a compact linear program, bucketed and solved on the CPU while the GPUs train → so the plan takes seconds and is fully overlapped, off the critical path.

What it costs (trade-offs and scope)

The benefit scales with how skewed the data is. FlexSP wins by accelerating the short-sequence majority; on a corpus of uniform lengths there is no heterogeneity to exploit and it reduces to ordinary sequence parallelism, which is also why its gains are largest on the most long-tailed datasets.
Bucketing trades some accuracy for tractability. Representing many similar lengths by one bucket length makes the optimization solvable but introduces estimation bias; an adaptive dynamic-programming bucketing minimizes it, yet it is nonzero.
The plan is only as good as its cost models. The optimization relies on accurate models of compute, communication, and memory calibrated to the cluster’s bandwidths and kernels; on a different interconnect the plan and the gains shift.
It adds a CPU solving stage. Each step solves a mixed-integer program (roughly 5 to 15 seconds); this is multi-process-parallelized and overlapped with training so it is hidden, but it is real work and depends on a capable solver.
The gains come from the interconnect gap. Small groups win because the links inside a machine are far faster than those between machines; on a cluster with a uniformly fast fabric the advantage of small groups would shrink.
It targets the sequence-parallelism dimension. FlexSP makes sequence-parallel grouping heterogeneous atop sharded data parallelism; it composes with, rather than replaces, tensor and pipeline parallelism, and assumes sequence packing with its gradient equivalence.

The evaluation

Three ways to split varied-length sequences across GPUs

Property	Homogeneous SP (Megatron / DeepSpeed)	Per-batch adaptive
Parallel group sizes		One size per batch
How short sequences are handled		Still uniform within a batch
How the split is chosen	Static, hand-tuned	Adapted per batch
Training speedup	Baseline	Up to 1.34 times

Three ways to split varied-length sequences across GPUs, evaluated on 64 A100 GPUs across model sizes, context lengths, and the GitHub, CommonCrawl, and Wikipedia datasets. By forming heterogeneous sequence-parallel groups per step and solving for the optimal grouping and assignment, FlexSP routes short sequences to small single-machine groups (cutting All-to-All communication by up to 5.86 times) while keeping large groups for long ones, for up to 1.72 times speedup over DeepSpeed and 1.98 times over Megatron-LM — largest on the most long-tailed dataset (Wikipedia). A per-batch-adaptive variant, homogeneous within each batch, reaches only up to 1.34 times, and FlexSP scales more robustly as the cluster grows because it leans on fast intra-machine links.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

FlexSP is a clean first-principles result: it notices that real training corpora are long-tailed, that a one-size-fits-all sequence split therefore drags the short-sequence majority onto a slow oversized group, and that letting each step form heterogeneous groups — big for the rare long sequences, small and fast for the common short ones — recovers most of the lost time. It makes the grouping-and-assignment choice a mixed-integer program solved in seconds via adaptive bucketing, hides that solving behind the training, and delivers up to 1.98 times faster training with markedly better scaling. Read it with its scope. The gains depend on the data being skewed; on uniform-length corpora there is nothing to exploit. The optimization rests on cost models calibrated to the cluster, and its win comes specifically from the gap between fast intra-machine and slow inter-machine links — on a uniformly fast fabric small groups would matter less. Bucketing introduces a small estimation bias, and the per-step CPU solve, though overlapped and parallelized, is real work depending on a capable solver. It optimizes the sequence-parallelism dimension atop sharded data parallelism and composes with, rather than replaces, other parallelism forms. The durable lesson is that systems built on a uniformity assumption leave large gains on the table when the real workload is skewed — and that matching the system’s structure to the actual data distribution, chosen by optimization rather than by hand, can pay off handsomely.

Background

Why matching the system to the real distribution keeps winning

It helps to see FlexSP as one instance of a pattern recurring across large-model systems: the biggest efficiency wins increasingly come not from a faster kernel but from noticing that the workload is not uniform and refusing to treat it as if it were. Training and serving pipelines were long built on tidy assumptions — every sequence the same length, every request the same cost, every expert equally busy — and reality is lumpy, skewed, and heterogeneous, so a system that adapts its structure to the actual distribution reclaims the slack those assumptions waste. The Mixture-of-Experts serving work covered elsewhere on this site pursues the same idea from the serving side, exploiting the uneven, structured demand across a model’s experts to schedule and batch work far more efficiently than a uniform policy would. The broad lesson is that as models and their workloads grow, the uniformity assumptions that once simplified system design become the dominant source of waste, and the systems that win are the ones that measure the real distribution — of sequence lengths, request sizes, expert loads — and let it drive how the work is partitioned and placed, dynamically and per step, rather than committing to a single static plan.

If you remember only three things

1. Long-context LLM training uses sequence parallelism to split each sequence across GPUs so it fits in memory, but splitting forces communication that grows with the group size, and once a group spans multiple machines it pays the slow inter-machine network (many times slower than intra-machine links). Existing systems assume uniform sequence lengths, pack sequences, and apply one fixed group size sized for the longest sequence. Real corpora are long-tailed — a few very long sequences among a majority of short ones — so the short majority is forced onto the oversized, communication-heavy group and wastes most of the training time.
2. FlexSP makes sequence parallelism heterogeneity-adaptive: for each step it forms multiple sequence-parallel groups of different sizes at once (large groups to fit long sequences in memory, small single-machine groups to run short ones with minimal communication) and carefully assigns each sequence to balance the load across groups. It co-optimizes grouping and assignment by modeling compute, communication, and memory as a mixed-integer linear program, made tractable by an adaptive dynamic-programming length-bucketing (default 16 buckets) and solved with SCIP; a sequence blaster chunks over-large batches into balanced, similar-length micro-batches, groups are switched on the fly from a cached NCCL pool, and the CPU solving is disaggregated and overlapped with GPU training so it is off the critical path.
3. On 64 A100 GPUs across model sizes, context lengths, and the GitHub, CommonCrawl, and Wikipedia datasets, FlexSP outperforms state-of-the-art systems by up to 1.72 times over DeepSpeed and 1.98 times over Megatron-LM, cutting All-to-All communication by up to 5.86 times by routing short sequences to small fast groups, with the largest gains on the most long-tailed dataset and more robust scaling as the cluster grows. Its limits: the benefit depends on data skew (no gain on uniform lengths), bucketing adds a small estimation bias, the plan relies on cluster-calibrated cost models, the per-step solve is real (if hidden) CPU work, the win rests on the intra- versus inter-machine bandwidth gap, and it optimizes the sequence-parallelism dimension atop sharded data parallelism.

Check yourself

Check your understanding

1. Why does a fixed sequence-parallelism group size waste time on real training data?
2. What are FlexSP's two key innovations for handling varied-length sequences?
3. How does FlexSP choose the groups and assignment, and keep that choice cheap?
4. Which statement best reflects FlexSP's results and limitations?

Sources

The paper is “FlexSP: Accelerating Large Language Model Training via Flexible Sequence Parallelism” by Yujie Wang, Shiju Wang, Shenhan Zhu, Fangcheng Fu, Xinyi Liu, Xuefeng Xiao, Huixia Li, Jiashi Li, Faming Wu, and Bin Cui, published at ASPLOS 2025, DOI 10.1145/3676641.3715998. The authors are at Peking University, Beihang University, and ByteDance; the system is open-source as part of Hetu-Galvatron.

This explainer is written from the full text of the paper (the authors’ arXiv version, 2412.01523, matching the ASPLOS proceedings). Everything above comes from the paper’s own sections: the introduction and preliminaries on sequence parallelism (Ulysses-style with its All-to-All communication), sequence packing, and other parallelism forms; the motivating example (64 devices, 192K context, heterogeneous versus homogeneous SP groups) and the observations (Observation 1 on large groups being inefficient due to communication, measured in Table 1 across sequence lengths and SP degrees with intra-node NVLink versus inter-node InfiniBand; Observation 2 on the long-tail sequence-length distributions of GitHub, CommonCrawl, and Wikipedia); the FlexSP design with its heterogeneous groups and time-balanced assignment, the mixed-integer linear programming formulation of compute, communication, and memory costs, the adaptive dynamic-programming sequence bucketing (default 16 buckets) and its lower estimation bias versus naive bucketing, and the sequence-blaster micro-batch chunking with its three takeaways (fewer micro-batches, smaller length variance within a micro-batch, evenly distributed token counts); the implementation on PyTorch with SCIP, NCCL, and flash-attn, the hot-switching NCCL group pool, and the disaggregated overlapped solving (roughly 5 to 15 seconds, hidden); and the experiments against Megatron-LM, DeepSpeed, and FlexSP-BatchAda showing up to 1.72 and 1.98 times end-to-end speedups, the up-to-5.86-times All-to-All reduction, the correlation with dataset skew, and the scalability results across 16, 32, and 64 GPUs.