NSP: Accelerating Variable-Length LLM Training via Nested Sequence Parallelism Yi’ou Wang1 , Xiaoyang Li1,∗ , Yijie Zheng1 , Shouda Liu1 , Yuxuan Wang1 1
arXiv:2609.22755v1 [cs.DC] 19 Sep 2026
∗
ByteDance Seed
Corresponding author
Abstract Long-context LLM training on long-tailed corpora faces a central communication–balance tradeoff. Such sequence-length heterogeneity makes any single sequence-parallelism (SP) degree a poor fit for the workload: a small degree leaves the few long sequences badly imbalanced, while a large degree forces the many short sequences that dominate the workload to pay excessive communication. Existing dynamic-SP systems mix SP degrees within a batch, but to run several groups at once they partition the GPUs into disjoint groups, which reintroduces imbalance across groups and forces costly micro-batch workarounds. We present NSP, a sequence-parallel training system that resolves this tradeoff by nesting differently sized SP groups on shared GPUs within a single training iteration. This lets long sequences use larger SP groups while keeping short sequences on smaller ones, so communication is incurred only where needed and load is balanced per GPU rather than per group. NSP realizes this idea with a tree-structured routing planner that assigns sequences under memory constraints and an executor that exploits the resulting hierarchy through inter-level phase streaming and tree-level recomputation. NSP supports common SP backends and requires no model changes. We evaluate NSP on Qwen3-MoE workloads with up to 384K-token contexts across multiple long-tail datasets on an internal production GPU cluster. Across these settings, NSP consistently improves end-toend training throughput, outperforming Static SP by up to 1.48× and FlexSP by up to 1.16×. Date: September 22, 2026 Correspondence: Xiaoyang Li at [email protected]
1
Introduction
can fit and run efficiently. Real training corpora, however, are not composed of uniformly long sequences. This distribution creates a dilemma when choosing a static SP degree. A small SP degree leaves the long sequences imbalanced or memory-constrained; a large SP degree improves their balance but forces the many short sequences that dominate the batch to pay communication they do not need. The top row of Figure 1 illustrates this dilemma.
Large language models (LLMs) are increasingly trained with long context windows, ranging from a few thousand tokens to hundreds of thousands or more [12, 15, 26]. Such contexts place heavy pressure on the memory and compute resources of a single GPU. Sequence parallelism (SP) [15, 26] is an essential technique for long-context LLM training: it shards each sequence along the token dimension across GPUs, reducing per-GPU activation memory and attention computation so that longer sequences
Dynamic-SP systems [34] exploit this mismatch by 1
giving each sequence its own SP degree, reserving large groups for the few long sequences while the dominant short ones run in small, cheaper groups. To execute multiple groups concurrently, however, these systems partition the GPUs among the groups, so that the groups are necessarily disjoint and each GPU still serves a single SP degree within a pass. Balancing the running time of these disjoint groups then requires either packing many short sequences into one group or dividing the batch into successive micro-batches, both of which forfeit much of the efficiency that the smaller groups were intended to deliver. The bottom-left panel of Figure 1 shows this residual imbalance.
Computation (long / medium / short seq.) L
8×
Communication
8×
L
M
S
8×
L
M
S
8×
L
M
S
8×
L
M
S
SP=16
L
8×
SP=32 8×
M
S
8×
M
S
SP=16
Larger SP: heavy comm
Smaller SP: imbalance
Single SP Degree SP=32
L
8×
SP=16 SP=8
8×
L
M
S
SP=16
L
8×
Our key idea is to allow SP groups of different sizes to share the same GPUs within a single forward– backward pass, rather than partitioning the GPUs among them: a few large groups for the long sequences are nested within many small groups for the short ones. Communication is then incurred only where it is necessary, and load is balanced per GPU rather than per group. Each GPU can participate in a large group for a long sequence and use the remaining time and memory budget for short sequences in smaller groups, all within a single pass, reducing the need for extra micro-batches.
8×
M
S
SP=8
8×
M
S
SP=8
Disjoint SP: group-level imbalance
8×
L
M
S
8×
L
M
S
8×
L
M
S
Nested SP (NSP, ours): balanced load, low comm
Multiple SP Degree
Figure 1 Per-GPU latency and memory under different SP strategies on a long-tail batch. A small Static SP leaves work badly imbalanced (top-left); a large Static SP balances work but pays heavy communication (top-right). Disjoint dynamic-SP groups still leave imbalance across groups (bottom-left). Nesting differently sized SP groups on shared GPUs within one pass keeps communication low and balances load per GPU (bottom-right).
However, turning this idea into a practical system raises two challenges:
throughput. The Executor takes this routing plan as input and executes it efficiently. It further exploits the tree structure through two optimizations. First, inter-level phase streaming overlaps computation and communication across tree levels to reduce critical-path stalls. Second, tree-level recomputation selectively rematerializes activations across tree levels to reduce attention recomputation under memory limits. Figure 1 illustrates the resulting behavior of NSP on a simple long-tail batch, compared with Static SP and disjoint dynamic SP.
• Combinatorial planning. The system must decide where each sequence should run among many possible SP group sizes and placements, while satisfying per-GPU memory constraints and controlling compute and communication cost. • Efficient execution. Once sequences are assigned to specific SP groups, the runtime must execute the resulting nested plan efficiently, without turning the richer execution structure into new critical-path bottlenecks.
We implement NSP on top of an internal PyTorchbased distributed training framework with Fully Sharded Data Parallel (FSDP)-style model-state sharding, and evaluate it on an internal 64-GPU production cluster across two long-context LLM workloads, three datasets, and maximum context lengths of 192K and 384K tokens. In our end-to-end results, NSP outperforms Static SP by up to 1.48× and FlexSP by up to 1.16×.
To address these challenges, we present Nested Sequence Parallelism (NSP), a sequence-parallel training system for long-tail workloads. NSP is organized into two components: a Planner and an Executor. The Planner defines a structured schedule space by constraining admissible SP groups to an SP tree, where each node represents one aligned SP group and each root-to-leaf path represents the groups that may share a GPU. Within this tree-structured space, a profiled cost model scores candidate routing plans. The Planner then applies a beam-search routing heuristic to select a low-cost plan for higher training
In summary, this paper makes the following contributions: • We show that the disjoint-group restriction of existing dynamic SP is the root of its micro2
batch workarounds, and that nesting differently sized SP groups within a single pass removes them.
ken dimension across k GPUs, so that per-GPU activation memory and compute both scale down with k. The difficulty is attention: every query must attend over the entire sequence, whose tokens are scattered across the k GPUs. DeepSpeed-Ulysses [15] resolves this with all-to-all collectives. Before attention, an all-to-all reshuffles the activations from a sequencesharded layout to a head-sharded layout, so that each GPU holds the full-length sequence for a subset of attention heads; after attention, a second all-to-all restores the sequence-sharded layout for the remaining layers. Attention is therefore computed locally over complete sequences, but the two all-to-all collectives sit on the attention critical path, and their cost grows relative to the shrinking per-GPU compute as k increases. Other SP realizations replace this all-to-all with point-to-point ring exchanges [26] or combine the two along orthogonal dimensions [7]; NSP builds on Ulysses and we defer the treatment of alternative SP implementations to Section 4.
• We introduce a cost-guided Planner that constrains multi-SP scheduling to an SP tree and assigns sequences to nested SP groups under a memory budget. • We design an Executor that runs nested SPtree plans efficiently through inter-level phase streaming and tree-level recomputation. • We implement NSP and show that it consistently outperforms static and dynamic SP baselines on long-tail long-context workloads.
2
Background
2.1
Parallelism for Long-Context Training
Training long-context LLMs is constrained by the device memory needed to hold model states (parameters, gradients, optimizer states) and the activations of long sequences. Two forms of parallelism address these two pressures and are used throughout NSP.
2.2
Sequence Packing
Training corpora contain sequences of widely varying length. A simple way to assemble fixed-shape batches is to pad every sequence to a common length, but padding wastes computation and memory on filler tokens, which is especially severe under the long-tail length distributions of real corpora. Packing [21] concatenates multiple variable-length sequences into a single packed input up to a fixed token budget, without padded tokens. It uses a segmented attention mask and adjusted position indices so that each original sequence is processed independently and tokens from different sequences do not attend to one another. Packing therefore lets a batch be defined by a token budget rather than a fixed number of sequences. It is the standard choice for longcontext training and the default setting throughout this paper; accordingly, NSP treats a batch as a collection of variable-length sequences packed within per-GPU token budgets, rather than as a padded rectangular tensor.
Fully Sharded Data Parallel Plain data parallelism replicates the full model on every GPU and splits the batch across GPUs, synchronizing gradients with an all-reduce. Replicating the model states is wasteful, so FSDP-style training [30, 39] partitions parameters, gradients, and optimizer states across GPUs. Each GPU all-gathers the parameters of a layer on demand during the forward and backward passes and reduce-scatters the corresponding gradients. This reduces the per-GPU model-state cost at the price of extra parameter communication. When each microbatch provides enough computation, parameter allgathers can be prefetched and reduce-scatters can be overlapped with computation from neighboring layers, hiding much of this communication. Beyond memory reduction, FSDP keeps a largely dataparallel execution model and exposes configurable wrapping and sharding choices, making it flexible across model structures and hardware topologies. NSP relies on FSDP to hold model states and treats it as orthogonal to how sequences are parallelized.
2.3
Activation Recomputation
Even with sequence parallelism, the activations stored for the backward pass can dominate memory in long-context training. Activation recomputation (also called gradient checkpointing) trades compute for memory: selected activations are discarded during the forward pass and recomputed on the fly during the backward pass [4, 20]. Recomputation can be applied at different granularities, from
Sequence Parallelism As context length grows, the activations of a single sequence may no longer fit on one GPU, and per-GPU work becomes severely imbalanced across sequences of different lengths. Sequence parallelism shards one sequence along the to3
tokens and FLOPs simultaneously: balancing tokens per GPU leaves FLOPs severely unbalanced, since a single long sequence carries more FLOPs than a batch of short sequences of the same token count. A FLOPs-balanced policy, on the other hand, must pile many short sequences onto the GPUs that host them, placing heavy pressure on activation memory and even causing OOMs.
6K 25
K
8K 12
64
K 32
K 16
8K
4K
2K
1K
2
6
51
25
12
8
Dataset A Dataset B Dataset C
Sequence length (tokens)
(a) Corpus sequence-length distribution. 62%
Tokens
13%
FLOPs
0.0
12%
37%
0.2
0.4 0.6 fraction (per-batch normalized) <16K 32–64K 128–256K 16–32K 64–128K ≥256K
Sequence parallelism was originally introduced to relieve the memory pressure of long sequences, but it also offers an unexpected handle on the balance problem. By sharding a sequence along the token dimension across k GPUs, SP balances both tokens and FLOPs within that sequence across these k GPUs, so the remaining imbalance is confined to the distribution of sequence workloads across SP groups. As the uniform SP degree increases, the best attainable balance improves monotonically (proved in Section A). However, larger SP groups also incur higher communication overhead, much of it wasted on the short sequences that gain little from the extra sharding. This makes balance versus communication a fundamental dilemma for static SP in long-tail training.
13%
44%
0.8
1.0
(b) Token vs. FLOP composition for one representative batch (Qwen3-MoE-30B forward–backward pass). Figure 2 Real corpora are long-tailed (a), and because attention cost is quadratic, even within a single batch the short sequences that dominate the token count contribute little of the compute, while a few long sequences dominate the FLOPs (b). Balancing tokens and FLOPs at once is thus impossible without sharding long sequences.
3.2
whole transformer layers down to individual modules. In long-context training, recomputing attention is much more expensive than recomputing most other modules, so unnecessary attention recomputation directly hurts iteration time. NSP later exploits recomputation at the granularity of its SP structure (Section 4).
3
Motivation
3.1
Observation 1: A Mismatch Between Real-World Sequence Distributions and Static Parallelism
Observation 2: Existing Dynamic SP Mitigates the Tradeoff, But Imposes Costly Restrictions
Prior work proposes dynamic sequence parallelism [34]. The key observation is that a large SP size is only needed for the few sequences that require it. Dynamic SP routes samples of different lengths to SP groups of different sizes, lowering the communication overhead on short sequences while still scaling out long ones. Existing dynamic-SP designs, however, are limited by a strong structural restriction: their SP groups must be disjoint within a single forward–backward pass. FlexSP’s insight is to run short sequences in small, communication-cheap groups and reserve large groups for the few long ones; its MILP solver then assigns samples to balance per-group time. Yet disjointness fixes the GPU partition for the entire pass, so balance and communication pull against each other: tightening balance pushes the plan back toward large, homogeneous groups that lose the small-group savings, unless FlexSP splits the batch into more micro-batches—at the cost of fixed overheads, smaller GEMMs, and weaker communication hiding.
Real-world training corpora are not uniform in length; they follow long-tail distributions. Figure 2a plots the length distribution of three anonymized public sequence-length traces: roughly 94.0% of the samples are shorter than 8K tokens, while sequences longer than 128K account for only 0.12%. Because batches are sampled from these traces, the per-batch length distribution inherits the same long-tail shape. This skew creates a hard balance problem for distributed training. Due to the quadratic cost of attention, even a workload dispatcher [36] cannot equalize 4
3.3
Opportunity: Nesting Different SP Groups in One Pass
NSP Planner (Sec 4.1) Cost Estimation
All of these costs trace back to one assumption: that the SP groups active in a single pass must be disjoint, so each GPU belongs to exactly one group. Disjointness fixes the GPU partition for the whole step and is what forces balance to be restored through extra micro-batches. It also points to the way out—letting SP groups of different sizes share GPUs within one forward–backward pass.
SP-tree Routing
Long-tail samples
SP Tree Topology
Ulysses / Ring / …
Table 1 Notation used throughout Section 4.
Figure 3 gives an overview of NSP. NSP consists of two modular components—an NSP Planner and an NSP Executor— that together adapt the parallel plan to the length distribution of each batch. Given a batch, the Planner (Section 4.1) produces an SP-tree routing plan that assigns each sample to a node of the SP tree. The assignment optimizes a profiled cost-model proxy that combines computation and communication costs. The Executor (Section 4.2) takes the resulting SP tree plan and executes it efficiently with inter-level phase streaming and tree-level recomputation. The rest of this section describes the Planner (Section 4.1) and Executor (Section 4.2).
4.1.1
Tree-level Recomputation
Figure 3 NSP system overview. The Planner routes a long-tail batch onto an SP tree using profiled cost estimates, producing an SP-tree plan. The Executor realizes the plan through inter-level phase streaming and tree-level recomputation while invoking existing SP backends.
Nested Sequence Parallelism
NSP Planner
Inter-level Phase Streaming
SP backend invocations
This nested structure also creates new system-level opportunities. First, communication in one nested group can be overlapped with computation in another nested group. Second, because the sequences assigned to different groups can differ in compute and storage cost, the runtime can choose recomputation policies at group granularity. These opportunities motivate the routing and executor design of NSP (Section 4).
4.1
NSP Executor (Sec 4.2)
Batch
Such nesting promises a better communication– balance tradeoff. Communication stays low where it should: short sequences keep using small groups, while only the few long sequences use large ones. Balance is then recovered per GPU rather than per group, since a GPU can contribute to a large group for long sequences and use the remaining time and memory budget for short sequences in smaller groups. Both happen within a single pass, reducing the need to force balance through extra micro-batches.
4
SP-tree plan
Symbol
Meaning
P p spmax B j Lj T u R(u) |R(u)| a Btok Tokp (a) Cp [ Load(a)
Total number of GPUs (a power of two) GPU index, p ∈ {0, . . . , P − 1} Largest admissible SP degree on the cluster Input batch, B = {1, . . . , M } Sample index, j ∈ B Sequence length of sample j The SP tree Tree node in T GPU-rank set associated with node u SP degree of node u Routing map from samples to SP tree nodes Per-GPU token budget Sharded token load on GPU p Cost-model load on GPU p Cost-model iteration-time proxy for plan a
tion order for each sample in the batch. However, without further structure, this problem is impractical to solve efficiently and implement directly in a distributed training system. To balance expressiveness and implementability, NSP extracts a structured nested execution space represented by an SP tree. Table 1 summarizes the notation used throughout this section. Definition 1 (SP tree). Let spmax = 2K be the largest admissible SP degree. The SP tree T is a complete binary tree with node set {(d, i) | 0 ≤ d ≤ K, 0 ≤ i < 2d }. A node u = (d, i) is associated with the contiguous, aligned rank set R(u) = {r | i2K−d ≤ r < (i + 1)2K−d }.
In its most general form, nested-SP routing would jointly choose the cooperating GPU set and execu-
Each node u defines one admissible SP group with size |R(u)| = 2K−d . A sample assigned to u is 5
SP=8
Leaf to root chain of GPU5
Disjoint SP plan as an antichain SP=4
4-7
0-3
0-1
SP=2
SP=1
Consider an input batch B = {1, . . . , M }, where sample j has length Lj . A routing plan is a mapping a : B → T , where a(j) is the tree node selected for sample j. A plan is evaluated by two quantities: a per-GPU memory load and an overall iteration-time proxy.
0-7
0
2-3
1
2
3
4
The memory load uses the sharded token count, accumulated over the leaf-to-root path containing each GPU: X Lj . Tokp (a) = |R(a(j))|
6-7
4-5
5
6
7
j∈B: p∈R(a(j))
Figure 4 The SP tree for spmax = 8. Each depth fixes an SP degree, and every node u is annotated with its contiguous aligned rank range R(u). A GPU participates in work along one leaf-to-root chain (highlighted for g5 ), nesting SP degrees 8, 4, 2, 1; a disjoint-group SP plan is an antichain of the tree (dashed).
[ The iteration-time proxy Load(a) is a conservative estimate produced by the cost model defined below. The routing objective minimizes it subject to a perGPU token budget: [ min Load(a) a
s.t.
executed with sequence parallelism over the ranks R(u). Leaf nodes correspond to single-GPU execution. A per-batch routing plan maps each sample to one node of T , allowing samples in the same batch to use different SP degrees. When multiple routed nodes are active, a GPU may participate in work along the ancestor chain from its leaf to the root. Figure 4 illustrates these node ranges, leaf-to-root chains, and disjoint-group antichains.
∀p ∈ {0, . . . , P − 1}.
The token constraint acts as a memory safety guard; the implementation sets Btok from the available activation memory and the desired token-balance ratio. Estimating the time load is more involved than in disjoint dynamic-SP systems. A disjoint-group plan places each GPU in a single SP group, so each sample is processed at one SP degree and a single aggregate per-sample cost is enough. The SP tree instead co-locates several SP levels on the same GPUs, so NSP estimates the time load with a hierarchical cost model whose parameters are obtained from offline profiling, built from the SP-backend level up to the full iteration.
The resulting SP-tree routing space is expressive while remaining regular. It generalizes disjointgroup routing, since any disjoint-group plan corresponds to an antichain of the SP tree. This restriction is not merely heuristic: in an abstract homogeneous setting, Section B shows that any laminar power-of-two schedule can be relabeled into an SP tree schedule with the same makespan. On real clusters, the tree also serves as a hardware-conscious placement template, keeping smaller groups within high-bandwidth local domains and using larger groups only when wider cooperation is needed. Together, these properties make the SP-tree space easier to search and execute, while preserving structured schedules and providing a regular hierarchy for system optimizations. 4.1.2
Tokp (a) ≤ Btok ,
SP-backend level. Each backend exposes a phase-
cost interface. For a node u under plan a, the backend bound to u predicts three attention phase costs—a communication prologue Tuprol (a), local attention compute Tucomp (a), and a communication epilogue Tuepil (a)—separately for the forward and backward passes, and we write τu (a) for this phasecost tuple. For each backend, NSP first derives the attention FLOPs and communication volume from the routed inputs. It then combines these quantities with backend-specific offline profiles of compute throughput and communication bandwidth to predict the compute and communication phase times.
Routing Objective and Cost Model
We define the per-batch routing objective over the SP tree, together with the cost model that evaluates a candidate plan. This specifies what the runtime heuristic in Section 4.1.3 optimizes.
Tree level. For a pass σ ∈ {fwd, bwd} and a GPU
p, let Πp (a) be the active nodes on p’s leaf-to-root path. NSP aggregates their backend phase costs into
6
a conservative per-GPU attention time Algorithm 1: Prefix-Completion SP-Tree Routing X Input: Samples B with lengths {Lj }; SP tree T ; Aσp (a) = Tuprol,σ (a) + Tucomp,σ (a) + Tuepil,σ (a) , prefix size k; beam width w; token u∈Πp (a) budget Btok Output: SP-tree plan a where the three terms are the backend prologue, comSort B by nonincreasing sequence length; pute, and epilogue costs for pass σ. H ← first k samples in B; R ← B \ H; Model level. Non-attention work (projections and Q ← {(∅, 0, 0)}; MLP) uses no SP collectives, so it is added on top // Long-prefix tree search of the attention time rather than inside it. From foreach j ∈ H do the token count on p and the profiled dense-compute N ← ∅; throughput, NSP estimates a non-attention time foreach S = (a, C, T) ∈ Q do Opσ (a) and forms the per-GPU stage time Dpσ (a) = foreach u ∈ FeasibleNodes(T , S, j) do Aσp (a) + Opσ (a). a′ ← a ∪ {(j, u)}; C′ , T′ ← CostModel(a′ ); Iteration level. The routing proxy covers the forward if maxp Tp′ ≤ Btok then (fwd) and backward (bwd) stages. Each stage reuses S ′ ← (a′ , C′ , T′ ); the levels above with its own phase costs, and the N ← N ∪ {S ′ }; proxy sums the per-stage critical GPU: Q ← the w states in N with the smallest maxp Cp ;
[ Load(a) = max Dpfwd (a) + max Dpbwd (a). p
p
// Greedy short-tail completion C ← ∅; foreach S ∈ Q do Utail ← effective leaf units in T ; foreach j ∈ R in nonincreasing Lj do u ← least-loaded unit in Utail under S that satisfies Btok ; if no such u exists then u ← unit in Utail with the smallest token load; Extend S by assigning j to u and updating (C, T);
All throughput and bandwidth coefficients are profiled offline per model, deployment environment, backend, and SP degree, and are treated as fixed during routing, so the online loop only combines them with the current length distribution. The model is used only as a routing proxy, not as a standalone latency predictor. 4.1.3
Routing Algorithm
The routing problem in Section 4.1.2 is combinatorially hard to solve exactly. A plan assigns each sample in the batch to one SP tree node, so the number of candidate plans grows exponentially with the batch size, and each assignment couples all GPUs covered by the chosen node through both the estimated load and the token budget. Even with a single fixed SP degree, balancing the per-GPU load already reduces to makespan minimization on parallel machines, which is NP-hard [10]. The full problem is larger, adding a per-sample SP-size choice and the token budget, and its objective is the cost-model load rather than a simple closed-form cost. Because routing runs once per batch inside the training loop, NSP forgoes exact optimization and uses a beam-search routing heuristic that exploits the structure of long-tail batches.
C ← C ∪ {S}; Srst ← arg min(a,C,T)∈C maxp Cp ;
return Srst ;
summaries induced by a: Cp is the load that the cost model attributes to GPU p, and Tp = Tokp (a) is its sharded token load. The cost model summarizes the conservative forward–backward estimate into a single scalar per GPU, while execution-time overlap and tree-level recomputation are left to the Executor. Initially, S0 = (∅, 0, 0). Extending a state with assignment j 7→ u produces a new state S ′ = (a′ , C′ , T′ ), where a′ = a ∪ {(j, u)}. The load summary C′ is updated using the cost-model estimate for placing sample j on node u, while the token-load summary T′ adds the corresponding sharded token load on the GPUs covered by u. A complete plan is feasible if every sample is assigned and every GPU satisfies the
The Planner incrementally constructs an SP-tree plan. We denote a search state by S = (a, C, T). Here a ⊆ B × T is a partial assignment from processed samples to tree nodes. The vectors C = (C0 , . . . , CP −1 ) and T = (T0 , . . . , TP −1 ) are cached 7
token budget. Among feasible plans, the Planner selects the plan with the smallest critical load maxp Cp , a tractable per-GPU surrogate for the cost-model ob[ jective Load(a).
poses an additional overlap dimension: its nodes carry independent attention work once their inputs are ready, so the Executor can overlap one node’s communication with another node’s compute.
To bound the cost of per-batch routing, NSP restricts the candidate space along three dimensions. First, it applies fine-grained search only to the longest prefix of the batch and completes the remaining short tail greedily. This focuses the search on the samples that have the largest impact on the critical path, while relying on coarse placement for the many short samples whose marginal impact on the final iteration time is smaller. Second, it retains a bounded beam of partial plans after each expansion. This preserves multiple residual-capacity profiles without enumerating all prefixes. Third, it processes samples from long to short and can enforce a monotoneSP constraint: if sample j is longer than sample j ′ , then j cannot use a smaller SP degree than j ′ . This concentrates the search on plans consistent with the long-tail structure: long samples tend to use larger SP groups to reduce per-GPU compute and memory pressure, while short samples fill lower levels with less communication.
This inter-level streaming complements, rather than replaces, any overlap a backend already performs internally, and it leaves the backend’s own execution unchanged. It is useful because a backend hides its own communication only partially, if at all: Ulysses [15], for instance, places its all-to-all directly on the attention critical path. Part of the communication therefore tends to stay exposed. The Executor streams the three backend phases of Section 4.1.2—communication prologue, attention compute, and communication epilogue—across the active nodes on a GPU’s path and the available physical resources. Let ui and ui+1 be two consecutive nodes in the Executor’s issue order. The communication prologue of ui+1 can run during the attention compute of ui ; once ui finishes compute, its communication epilogue is enqueued and can overlap with compute from later nodes. The Executor maintains one GPU-compute queue for attention kernels and one communication queue for collective phases. Within a node, the communication prologue, compute, and communication epilogue keep their dependency order across queues; across nodes, each queue follows the issue order, so a later collective does not steal bandwidth from an earlier one on the chosen critical path. The compute and communication queues then proceed concurrently. This streaming is an execution-time optimization: the Planner uses the conservative serial sum in Section 4.1.2, and the Executor then masks the remaining exposed communication through inter-level streaming.
Algorithm 1 summarizes the resulting routing heuristic. The algorithm first separates the batch into a searched prefix and a greedily completed tail. It then performs a level-by-level expansion over the prefix: each retained state represents one residual capacity profile of the SP tree, and each expansion places the next sample on a feasible tree node. Token-infeasible states are filtered immediately, and only the w states with the smallest critical load survive to the next prefix sample. After the prefix search, the heuristic completes each surviving state by packing tail samples into the effective leaf units under the current load and token constraints. The returned plan is the completed state with the smallest critical load.
4.2
4.2.2
Beyond communication, long-context training also suffers from expensive attention recomputation under activation-memory limits, which NSP reduces with a second tree-enabled optimization. Activation checkpointing trades memory for recomputation, and existing policies make the save/recompute decision at layer granularity [4, 20, 33]. This is a poor fit for attention: saving attention outputs raises a decoder layer’s activation memory by more than 1×, which is often infeasible under tight budgets, so systems fall back to full recomputation—yet attention is the most expensive part to recompute for long sequences.
NSP Executor
The Executor turns the populated SP tree from the Planner into a per-node execution plan. It then applies two tree-structured optimizations: inter-level phase streaming and tree-level recomputation. 4.2.1
Tree-Level Recomputation
Inter-level Phase Streaming
NSP reduces exposed communication in two stages. Tree-shaped routing first lowers the communication on the critical path (Section 4.1.1); the Executor then hides the communication that remains through inter-level phase streaming (Figure 5). The tree ex-
The tree-shaped execution plan naturally exposes 8
Computation (long / medium / short seq.) comp
L
comm
Communication
Node split
M
…
S
P
E N1
SP=32
budget cut
Flat comp comm
M
S
SP=8
SP=1
L
P
Nested comp
L
P
N2
less comm
N3
N4
N5
N6
E P SP=32
comm
recompute cost saved bytes recompute node save node
N0
cut
M
P
S
Figure 6 Save/recompute marking for tree-level recomputation. NSP traverses the routed SP tree from upper levels to lower levels, saves the high-return nodes admitted by the activation budget, and recomputes the remaining lower-level nodes during backward.
less exposed comm
E SP=32
SP=1 SP=8
Overlapped
marked recompute are released once they are no longer needed in the forward pass. During backward, NSP recomputes attention only for the nodes marked recompute and combines the recomputed slices with the saved ones to reconstruct the full NSP attention result.
Figure 5 Reducing exposed communication in two stages. Disjoint-group SP runs a long sequence with a large prologue and epilogue. Nested routing first assigns samples to different SP levels, reducing the communication on the critical path. The Executor then streams levels so that communication phases overlap with attention compute from other levels; reordering the issue order further shortens the exposed tail.
The plan is generated consistently across ranks. Under the uniform-token assumption used by NSP slicing, all ranks in an SP group use the same per-node byte estimate and the same budget. The plan is therefore group-uniform: a tree node is either saved by all ranks in the corresponding SP group or recomputed by all of them, so no rank enters a collective that another rank skips. Because recompute is assigned mainly to lower-level short-sequence nodes, the residual recomputation cost is small compared with the original attention workload and has limited impact on the Planner’s routing objective. The Planner therefore need not be made recompute-aware: keeping routing and the save/recompute decision decoupled costs little in practice and keeps the routing objective simple.
level-granular recomputation units. The routed tree usually exhibits a length gradient: long, expensive samples concentrate near upper levels, while short samples populate lower levels. This structure makes the save/recompute decision cost-effective. Upperlevel long samples have the highest recomputation cost per saved byte because recomputation repeats both attention compute and SP communication; at the same time, their attention outputs are sharded across larger SP groups, so the per-GPU storage increase is moderated. NSP therefore spends the activation budget on these high-return levels and recomputes cheaper short-sequence levels during backward.
5
Given this tree order, the Executor derives a save/recompute plan from an activation-memory budget. It traverses the SP tree from the root toward the leaves and accumulates the estimated bytes of saved attention-output slices. Once the budget is reached, the remaining lower levels are marked recompute; the levels already admitted into the budget are marked save. If the budget cut falls inside a tree node, NSP can further split the node by sample subset and save the subset with larger recomputation cost.
Implementation
We implement NSP on top of an internal PyTorchbased distributed training framework that provides FSDP-style model-state sharding. NSP runs above this framework and supplies the sequence-parallel execution layer. We use NCCL for collective communication and FlashAttention kernels for attention computation [5]. SP-tree group management. NSP creates NCCL communication groups dynamically and caches them for reuse, similar to prior dynamic-SP systems. The SP-tree structure does not increase the number of
At runtime, the forward pass extracts and stores only the attention output slices marked save; slices 9
groups compared with disjoint dynamic-SP solutions: all groups are aligned tree nodes, and each GPU participates in at most one group per tree level.
candidate SP levels to only the maximum SP level and the leaf level (SP = 1), disabling all intermediate SP-tree levels. Hardware Environments. We evaluate on an internal
Execution engine. NSP separates the input layout from the execution layout. The input pipeline may deliver packed sequences in a conventional fixed-SP layout, chosen only to satisfy the memory requirement of the longest sample and to remain compatible with upstream modules. Before LLM execution, NSP uses one all-to-all to repartition the same samples according to the current SP-tree routing plan; after LLM execution, a symmetric all-to-all restores the conventional layout for loss computation and downstream modules. This boundary keeps the input pipeline independent of nested SP and lets NSP act as a drop-in execution layer.
64-GPU NVIDIA production cluster. To avoid exposing deployment-specific hardware details, we report normalized performance and omit the specific GPU type. All systems in each experiment run under the same hardware, software stack, model configuration, and input batches. Experimental Workloads. We conduct experiments
on two long-context LLM workloads: Qwen3-MoE30B and Qwen3-MoE-235B [29]. We use three public sequence-length traces from the FlexSP artifact [34], anonymized as Dataset A, Dataset B, and Dataset C. These traces have different sequence-length distributions and contain only sampled sequence lengths, not raw text; we use them solely to reproduce realistic long-tail batch compositions for system-performance evaluation. To evaluate both moderate and extreme long-context settings, we run each workload with maximum context lengths of 192K and 384K tokens. Sequences longer than the configured maximum context length are filtered out during training.
NSP supports multiple SP backends via a common backend interface. Each backend implements forward and backward communication prologue, compute, and communication epilogue routines. NSP composes these routines over the SP tree to enable streamed execution and tree-level recomputation. A backend also exposes cost-model hooks, which NSP queries and aggregates when evaluating a candidate routing plan.
Protocols. We use mixed-precision training with AdamW. Each training batch contains 2M tokens, and all systems use sequence packing. We use an FSDP size of 64, evenly sharding model states across all GPUs. To reduce memory pressure for Qwen3MoE-235B, we use an expert-parallel size of 8 and offload optimizer states to CPU memory. For a fair comparison, we implement all baselines in the same framework so that they share the same model kernels, communication backend, data pipeline, and training configuration unless otherwise stated. We tune each baseline and report its best observed performance. For Static SP, we sweep SP degrees that fit in memory and select the fastest configuration. For FlexSP, we first run offline profiling for each model and provide the profiled parameters to its solver. Although FlexSP can automatically determine microbatch partitioning and SP assignment, we observe that it sometimes over-partitions a batch and degrades performance. We therefore manually sweep the maximum number of micro-batches, up to five as used in the open-source FlexSP implementation, and report the best result. We use the first 50 iterations for warmup, average iterations 50–100, and normalize each result to Static SP under the same configuration.
Planner integration. The online planner runs on CPU and uses the heuristic described in Section 4.1.3. Its runtime is at the millisecond scale in our experiments, so it is negligible compared with a training iteration. Our implementation can prefetch upcoming batch metadata and generate routing plans ahead of execution.
6
Evaluation
6.1
Setup
Baseline Systems. We compare NSP with existing
sequence-parallel training strategies. Static SP follows the DeepSpeed-Ulysses style sequence-parallel execution [15] and uses a single SP degree for all samples. FlexSP [34] is a batch-adaptive disjoint-group baseline: it partitions a batch into micro-batches and uses an MILP solver to select a set of disjoint SP groups for each micro-batch. We reproduce FlexSP’s SP strategy based on its published description and open-source implementation, including micro-batch partitioning and MILP-based SP assignment. To isolate the benefit of full tree nesting, we also introduce NSP-2L, an ablated variant of NSP. NSP-2L uses the same planner and executor as NSP, but restricts the 10
Static SP
FlexSP
NSP-2L
Normalized Time
Qwen3-MoE-235B, Max Seq=384K 1.0
0x x 1.1 .13 7x, x, 1 1.3 1.41
2x x 1.1 1.13 0x, x, 1.2 1.21
Qwen3-MoE-30B, Max Seq=384K 4x x 1.0 .08 2x, x, 1 1.3 1.37
0.5
1.0
6x x 1.0 1.07 8x, x, 1.0 1.09
4x x 1.1 1.16 5x, x, 1.4 1.48
0.0 Dataset A
Dataset B
Dataset C
Dataset A
Qwen3-MoE-235B, Max Seq=192K Normalized Time
4x x 1.1 1.15 4x, x, 1.2 1.25
0.5
0.0
1.0
NSP
4x 1.0 1.10x , 5x, 1.3 1.42x
6x 1.0 1.12x , 0x, 1.3 1.38x
Dataset B
Dataset C
Qwen3-MoE-30B, Max Seq=192K 5x x 1.0 .09 2x, x, 1 1.3 1.37
0.5
1.0
3x x 1.1 1.15 3x, x, 1.4 1.46
7x 7x 1.0 1.0 0x, 0x, 1.2 1.2
1x x 1.1 1.13 4x, x, 1.2 1.26
0.5
0.0
0.0 Dataset A
Dataset B
Dataset C
Dataset A
Dataset B
Dataset C
Figure 7 End-to-end normalized iteration time across model sizes, maximum context lengths, and datasets. Values are normalized to Static SP under the same configuration; lower is better.
6.2
End-to-End Performance
tion and communication at a finer granularity within the same batch.
Figure 7 shows the end-to-end iteration time of NSP and the baselines. The evaluation covers two Qwen3MoE workloads, three long-context datasets, and maximum context lengths of 192K and 384K. Across these settings, NSP consistently achieves the best performance, with up to 1.48× speedup over Static SP and up to 1.16× speedup over FlexSP. This demonstrates that nested sequence parallelism remains effective across diverse workload distributions.
For Qwen3-MoE-235B, Static SP requires attention recomputation under the memory limit, while NSP reduces much of this overhead through tree-level recomputation. Compared with FlexSP, NSP also avoids excessive micro-batch partitioning. This is important for Qwen3-MoE models, where FSDP communication is non-negligible: splitting a batch into many small micro-batches reduces the amount of work available to hide FSDP communication and can lower device utilization.
Static SP performs poorly because a single SP degree cannot match the varied sequence lengths within a batch. A large SP degree can reduce the critical load of long samples by spreading their attention work across more GPUs, but applying the same degree to short samples introduces unnecessary communication. Conversely, choosing a smaller SP degree reduces communication for short samples but leaves long samples with severe load imbalance and, in extreme cases, memory pressure. FlexSP mitigates this problem by forming disjoint SP groups and adapting the SP assignment at the micro-batch level. However, this disjoint-group structure still requires splitting a batch into multiple micro-batches to improve load balance, which reduces the amount of work per execution unit and can lower communication overlap and device utilization. In contrast, NSP places samples with different lengths at appropriate positions in the SP tree. Because the tree allows nested rather than disjoint SP groups, NSP can balance computa-
Finally, NSP-2L confirms that coarse long/short separation is not sufficient. NSP-2L improves over Static SP in many cases by avoiding some unnecessary large-SP communication for short samples, but it remains below full NSP and can be worse than FlexSP when many medium-length samples are present. By disabling intermediate SP levels, NSP-2L must place such samples either at the maximum SP level or at the leaf level, missing the communication–balance trade-off points provided by the full tree. The gap between NSP-2L and NSP therefore isolates the value of complete SP-tree nesting.
6.3
Case Study
Aggregate time breakdown. Figure 8 breaks down the relative step time for a representative longcontext iteration. We include Static SP-32, Static 11
Attn fwd Attn remat
Attn bwd Non-attn
SP exposed comm FSDP exposed comm
Static SP-32 leaves the workload imbalanced, which appears as a long communication-wait component. Static SP-64 improves balance for the longest samples and reduces communication waits, but it forces short samples to communicate across the full SP group, increasing exposed SP communication. FlexSP forms disjoint SP groups of different sizes and uses its MILP solver to assign samples across them so as to balance per-group execution time, confining much of the SP communication to smaller groups and reducing exposed SP communication. It also sorts and splits the batch into three microbatches, lowering the length variance within each and thus improving balance. However, the nonattention compute becomes larger, reflecting the extra overhead introduced by micro-batch splitting.
Comm wait Other exposed
1.00x
Static SP (32)
0.89x
Static SP (64)
0.83x
FlexSP
0.69x
NSP
0.0
0.2
0.4 0.6 0.8 Relative step time
1.0
Figure 8 Relative step-time breakdown for a representative iteration. SP/FSDP communication reports exposed pure communication, and communication waits are grouped separately. Values are normalized to Static SP-32.
SP=64
33%
SP=32
25%
SP=16
NSP’s gain comes from two effects. First, compared with the full attention recomputation that a tight activation budget would otherwise force, tree-level recomputation spends about one quarter more activation storage to avoid roughly two thirds of the recomputation overhead. Second, the routing policy, nested structure, and compute–communication overlap together reduce exposed SP communication. Overall, the breakdown shows that NSP’s end-toend speedup comes from jointly reducing exposed SP communication, communication waits, and recomputation overhead.
0.35
33%
4%
0.30
SP=8 <SP8
9%
30%
11%
30%
SP=8 <SP8
0.25
SP=16 4%
<SP8
11%
SP=8
7%
0.20
4%
<SP8
10%
SP=32
25%
SP=16
4%
SP=8
33%
Fraction
SP=8
0.15
25%
Routing distribution. To further explore NSP’s flexible strategy, Figure 9 visualizes the actual routing plan for the same iteration. Long samples are mostly placed on upper tree nodes such as SP=64 and SP=32, while short samples concentrate on lowerlevel and sub-SP8 nodes to avoid unnecessary largegroup communication. Medium-length samples are spread across SP=16 and SP=8 depending on residual subtree capacity. This distribution shows how NSP uses different SP-tree levels for different samples, achieving good load balance while keeping communication low.
4%
<SP8
8%
SP=8
0.10
4%
<SP8
11%
SP=16
4%
SP=8
25%
4%
<SP8
0.05
14%
SP=8
4%
<SP8
26%
0-8
0.00 8-32
32-128
128+
Sequence length bucket (K tokens)
Figure 9 Node-level routing distribution across the SP tree for a representative NSP iteration. Each column sums to 100% within a sequence length bucket.
6.4
Ablation
Figure 10 reports a leave-one-out ablation of NSP’s main components on Dataset A using the same internal 64-GPU cluster, with Static SP as the reference (dashed line). Starting from full NSP, we disable one component at a time: tree-level recomputation (w/o Recompute), phase-streaming overlap (w/o Overlap), the routing heuristic, which falls back to a greedy assignment (w/o Router), and the full cost model, which is replaced by an attention-
SP-64, FlexSP, and NSP to expose the main tradeoffs behind the aggregate results. All configurations process the same samples with the same model, so the main forward–backward compute is directly comparable. The differences come from recomputation, exposed communication, and communication waits, revealing how effectively each system balances work and hides communication. 12
0.6
0.94x
0.6 Qwen-MoE-235B
Qwen-MoE-30B
Figure 10 Leave-one-out ablation of NSP’s components on Dataset A using the same internal 64-GPU cluster. Bars show iteration time relative to Static SP (dashed line); lower is better, and each bar is annotated with its speedup over Static SP.
7.2
forward-only proxy (w/o Cost Model). Every component contributes positively, but their relative importance shifts with the operating regime, which is in turn determined by model size and memory pressure.
Long-Context Sequence Parallelism
Sequence and context parallelism make long-context training feasible by sharding sequence activations or attention work across GPUs. Megatron-style SP reduces activation memory, DeepSpeed-Ulysses uses all-to-all sequence partitioning, Ring Attention uses ring exchanges, USP combines the two communication patterns, and LoongTrain adds two-dimensional head-context parallelism for long sequences [7, 13, 15, 20, 26]. FlashAttention and FlashAttention-2 reduce attention memory traffic, while DistFlashAttn brings memory-efficient attention to distributed execution with token-level load balancing, overlap, and rematerialization-aware checkpointing [5, 6, 23]. These systems provide the SP/CP backends that NSP reuses, but they generally fix the SP/CP degree or communication pattern for a sample group. NSP instead asks how sequences of different lengths should share one decoder-layer execution, and organizes standard SP backends into nested groups so short sequences avoid the synchronization cost of the longest ones.
On the memory-bound Qwen3-MoE-235B, recomputation is the dominant contributor, with a 9.6% degradation when removed, since NSP’s tree-level recomputation avoids the full attention recomputation that the memory limit would otherwise force. The cost model contributes the least, at 0.7%, because under memory pressure the token-balance constraint already balances the non-MLP cost implicitly, so an attention-forward-only proxy suffices to reach a nearbalanced plan. The compute-bound Qwen3-MoE-30B reverses this ordering. Here the cost model becomes the dominant contributor: removing it degrades iteration time by 33.5% and even falls behind Static SP. With 32 attention heads and a maximum SP degree of 32, the two top-level SP=32 groups are prone to severe imbalance, and the full cost model is needed to jointly balance attention and non-attention (MLP) cost. Disabling recomputation, in contrast, has no measurable effect because the model fits in memory and never triggers recomputation. Overall, these results show that NSP’s components are complementary rather than redundant, with recomputation most critical in the memory-bound regime and the cost model most critical in the compute-bound, balancedominated regime.
7
Basic Parallelisms in LLM Training
Modern LLM training stacks combine sharded data parallelism (e.g., DeepSpeed, ZeRO, and FSDP), tensor/model parallelism, and pipeline parallelism [14, 20, 25, 27, 30, 31, 33, 39]. Large-scale and automatic systems compose or search these axes across device meshes [16, 22, 28, 32, 37, 40], while MoE training adds expert parallelism as another scaling dimension [8, 17]. Hybrid systems further optimize overlap, elasticity, full-stack scaling, pipeline bubbles, and rematerialization/offloading choices [1, 2, 9, 19, 38]. NSP is complementary to these outer parallelism choices: it operates inside sequence-parallel decoderlayer execution, routing variable-length sequences across nested SP groups while leaving data, tensor, pipeline, and expert parallelism unchanged.
1.21x
1.25x
0.8
1.25x
1.40x
0.8
1.30x
1.0 1.35x
1.0 1.29x
1.2
7.1
w/o Router w/o Cost Model
1.01x
w/o Recompute w/o Overlap
1.2
1.41x
Relative Time
Static SP NSP
7.3
Dynamic and Variable-Length Training
Dynamic long-context systems address long-tail imbalance by reassigning samples or micro-batches across ranks. FlexSP uses disjoint SP groups of different sizes; Hydraulis, ByteScale HDP, HotSPa, and multimodal workload dispatchers rebalance variablelength work through micro-batch, data/contextparallel, or hybrid-strategy assignment [11, 12, 24, 34, 36]. Other systems extend variable-length balancing across the training stack: Zeppelin combines hierarchical sequence partitioning with NIC-aware
Related Work 13
routing and module remapping; WLB-LLM balances 4D parallel training with packing and per-document context sharding; and DCP partitions data and computation blocks using input-dependent lengths and attention patterns [3, 18, 35]. These designs mainly balance at the sample or microbatch granularity across ranks. That flexibility can couple the scheduler to the outer training stack: for example, HDP relies on replicated parameters and is limited to ZeRO-1 rather than full parameter sharding, with additional constraints for pipeline parallelism [12]. NSP takes the opposite route: it keeps consecutive layers globally synchronized and confines heterogeneity inside each decoder layer by nesting differently sized SP groups. Thus NSP recovers balance within the layer while remaining compatible with full sharding, pipeline parallelism, and tensor parallelism.
8
Conclusion and Future Work
NSP addresses the communication–balance tradeoff in variable-length long-context LLM training. Instead of assigning one SP degree to an entire batch or partitioning GPUs into disjoint dynamic-SP groups, NSP nests SP groups of different sizes on shared GPUs within a training iteration. This turns SPdegree selection into routing over an SP tree and gives the executor a regular hierarchy for interlevel phase streaming and tree-level recomputation. Across Qwen3-MoE workloads and long-tail datasets on an internal production GPU cluster, NSP consistently improves end-to-end throughput over Static SP and FlexSP. Several directions remain open. First, NSP currently routes each training step independently; extending the planner across gradient-accumulation windows or global batches could further smooth rare longsequence bursts. Second, NSP could be jointly optimized with other parallelization strategies, such as data, tensor, pipeline, and expert parallelism, rather than treating the SP hierarchy in isolation.
14
References
and Intractability: A Guide to the Theory of NP-Completeness. W. H. Freeman, 1979.
[1] Sanjith Athlur, Nitika Saran, Muthian Sivathanu, Ramachandran Ramjee, and Nipun Kwatra. Varuna: Scalable, low-cost training of massive deep learning models. In Proceedings of the European Conference on Computer Systems (EuroSys), pages 472–487, 2022. doi: 10.1145/3492321.3519584.
[11] Hao Ge, Fangcheng Fu, Haoyang Li, Xuanyu Wang, Sheng Lin, Yujie Wang, Xiaonan Nie, Hailin Zhang, Xupeng Miao, and Bin Cui. Enabling parallelism hot switching for efficient training of large language models. In Proceedings of the ACM SIGOPS 30th Symposium on Operating Systems Principles (SOSP), pages 178–194, 2024. doi: 10.1145/3694715. 3695969.
[2] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, and Chao Yang. Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning. In Proceedings of the ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), pages 178–191, 2024. doi: 10.1145/3620666.3651379.
[12] Hao Ge, Junda Feng, Qi Huang, Fangcheng Fu, Xiaonan Nie, Lei Zuo, Haibin Lin, Bin Cui, and Xin Liu. ByteScale: Communication-efficient scaling of LLM training with a 2048K context length on 16384 GPUs. In Proceedings of the ACM SIGCOMM 2025 Conference, pages 963–978, 2025. doi: 10.1145/ 3718958.3754352. [13] Diandian Gu, Peng Sun, Qinghao Hu, Ting Huang, Xun Chen, Yingtong Xiong, Guoteng Wang, Qiaoling Chen, Shangchun Zhao, Jiarui Fang, Yonggang Wen, Tianwei Zhang, Xin Jin, and Xuanzhe Liu. LoongTrain: Efficient training of long-sequence LLMs with head-context parallelism. arXiv preprint arXiv:2406.18485, 2024.
[3] Chang Chen, Tiancheng Chen, Jiangfei Duan, Qianchao Zhu, Zerui Wang, Qinghao Hu, Peng Sun, Xiuhong Li, Chao Yang, and Torsten Hoefler. Zeppelin: Balancing variable-length workloads in data parallel large model training. In Proceedings of the 21st European Conference on Computer Systems (EuroSys), pages 1879–1893, 2026. doi: 10.1145/ 3767295.3769369.
[14] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. GPipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, pages 103–112, 2019.
[4] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016. [5] Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
[15] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. DeepSpeed Ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023.
[6] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pages 16344–16359, 2022.
[16] Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. In Proceedings of Machine Learning and Systems (MLSys), volume 1, pages 1–13, 2019.
[7] Jiarui Fang and Shangchun Zhao. USP: A unified sequence parallelism approach for long context generative AI. arXiv preprint arXiv:2405.07719, 2024.
[17] Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, MarieAnne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
[8] William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. [9] Weiqi Feng, Yangrui Chen, Shaoyu Wang, Yanghua Peng, Haibin Lin, and Minlan Yu. Optimus: Accelerating large-scale multi-modal LLM training by bubble exploitation. In Proceedings of the USENIX Annual Technical Conference (ATC), pages 161–177, 2025.
[18] Chenyu Jiang, Zhenkun Cai, Ye Tian, Zhen Jia, Yida Wang, and Chuan Wu. DCP: Addressing in-
[10] Michael R. Garey and David S. Johnson. Computers
15
put dynamism in long-context training via dynamic context parallelism. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles (SOSP), pages 221–236, 2025. doi: 10. 1145/3731569.3764849.
[26] Hao Liu, Matei Zaharia, and Pieter Abbeel. RingAttention with blockwise transformers for near-infinite context. In International Conference on Learning Representations (ICLR), 2024. [27] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline-parallel DNN training. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139, pages 7937–7947, 2021.
[19] Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, Yulu Jia, Sun He, Hongmin Chen, Zhihao Bai, Qi Hou, Shipeng Yan, Ding Zhou, Yiyao Sheng, Zhuo Jiang, Haohan Xu, Haoran Wei, Zhang Zhang, Pengfei Nie, Leqi Zou, Sida Zhao, Liang Xiang, Zherui Liu, Zhe Li, Xiaoying Jia, Jianxi Ye, Xin Jin, and Xin Liu. MegaScale: Scaling large language model training to more than 10,000 GPUs. In Proceedings of the USENIX Symposium on Networked Systems Design and Implementation (NSDI), pages 745–760, 2024.
[28] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using Megatron-LM. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2021. doi: 10.1145/3458817.3476209.
[20] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. In Proceedings of Machine Learning and Systems (MLSys), volume 5, pages 341–353, 2023.
[29] Qwen Team. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. [30] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2020. doi: 10.1109/SC41405. 2020.00024.
[21] Mario Michael Krell, Matej Kosec, Sergio P. Perez, and Andrew Fitzgibbon. Efficient sequence packing without cross-contamination: Accelerating large language models without impacting performance. arXiv preprint arXiv:2107.02027, 2022. [22] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations (ICLR), 2021.
[31] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), pages 3505– 3506, 2020. doi: 10.1145/3394486.3406703.
[23] Dacheng Li, Rulin Shao, Anze Xie, Eric P. Xing, Xuezhe Ma, Ion Stoica, Joseph E. Gonzalez, and Hao Zhang. DISTFLASHATTN: Distributed memory-efficient attention for long-context LLMs training. In Conference on Language Modeling (COLM), 2024.
[32] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hechtman. Mesh-TensorFlow: Deep learning for supercomputers. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, pages 10414–10423, 2018.
[24] Haoyang Li, Fangcheng Fu, Sheng Lin, Hao Ge, Xuanyu Wang, Jiawen Niu, Jinbao Xue, Yangyu Tao, Di Wang, Jie Jiang, and Bin Cui. Hydraulis: Balancing large transformer model training via codesigning parallel strategies and data assignment. Proceedings of the ACM on Management of Data, 3 (6):1–30, 2025. doi: 10.1145/3769802.
[33] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
[25] Zhuohan Li, Siyuan Zhuang, Shiyuan Guo, Danyang Zhuo, Hao Zhang, Dawn Song, and Ion Stoica. TeraPipe: Token-level pipeline parallelism for training large-scale language models. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139, pages 6543–6552, 2021.
[34] Yujie Wang, Shiju Wang, Shenhan Zhu, Fangcheng Fu, Xinyi Liu, Xuefeng Xiao, Huixia Li, Jiashi Li, Faming Wu, and Bin Cui. FlexSP: Accelerating large language model training via flexible sequence parallelism. In Proceedings of the 30th
16
International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), pages 421–436, 2025. doi: 10.1145/ 3676641.3715998. Open-source: https://github. com/PKU-DAIR/Hetu-Galvatron (branch: flexsp). [35] Zheng Wang, Anna Cai, Xinfeng Xie, Zaifeng Pan, Yue Guan, Weiwei Chu, Jie Wang, Shikai Li, Jianyu Huang, Chris Cai, Yuchen Hao, and Yufei Ding. WLB-LLM: Workload-balanced 4d parallelism for large language model training. In Proceedings of the 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI), pages 785–801, 2025. URL https://www.usenix.org/conference/ osdi25/presentation/wang-zheng. [36] Bangjun Xiao, Yijie Zheng, Lei Shi, Xiaoyang Li, Faming Wu, Tianyu Li, Xuefeng Xiao, Yang Zhang, Yuxuan Wang, and Shouda Liu. OrchMLLM: Orchestrate multimodal data with batch postbalancing to accelerate multimodal large language model training. arXiv preprint arXiv:2503.23830, 2025. [37] Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. GSPMD: General and scalable parallelization for ML computation graphs. arXiv preprint arXiv:2105.04663, 2021. [38] Tailing Yuan, Yuliang Liu, Xucheng Ye, Shenglong Zhang, Jianchao Tan, Bin Chen, Chengru Song, and Di Zhang. Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism. In Proceedings of the USENIX Annual Technical Conference (ATC), pages 545–561, 2024. [39] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. PyTorch FSDP: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16(12):3848–3860, 2023. doi: 10.14778/3611540.3611569. [40] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P. Xing, Joseph E. Gonzalez, and Ion Stoica. Alpa: Automating inter- and intra-operator parallelism for distributed deep learning. In Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI), pages 559–578, 2022.
17
Appendix A
lighter half: Tk = S(k) ∪ S(m+k) ,
Pairing Proof for Static SP Imbalance
where S(k) is the set of sequences assigned to the kth largest group. Each new group is feasible because its token load is at most sC + sC = 2sC. Its work is W(k) + W(m+k) , and the average work per new group is 2W̄ .
This appendix formalizes the pairing argument used in Section 3: increasing a uniform SP degree reduces the load-imbalance contribution caused by rare long sequences. The argument is best understood after collapsing each fixed-size SP group into one scheduling unit. Suppose a cluster has P GPUs and each GPU has token capacity C. Under a static SP degree s, the cluster contains D = P/s identical SP groups, each with aggregate token capacity sC. Let Wq denote the aggregate attention work assigned to group q. Assuming that tokens and attention work are evenly partitioned within each SP group, a GPU in group q receives work Wq /s, while the global average per-GPU work is W̄ /s. Thus the per-GPU imbalance ratio is exactly the same as the imbalance across collapsed SP groups: Rs =
maxq Wq , W̄
Because both terms are sorted in non-increasing order, the largest paired load in this constructed assignment is W(1) + W(m+1) . Hence the constructed assignment has imbalance R′ =
1 X Wq . D q=1
Substituting into R′ gives W(1) + W(m+1) − 2W̄ 2W̄ (1 + ρ)(W(1) − W̄ ) 1+ρ = = (Rs − 1). 2 2W̄
R′ − 1 =
The proof therefore isolates the inter-group load imbalance caused by static SP assignment; additional intra-group or system-overhead imbalance is outside this idealized bound.
Therefore R2s − 1 ≤ (1 + ρ)(Rs − 1)/2.
Before stating the pairing bound, sort the group loads as W(1) ≥ · · · ≥ W(D) . For an even number D = 2m of groups and W(1) > W̄ , define ρ=
This bound gives a monotonicity result without any distributional assumption: since W(m+1) ≤ W(1) , we have ρ ≤ 1, and therefore R2s ≤ Rs . If W(m+1) < W(1) , then ρ < 1 and the same pairing argument proves strict improvement. In long-tail batches, the heavier half of the SP groups is typically dominated by a small number of long, high-density sequences, while the lighter half consists mostly of short-sequence groups. When this separation gives ρ ≤ ρ0 < 1, each SP doubling contracts the excess imbalance by a factor at most (1 + ρ0 )/2. The special case W(m+1) ≤ W̄ has ρ ≤ 0 and recovers the intuitive “halving” bound.
W(m+1) − W̄ . W(1) − W̄
This parameter lies in the closed interval ρ ∈ [−1, 1]. The upper bound follows from W(m+1) ≤ W(1) . For the lower bound, 2mW̄ =
2m X
W(i) ≤ mW(1) + mW(m+1) ,
i=1
which implies W(m+1) − W̄ ≥ −(W(1) − W̄ ). Lemma 1 (Pairing bound for SP doubling). Under the notation above, there exists a feasible assignment at SP degree 2s whose imbalance ratio R2s satisfies R2s − 1 ≤
W(1) + W(m+1) . 2W̄
Since R2s is the best imbalance attainable at degree 2s, it is no worse than this feasible construction: R2s ≤ R′ . It remains to rewrite the excess imbalance of R′ . From the definition of ρ, W(m+1) = W̄ + ρ W(1) − W̄ .
D
W̄ =
k = 1, . . . , m,
B
From Feasible Solutions to SP-Tree Solutions
This appendix justifies the tree restriction used in Section 4.1.1 at the level of whole schedules. In an abstract homogeneous setting, any feasible malleable schedule whose cooperation groups are laminar and have power-of-two sizes can be relabeled into one
1+ρ (Rs − 1). 2
Proof. Construct a degree-2s assignment by pairing the heavier half of the degree-s groups with the 18
that uses only SP tree nodes without increasing the makespan. We model the problem on P = 2K workers labeled 0, . . . , P − 1. Each job j runs on a worker group Gj over some time interval, and a schedule is feasible if at every instant the groups of concurrently running jobs are pairwise disjoint. A schedule is a tree schedule if every group Gj equals the worker set R(uj ) of some SP tree node uj .
B.1
Because workers are homogeneous, π preserves every job’s running time. As a bijection, π keeps disjoint groups disjoint, so concurrently running jobs remain conflict-free at every instant and feasibility is preserved. The relabeled schedule keeps all timings, so its makespan equals the original. The laminar condition is tight. The crossing groups {0, 1} and {1, 2} both have power-of-two size, yet they overlap without nesting; no relabeling can make both aligned dyadic intervals, because a bijection preserves partial overlap while any two SP tree nodes are nested or disjoint. This crossing configuration is exactly what the laminar assumption rules out. The conversion relabels workers and thus relies on homogeneity; across a hierarchical inter-node fabric the relabeling should be read up to the bandwidth domain within which workers are interchangeable.
Laminar Groups on Homogeneous Workers
We assume the running time of a job depends only on the number of workers |Gj |, not on their identities (e.g., GPUs fully connected by NVLink within a node). A family of groups is laminar if any two of its members are either nested or disjoint. Lemma 2 (Dyadic packing). A multiset of block sizes, each a power of two and summing to 2k , can be placed as pairwise-disjoint aligned dyadic intervals that tile [0, 2k ). Proof. Induction on k. For k = 0 the sum is 1 and the single unit block tiles [0, 1). For k ≥ 1, every block of size larger than one is even, and the total 2k is even, so the number of unit blocks is even; pair them into size-2 units. All sizes are now even, and halving them yields a multiset of powers of two summing to 2k−1 , which tiles [0, 2k−1 ) by induction. Mapping each interval [a, b) 7→ [2a, 2b) and splitting each paired unit back into two unit blocks tiles [0, 2k ). Theorem 1. On P = 2K homogeneous workers, let a feasible schedule satisfy (i) every group has powerof-two size and (ii) the family of all job groups is laminar. Then there is a relabeling of workers under which the schedule becomes a tree schedule with the same makespan. In particular, an optimal such schedule admits an optimal tree schedule. Proof. Add the full set and all singletons to the group family; it remains laminar, and its inclusion order is a forest rooted at the full set. Process the sets top down. For a set S of size 2k , its maximal proper subsets in the family are disjoint and have power-of-two sizes; together with the elements of S that lie in no such subset (treated as unit blocks), their sizes sum to 2k . By Lemma 2 they tile the 2k block assigned to S, placing each child in an aligned dyadic sub-block. Recursing yields a global relabeling π that maps every group to an SP tree node. 19