Как работает sequence parallelism в контексте LLM?

Краткий тезис

Sequence parallelism — это метод распределённого обучения, при котором длинная входная последовательность разрезается на части (chunks) между несколькими GPU. Каждый GPU обрабатывает свою часть последовательности, а для обмена скрытыми состояниями между соседними частями используется ring attention или его вариации. Этот подход необходим для обучения LLM на ultra-long context (>100k токенов), когда полная последовательность не помещается в память одного GPU. Sequence parallelism часто комбинируется с tensor parallelism и pipeline parallelism для построения эффективной 3D-параллельной стратегии.

1. Проблема: длинные последовательности и ограничения GPU

Современные LLM (GPT-4, Llama 3, Gemini) обучаются на контекстах длины 128k, 1M и более токенов. При обучении с полным вниманием (full attention) каждый слой вычисляет матрицу внимания размером (seq_len × seq_len). Для seq_len = 1M это триллионы операций и сотни гигабайт памяти только для одного тензора.

Основные ограничения одного GPU

• ОЗУ (VRAM): даже H100 (80 ГБ) не вмещает 1M токенов с полным attention.
• Вычислительная нагрузка: GPU|один GPU будет считать слишком долго.
• Пропускная способность между GPU: требуется эффективный обмен данными.

Традиционные подходы (Data Parallel|data parallelism, tensor parallelism) не решают проблему длины напрямую:

• Data parallelism копирует модель на каждый GPU и раздаёт разные батчи, но каждый GPU всё равно видит полную последовательность.
• Tensor parallelism разрезает слои по скрытым размерностям, но не по длине последовательности.

Sequence parallelism вводит новое измерение разрезания — по оси последовательности.

2. Базовый принцип: разрезание последовательности

Пусть входная последовательность длины L разбивается на P частей (по числу GPU). Каждый GPU получает chunk размера L/P.

Обозначения

• GPU i получает токены от i * (L/P) до (i+1) * (L/P) - 1.
• Каждый GPU хранит свою часть входных эмбеддингов, скрытых состояний и градиентов.

Проблема для вычисления self-attention токену на GPU i нужно обращаться к токенам на других GPU (полное внимание требует глобальной информации). Простое решение — обмениваться всеми скрытыми состояниями между всеми GPU, но это приводит к квадратичному росту коммуникаций (каждый GPU отправляет свои состояния всем остальным) и узким горлышкам.

3. Ring Attention — эффективная коммуникация

Ring Attention организует GPU в логическое кольцо. Каждый GPU передаёт свой chunk следующему по кольцу, а принимает от предыдущего. На каждом шаге GPU вычисляет частичное внимание, используя свои локальные состояния и полученные состояния от соседа.

Основные шаги ring attention для одного слоя:

1. Каждый GPU загружает свой chunk Q, K, V (query, key, value) из локальной памяти.
2. GPU вычисляет начальную часть attention (локальный score = Q @ K.T).
3. GPU передаёт свой chunk K и V следующему по кольцу (например, через NCCL).
4. Принимает chunk K и V от предыдущего GPU.
5. Добавляет вклад от нового chunk (softmax с коррекцией) и обновляет выход O.
6. Повторяет шаги 3-5, пока не обойдёт все GPU (P-1 обменов).

Ключевые оптимизации

• Коммуникации перекрываются с вычислениями: пока GPU считает attention для текущего chunk, следующий chunk уже передаётся.
• Нет необходимости в all-to-all: каждый GPU общается только с соседями, сложность O(P) вместо O(P²).
• Поддержка очень длинных последовательностей: при фиксированном числе GPU память растёт линейно с L, а время — пропорционально L (при фиксированном L/P).

4. Сравнение с другими видами параллелизма

Важно sequence parallelism не заменяет другие виды, а дополняет их — при обучении GPT-4 с 1M контекстом использовалась 3D-конфигурация: tensor parallelism (внутри узла) + sequence parallelism (между узлами для длинных последовательностей) + pipeline parallelism (между стойками).

5. Практическая реализация: псевдокод ring attention

import torch
import torch.nn.functional as F
from torch.distributed import ring_send, ring_recv

def ring_attention(q_chunk, k_chunk, v_chunk, rank, world_size):
    """
    q_chunk, k_chunk, v_chunk: локальные части attention (batch, num_heads, local_len, d_head)
    """
    local_len = q_chunk.size(2)
    output = torch.zeros_like(q_chunk)
    # Инициализация softmax denominator
    lse = torch.full((q_chunk.size(0), q_chunk.size(1), local_len), float('-inf'))
    normalizer = torch.zeros_like(lse)

    # Первый шаг: локальный attention
    attn_scores = torch.matmul(q_chunk, k_chunk.transpose(-2, -1))  # (batch, heads, local_len, local_len)
    attn_weights = F.softmax(attn_scores, dim=-1)
    output += torch.matmul(attn_weights, v_chunk)
    # Обновляем lse и normalizer для глобального softmax
    lse = attn_scores.max(dim=-1).values.unsqueeze(-1)
    exp_scores = torch.exp(attn_scores - lse)
    normalizer = exp_scores.sum(dim=-1)

    # Цикл по остальным GPU
    recv_k = torch.empty_like(k_chunk)
    recv_v = torch.empty_like(v_chunk)
    for step in range(1, world_size):
        # Отправляем свои K,V следующему GPU, получаем от предыдущего
        next_rank = (rank + 1) % world_size
        prev_rank = (rank - 1) % world_size
        send_req = ring_send(k_chunk, next_rank)
        recv_req = ring_recv(recv_k, prev_rank)
        # Асинхронная передача/приём
        send_req.wait()
        recv_req.wait()
        # Аналогично для V
        ring_send(v_chunk, next_rank).wait()
        ring_recv(recv_v, prev_rank).wait()

        # Вычисляем attention для полученного chunk
        attn_scores = torch.matmul(q_chunk, recv_k.transpose(-2, -1))
        local_lse = attn_scores.max(dim=-1).values.unsqueeze(-1)
        exp_scores = torch.exp(attn_scores - local_lse)
        local_normalizer = exp_scores.sum(dim=-1)

        # Слияние с предыдущим с помощью online softmax
        new_lse = torch.max(lse, local_lse)
        rescale_factor = torch.exp(lse - new_lse)
        lse = new_lse + torch.log(
            torch.exp(lse - new_lse) + torch.exp(local_lse - new_lse)
        )
        normalizer = normalizer * rescale_factor + local_normalizer * torch.exp(local_lse - new_lse)
        output = output * rescale_factor + torch.matmul(attn_weights, recv_v)

    # Финальная нормализация (optional, обычно делают позже в слое)
    output = output / normalizer.unsqueeze(-1)
    return output

Пояснения

• Используется online softmax, чтобы избежать двойного логарифмирования и численной ошибки.
• Передача K и V делается асинхронно (перекрытие с вычислениями).
• На практике в фреймворках (Megatron-LM, DeepSpeed, vLLM) используется оптимизированная версия с ядрами CUDA.

6. Комбинирование sequence parallelism с другими методами

3D параллелизм (3D parallelism) — стандарт для обучения больших LLM:

1. Tensor parallelism (TP) — внутри узла (NVLink). Разрезает скрытые размерности, уменьшает память под веса.
2. Sequence parallelism (SP) — между узлами (InfiniBand). Разрезает последовательность, уменьшает память под attention.
3. Pipeline parallelism (PP) — между стойками (Ethernet). Разрезает по слоям, уменьшает буфер активаций.

Пример конфигурации для 100k контекста на 64 GPU:

• TP=8 (один узел из 8 GPU, H800 с NVLink)
• SP=4 (четыре узла работают как кольцо, каждый узел обрабатывает 25k токенов)
• PP=4 (четыре стадии по 16 GPU каждая)

Важно SP и TP не конфликтуют, так как разрезают разные оси. TP внутри узла позволяет эффективно пересылать данные через NVLink, а SP между узлами — через InfiniBand. PP же добавляет асинхронность и уменьшает пиковую память.

7. Современные альтернативы и улучшения

Когда SP не нужен

• Длина контекста < 16k (влезает в один GPU, достаточно TP+DP)
• Разреженное внимание (Sparse Attention) — эффективнее использовать local attention
• Приоритет latency (инференс) — обычно используют tensor parallelism + speculative decoding

8. Производительность и ограничения

Метрики эффективности

• Memory reduction: теоретически в P раз меньше памяти для K/V cache и attention матриц. На практике из-за on-chip memory и коммуникационных буферов — ~0.8P.
• Scalability: при P=64 и L=1M достигается почти линейный прирост throughput (92-95%) с ring attention.
• Bubble при pipeline parallelism частично сглаживается за счёт микро-батчей.

Главное ограничение — пропускная способность сети.

• Для 1M токенов на 64 GPU с H800 (NVLink 900 GB/s, InfiniBand 400 Gb/s) коммуникация занимает ~40% времени.
• Решение: перекрытие коммуникаций с вычислениями (overlap), использование более быстрых сетей (NVLink 5.0, CX-8).

Пет-проект для закрепления

Задача Реализовать упрощённый sequence parallelism для обучения маленького transformer на длинных последовательностях (имитация нескольких GPU с помощью torch.distributed на одной машине).

Инструменты PyTorch, torch.distributed, NCCL (или Gloo), Python.

Шаги:

1. Создать простую модель (один слой self-attention, один feed-forward).
2. Написать функцию ring_attention как выше.
3. Инициализировать process group (world_size=4 с помощью torchrun --nproc_per_node=4 ...).
4. Сгенерировать случайную последовательность длины 4096 и разрезать на 4 части по 1024.
5. Для каждого forward pass:
   • Каждый worker вычисляет свою часть.
   • Вызвать ring_attention.
   • Выполнить всё остальное (LayerNorm, FFN) локально.
6. Посчитать loss, backward с синхронизацией градиентов (all-reduce).
7. Сравнить результаты с полным attention на одном GPU (без SP) — убедиться в одинаковости loss.

Ожидаемый результат Понимание, как именно данные передаются между GPU, почему необходима синхронизация, и как растёт скорость при увеличении числа GPU.

Связь с другими вопросами

Навигация

• Предыдущий: 424
• Следующий: 426
• Индекс: 00. Индекс разборов