Nested sequence parallelism speeds up training of long context large language models
NSP: Accelerating Variable-Length LLM Training via Nested Sequence Parallelism
Distributed, Parallel, and Cluster Computing
Summary
Training large language models with very long and varied-length text sequences is slow because computers must balance work fairly while communicating enough data. The authors present NSP, a new method that groups sequences by length in a nested way on shared GPUs, so longer sequences get more resources and shorter ones get less, reducing wasted communication. This leads to faster overall training without changing the model itself. Tests on long-text datasets show NSP outperforms previous methods in training speed.
What this means in practice
- •For machine learning engineers: Improve training speed of large language models with very long text inputs by balancing workload efficiently across GPU resources.
- •For cloud gpu infrastructure teams: Optimize GPU cluster usage during training of heterogeneous sequence-length language model workloads to raise throughput without needing model changes.
Authors
Yi'ou Wang, Xiaoyang Li, Yijie Zheng, Shouda Liu, Yuxuan Wang
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-to-end training throughput, outperforming Static SP by up to 1.48x and FlexSP by up to 1.16x.