HexiSeq: Accommodating Long Context Training of LLMs over Heterogeneous Hardware

TL;DR

HexiSeq is a system designed to enable efficient long-context training of large language models on heterogeneous GPU clusters by optimizing partitioning and scheduling based on device capabilities.

Contribution

It extends existing parallelism methods to heterogeneous hardware, formalizes the allocation as an optimization problem, and develops a scheduler for improved throughput.

Findings

HexiSeq improves training throughput by up to 1.72x on heterogeneous clusters.

It achieves near-homogeneous performance on mixed GPU setups.

HexiSeq effectively supports models from 3B to 70B parameters with long context lengths.

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.