AdaHC: Accelerating Multi-Token Prediction with Adaptive Head Chunking with Pipeline Parallelism

Multi-token prediction (MTP) architecture is widely adopted in LLMs. MTP blocks can be appended to the tail of model to predict additional tokens. However, when training with pipeline parallel, MTP leads to more pipeline bubbles and deteriorates the pipeline efficiency. Based on in-depth analysis of MTP architectures and loss functions, we have identified the parallel nature of the MTP blocks, and leverage it for superior pipeline scheduling. We propose AdaHC, an adaptive pipeline scheduling framework for accelerating LLMs training with MTP block(s). AdaHC splits the output heads into chunks and reassembles the chunks to generate balanced pipeline stages, and performs adaptive activation forwarding to preserve the numerical equivalence. Experimental results show that AdaHC improves the training throughput of SOTA LLMs with diverse MTP configurations by 1.35$\times$ on average. This work paves a new direction for practical pipeline training.

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.