Параллелизм последовательности Ulysses

Ulysses sequence parallelism (SP) trains on very long sequences by splitting them across multiple GPUs. To compute attention correctly, an all-to-all collective swaps the sharding dimension from sequence to attention heads. Each GPU then has the full sequence and computes attention locally over a subset of heads. A second all-to-all returns to the sequence-sharded layout so the rest of the forward pass continues locally on each chunk.

                        GPU 0                       GPU 1
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← each GPU holds half the sequence
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  attention        │ all N tokens  │          │ all N tokens  │  ← now each GPU has the full sequence
  (head-sharded)   │ heads 0..H/2  │          │ heads H/2..H  │  ← but only half the heads
                   └───────┬───────┘          └───────┬───────┘
                           └───────── all-to-all ──────┘
                   ┌───────────────┐          ┌───────────────┐
  forward          │ tokens 0..N/2 │          │ tokens N/2..N │  ← back to seq-sharded
  (seq-sharded)    │  all H heads  │          │  all H heads  │
                   └───────────────┘          └───────────────┘

В этом руководстве рассматривается компонент параллелизма последовательности Ulysses в ALST (Arctic Long Sequence Training). Полная система ALST также включает TiledMLP и выгрузку контрольных точек активаций, которые недоступны в Transformers. См. учебник DeepSpeed ALST для описания полной системы.

Настройка

Параллелизм последовательностей требует Accelerate v1.12.0 и как минимум 2 GPU. Настройте параллелизм последовательностей в ParallelismConfig Accelerate и передайте его в parallelism_config или в файл конфигурации Accelerate.

parallelism_config

Файл конфигурации Accelerate

Следующие поля важны для настройки параллелизма последовательностей.

Класс Trainer автоматически управляет шардированием DataLoader, генерацией position_ids, сдвигом меток и агрегацией потерь между рангами SP. Если вы пишете собственный цикл обучения, обратитесь к руководству Accelerate по Sequence Parallelism.

sp_backend должен быть установлен в значение "deepspeed" для использования параллелизма последовательности Ulysses.
sp_size — это количество GPU, которые параллельно обрабатывают одну последовательность. Каждый ранг SP получает уникальный поток данных из DataLoader, в отличие от тензорного параллелизма, где все ранги получают идентичные данные. Эффективное значение dp_world_size = world_size / sp_size, поэтому при 4 GPU и sp_size=4 dp_world_size=1 для расчётов размера батча. Последовательности также должны быть дополнены (padding) до кратного sp_size. Соответствующим образом установите pad_to_multiple_of в вашем data collator.

Количество голок внимания должно делиться на sp_size. Модель с 32 головками поддерживает значения sp_size 1, 2, 4, 8, 16 или 32.
```
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=sp_size,
)
```
sp_seq_length_is_variable управляет обработкой переменной длины последовательностей. Установите значение True (рекомендуется), если длина последовательностей различается между батчами. Установите значение False, если все последовательности дополняются до фиксированной длины, заданной параметром sp_seq_length.
sp_attn_implementation задаёт бэкенд внимания. Поддерживаемые значения: "sdpa", "flash_attention_2" или "flash_attention_3". Рекомендуется использовать FlashAttention, особенно при упаковке нескольких образцов в один батч. SDPA может некорректно обрабатывать внимание между границами образцов при их упаковке. Eager attention не поддерживается, поскольку его 4D attention_mask отбрасывается по соображениям экономии памяти и масштабируемости.

Сочетание с параллелизмом данных

Параллелизм последовательностей и параллелизм данных используют одни и те же GPU, а SP не требует дополнительного оборудования. Чтобы запустить оба режима, задайте dp_replicate_size или dp_shard_size так, чтобы dp_replicate_size × dp_shard_size × sp_size было равно общему количеству GPU.

Например, при 8 GPU и sp_size=4 установите dp_replicate_size=2 (2 × 1 × 4 = 8).

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_replicate_size=2,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

Дальнейшие шаги

Контекстный параллелизм с FSDP2 (cp_size) отличается от Ulysses (sp_size). См. Оценка с контекстным параллелизмом.
Руководство по Sequence Parallelism в Accelerate более подробно рассматривает реализацию Ulysses и показывает, как написать собственный цикл обучения.
Учебное руководство DeepSpeed ALST охватывает всю систему ALST, включая TiledMLP и выгрузку контрольных точек активаций.
Руководство по методам параллелизма показывает, как сочетать последовательный параллелизм с другими стратегиями, такими как ZeRO.
Блог-пост Ulysses Sequence Parallelism: Training with Million-Token Contexts объясняет, как работает Ulysses и как он интегрирован в Accelerate, Trainer и SFTTrainer.

Update on GitHub