ATAC: Anchor-tail aware context parallelism for LLM training
TL;DR
This work proposes ATAC, an anchor-tail aware framework for jointly optimizing packed-document construction and context-parallel sharding in large language model training and demonstrates that input-structure-aware co-design of packing and context parallelism is an effective approach to improving calibrated pipeline-level execution efficiency in large language model training.
Abstract
Large language model training commonly relies on multidimensional parallelism, including data, tensor, pipeline, and context parallelism, to support long-context and large-scale workloads. However, real pretraining corpora consist of highly heterogeneous variable-length samples, which create a complex coupling between the internal structure of packed documents and their actual execution cost. Existing approaches typically optimize upstream packing and downstream context parallelism separately, while paying limited attention to their coupled impact on execution block completion time, local load balance, and communication overhead. To address this issue, we propose ATAC, an anchor-tail aware framework for jointly optimizing packed-document construction and context-parallel sharding in large language model training. ATAC consists of two complementary components: WFAP, which constructs execution-friendly packed documents by jointly considering workload structure and sample fragmentation during packing, and ATP-CP, which performs hierarchical sharding by exploiting the anchor-tail structure of packed documents so as to mitigate intra-block stragglers while controlling additional key-value transfer overhead. In an A800-calibrated pipeline-level evaluation, ATAC improves normalized throughput over the strongest implemented baseline, achieving an overall geometric mean speedup of 1.93×\documentclass[12pt]{minimal} \usepackage{amsmath} \usepackage{wasysym} \usepackage{amsfonts} \usepackage{amssymb} \usepackage{amsbsy} \usepackage{mathrsfs} \usepackage{upgreek} \setlength{\oddsidemargin}{-69pt} \begin{document}$$\times $$\end{document}. Further analysis shows that WFAP substantially reduces runtime dispersion across packed documents, while ATP-CP provides a more effective balance between runtime equalization and communication overhead. Overall, this work demonstrates that input-structure-aware co-design of packing and context parallelism is an effective approach to improving calibrated pipeline-level execution efficiency in large language model training.