π€ AI Summary
This work addresses the challenge that existing long-context training systems for large language models rely on homogeneous GPU clusters, which poorly match real-world heterogeneous environments featuring mixed GPU models and non-uniform network bandwidth. To overcome this limitation, the authors propose a fully asymmetric joint partitioning strategy that integrates context parallelism (CP) and attention head parallelism (HP), dynamically allocating sequence shards and attention heads based on each deviceβs compute, memory, and communication capabilities. A constraint-optimization-based hierarchical scheduler is introduced to automatically discover the optimal parallel configuration. Evaluated across models ranging from 3B to 70B parameters with context lengths up to one million tokens, the system achieves an average throughput improvement of 1.11β1.36Γ (up to 1.72Γ) over state-of-the-art baselines, closely approaching the performance of homogeneous clusters under comparable FLOP budgets.
π Abstract
Long-context training of large language models (LLMs) is commonly distributed with Context Parallelism (CP) and Head Parallelism (HP), but existing training systems largely assume homogeneous GPU meshes. This paper extends CP and HP to heterogeneous GPU clusters with mixed GPU models and non-uniform network bandwidths, a common setting in production training. We introduce HexiSeq, a system that supports fully asymmetric CP--HP partitioning by assigning sequence shards and attention heads according to device compute, memory, and communication capabilities. We formalize heterogeneous CP--HP allocation as a constrained optimization problem and develop an efficient hierarchical scheduler for finding optimal schedules. We evaluate HexiSeq against state-of-the-art CP and HP baselines on both real and simulated heterogeneous clusters. Across models from 3B to 70B parameters and context lengths up to one million tokens, HexiSeq improves throughput by $1.11\times$ on average and up to $1.19\times$ on mixed H100--A100 testbeds, and by $1.36\times$ on average and up to $1.72\times$ in simulations with 32--128 GPUs spanning up to four GPU models. On FLOP-comparable pairs against homogeneous clusters, HexiSeq reaches throughput close to the strongest homogeneous baseline, showing that heterogeneous clusters can be used efficiently for long-context LLM training.