Paralelismo de sequência Ulysses

O paralelismo de sequência Ulysses (SP) treina em sequências muito longas dividindo-as entre várias GPUs. Para calcular a atenção corretamente, um coletivo all-to-all troca a dimensão de sharding de sequência para cabeças de atenção. Cada GPU então possui a sequência completa e calcula a atenção localmente sobre um subconjunto de cabeças. Um segundo all-to-all retorna ao layout com sharding por sequência, para que o restante da passagem de frente continue localmente em cada 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  │
                   └───────────────┘          └───────────────┘

Este guia aborda o componente de paralelismo de sequência Ulysses do ALST (Arctic Long Sequence Training). O sistema ALST completo também inclui TiledMLP e offloading de checkpoint de ativação, que não estão disponíveis no Transformers. Consulte o tutorial do DeepSpeed ALST para o sistema completo.

Configurar

O paralelismo de sequência requer o Accelerate v1.12.0 e pelo menos 2 GPUs. Configure o paralelismo de sequência no ParallelismConfig do Accelerate e passe-o para parallelism_config ou para um arquivo de configuração do Accelerate.

parallelism_config

Arquivo de configuração do Accelerate

Os campos a seguir são importantes para configurar o paralelismo de sequência.

O Trainer trata automaticamente do sharding do DataLoader, da geração de position_ids, do deslocamento de rótulos e da agregação da perda entre os ranks de SP. Se você estiver escrevendo um loop de treinamento personalizado, consulte o guia de Sequence Parallelism do Accelerate.

sp_backend deve ser definido como "deepspeed" para usar o paralelismo de sequência Ulysses.
sp_size é o número de GPUs que processam uma única sequência em paralelo. Cada rank de SP recebe um fluxo de dados exclusivo do DataLoader, diferente da paralelização de tensor, em que todos os ranks recebem dados idênticos. O dp_world_size = world_size / sp_size efetivo, portanto, com 4 GPUs e sp_size=4, dp_world_size=1 para os cálculos de batch size. As sequências também devem ser preenchidas até um múltiplo de sp_size. Defina pad_to_multiple_of no seu data collator de acordo.

O número de cabeças de atenção deve ser divisível por sp_size. Um modelo com 32 cabeças suporta sp_size de 1, 2, 4, 8, 16 ou 32.
```
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=sp_size,
)
```
sp_seq_length_is_variable controla o tratamento de comprimentos de sequência variáveis. Defina-o como True (recomendado) para comprimentos variáveis entre lotes. Defina-o como False quando todas as sequências forem preenchidas até um comprimento fixo especificado por sp_seq_length.
sp_attn_implementation define o backend de atenção. Os valores suportados são "sdpa", "flash_attention_2" ou "flash_attention_3". O FlashAttention é recomendado, especialmente ao empacotar múltiplas amostras em um lote. O SDPA pode atender incorretamente entre as fronteiras das amostras quando elas estão empacotadas. A atenção eager não é suportada porque sua attention_mask 4D é descartada por motivos de memória e escalabilidade.

Combinando com paralelismo de dados

O paralelismo de sequência e o paralelismo de dados usam os mesmos GPUs, e o SP não requer hardware adicional. Para executar ambos, defina dp_replicate_size ou dp_shard_size de modo que dp_replicate_size × dp_shard_size × sp_size seja igual ao número total de GPUs.

Por exemplo, com 8 GPUs e sp_size=4, defina 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",
    ),
)

Próximas etapas

O paralelismo de contexto com FSDP2 (cp_size) é separado do Ulysses (sp_size). Consulte Avaliar com paralelismo de contexto.
O guia de Sequence Parallelism do Accelerate aborda a implementação Ulysses com mais profundidade e mostra como escrever um loop de treinamento personalizado.
O tutorial do DeepSpeed ALST aborda o sistema ALST completo, incluindo TiledMLP e offloading de checkpoint de ativação.
O guia de métodos de paralelismo mostra como combinar o paralelismo de sequência com outras estratégias, como o ZeRO.
O post no blog Ulysses Sequence Parallelism: Training with Million-Token Contexts explica como o Ulysses funciona e como ele é integrado no Accelerate, Trainer e SFTTrainer.

Update on GitHub