AdaHC: Accelerating Multi-Token Prediction with Adaptive Head Chunking with Pipeline Parallelism
Abstract
Lay Summary
Modern large language models often use multi-token prediction (MTP) to learn more efficiently by predicting several future words at once. However, when training these models across multiple GPUs using pipeline parallelism, MTP unintentionally creates computational bottlenecks: the extra prediction work piles up on the last GPU, causing others to wait idle and slowing down training. We introduce AdaHC, a new scheduling method that solves this problem by recognizing that MTP prediction tasks are actually independent and can be flexibly rearranged. AdaHC automatically splits these tasks into smaller, equally-sized chunks and redistributes them across all GPUs—like rebalancing work across an assembly line—while using lightweight communication to ensure the model learns exactly the same way. Evaluated on state-of-the-art models, AdaHC improves training speed by 1.35× on average without changing model architecture or accuracy, offering a practical path toward faster, more scalable LLM training.