《解缚尤利西斯：基于头向分块的内存高效上下文并行方法》

摘要

在Transformer模型中高效处理长序列通常需要通过上下文并行化将计算任务分配到多个加速器上。该类方法的主流方案（如环形注意力或DeepSpeed Ulysses）虽能实现上下文维度的扩展，但未聚焦内存效率问题，从而限制了可支持的序列长度。更先进的技术（如全流水线分布式Transformer或激活值卸载）虽能进一步扩展上下文长度，但会以降低训练吞吐量为代价。本文提出UPipe——一种在注意力头层级进行细粒度分块的简洁而高效的上下文并行技术。该技术显著降低了自注意力机制的激活内存占用，突破了激活内存瓶颈，从而支持更长的上下文长度。在32B参数规模的Transformer模型中，我们的方法将注意力层的中间张量内存占用降低了87.5%，同时保持了与既有上下文并行技术相当的训练速度。在单台8×H100节点上训练Llama3-8B模型时，UPipe可支持500万标记的上下文长度，较现有方法提升超过25%。

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 8timesH100 node, improving upon prior methods by over 25%.

《解缚尤利西斯：基于头向分块的内存高效上下文并行方法》

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

摘要

Support