处理Transformer模型的长序列通常需要通过上下文并行化将计算拆分到多个加速器上。现有的主流方法，如Ring Attention或DeepSpeed Ulysses，虽然能在上下文维度上实现扩展，但并未专注于内存效率，这限制了它们所能支持的序列长度。更先进的技术，如完全流水线分布式Transformer或激活卸载，可以进一步扩展可能的上下文长度，但会牺牲训练吞吐量。

本文提出了UPipe，这是一种简单而有效的上下文并行技术，它在注意力头级别执行细粒度分块。该技术显著降低了自注意力的激活内存使用量，突破了激活内存的瓶颈，从而能够支持更长的上下文长度。我们的方法在注意力层中减少了中间张量的内存使用，同时与先前的上下文并行技术在训练速度上保持一致。

UPipe在单个节点上训练模型时，能够支持更长的上下文长度，相比先前的方法有显著提升。这种方法通过精细的内存管理，在扩展序列处理能力的同时，避免了训练吞吐量的损失，为处理超长序列提供了更高效的解决方案。