Развязанный Улисс: Эффективное использование памяти при параллелизме контекста через поблочное разделение по головам

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

February 24, 2026

Авторы: Ravi Ghadia, Maksim Abraham, Sergei Vorobyov, Max Ryabinin

cs.AI

Аннотация

Эффективная обработка длинных последовательностей моделями Transformer обычно требует распределения вычислений по акселераторам с помощью контекстного параллелизма. Доминирующие подходы в этом семействе методов, такие как Ring Attention или DeepSpeed Ulysses, позволяют масштабироваться по размерности контекста, но не ориентированы на эффективность использования памяти, что ограничивает длину последовательностей, которые они могут поддерживать. Более продвинутые техники, такие как Fully Pipelined Distributed Transformer или выгрузка активаций, могут дополнительно увеличить возможную длину контекста ценой пропускной способности обучения. В данной статье мы представляем UPipe — простой, но эффективный метод контекстного параллелизма, который выполняет мелкозернистое разбиение на уровне головок внимания. Этот метод значительно сокращает использование памяти активациями самовнимания, преодолевая барьер памяти активаций и открывая возможность работы с гораздо более длинными контекстами. Наш подход сокращает использование памяти промежуточными тензорами в слое внимания до 87,5% для 32-миллиардных Transformers, при этом соответствуя предыдущим методам контекстного параллелизма по скорости обучения. UPipe может поддерживать контекст длиной в 5 миллионов токенов при обучении модели Llama3-8B на одном узле с 8xH100, что более чем на 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%.