通过头级分块实现高效长序列训练，内存占用降低87.5%
Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking
- 在注意力头级别进行细粒度分块，减少激活内存
- 320亿模型激活内存最多降87.5%，训练速度不变
- 单卡支持500万词元上下文，比之前提升超25%
高效处理长序列的Transformer模型通常依赖于跨加速器的上下文并行。当前主流方法如Ring Attention或DeepSpeed Ulysses虽可扩展上下文维度，但未关注内存效率，限制了可支持的序列长度。更先进的技术如全流水分布式Transformer或激活卸载虽能进一步延长上下文，但以牺牲训练吞吐为代价。本文提出UPipe，一种简单而高效的上下文并行技术，在注意力头级别进行细粒度分块，显著降低自注意力的激活内存，突破激活内存瓶颈，实现更长上下文。该方法使32B Transformer的注意力层中间张量内存使用最多降低87.5%，且训练速度与此前方法相当。在单个8×H100节点上训练Llama3-8B时，可支持500万词元上下文，较之前方法提升超过25%。
原文摘要 · Abstract (English)
Efficiently processing long sequences with Transformer models usually requires splitting the computations across accelerators via context parallelism. The dominant approaches in this family of methods, such as Ring Attention or DeepSpeed Ulysses, enable scaling over the context dimension but do not focus on memory efficiency, which limits the sequence lengths they can support. More advanced techniques, such as Fully Pipelined Distributed Transformer or activation offloading, can further extend the possible context length at the cost of training throughput. In this paper, we present UPipe, a simple yet effective context parallelism technique that performs fine-grained chunking at the attention head level. This technique significantly reduces the activation memory usage of self-attention, breaking the activation memory barrier and unlocking much longer context lengths. Our approach reduces intermediate tensor memory usage in the attention layer by as much as 87.5$\%$ for 32B Transformers, while matching previous context parallelism techniques in terms of training speed. UPipe can support the context length of 5M tokens when training Llama3-8B on a single 8$\times$H100 node, improving upon prior methods by over 25$\%$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。