Paralelismo de secuencias de Ulysses

El paralelismo de secuencia de Ulysses (SP) entrena con secuencias muy largas dividiéndolas entre varias GPU. Para calcular la atención correctamente, una operación colectiva all-to-all intercambia la dimensión de particionado de secuencia a cabezas de atención. Cada GPU tiene entonces la secuencia completa y calcula la atención localmente sobre un subconjunto de cabezas. Un segundo all-to-all devuelve el diseño particionado por secuencia para que el resto de la pasada hacia adelante continúe localmente en cada fragmento.

                        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  │
                   └───────────────┘          └───────────────┘

Esta guía cubre el componente de paralelismo de secuencia Ulysses de ALST (Arctic Long Sequence Training). El sistema completo de ALST también incluye TiledMLP y la descarga de puntos de control de activación, que no están disponibles en Transformers. Consulte el tutorial de DeepSpeed ALST para el sistema completo.

Configurar

El paralelismo de secuencias requiere Accelerate v1.12.0 y al menos 2 GPUs. Configure el paralelismo de secuencias en ParallelismConfig de Accelerate y páselo a parallelism_config o a un archivo de configuración de Accelerate.

parallelism_config

Archivo de configuración de Accelerate

Los siguientes campos son importantes para configurar el paralelismo de secuencias.

El Trainer gestiona automáticamente el sharding del DataLoader, la generación de position_ids, el desplazamiento de etiquetas y la agregación de la pérdida entre los rangos de SP. Si estás escribiendo un bucle de entrenamiento personalizado, consulta en su lugar la guía de Sequence Parallelism de Accelerate.

sp_backend debe establecerse en "deepspeed" para usar el paralelismo de secuencia de Ulysses.
sp_size es el número de GPUs que procesan una secuencia única en paralelo. Cada rango de SP recibe un flujo de datos único del DataLoader, a diferencia de la paralelización de tensores, donde todos los rangos reciben datos idénticos. El dp_world_size = world_size / sp_size efectivo, por lo que con 4 GPUs y sp_size=4, dp_world_size=1 para los cálculos del tamaño de lote. Las secuencias también deben rellenarse hasta un múltiplo de sp_size. Establezca pad_to_multiple_of en su collator de datos en consecuencia.

El número de cabezas de atención debe ser divisible por sp_size. Un modelo con 32 cabezas admite sp_size de 1, 2, 4, 8, 16 o 32.
```
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=sp_size,
)
```
sp_seq_length_is_variable controla el manejo de longitudes de secuencia variables. Establézcalo en True (recomendado) para longitudes variables entre lotes. Establézcalo en False cuando todas las secuencias se rellenan hasta una longitud fija especificada por sp_seq_length.
sp_attn_implementation establece el backend de atención. Los valores admitidos son "sdpa", "flash_attention_2" o "flash_attention_3". Se recomienda FlashAttention, especialmente al empaquetar varias muestras en un lote. SDPA puede atender incorrectamente entre los límites de las muestras cuando estas están empaquetadas. La atención eager no está admitida porque su attention_mask de 4D se descarta por razones de memoria y escalabilidad.

Combinación con paralelismo de datos

El paralelismo de secuencias y el paralelismo de datos utilizan las mismas GPU, y el SP no requiere hardware adicional. Para ejecutar ambos, configure dp_replicate_size o dp_shard_size de modo que dp_replicate_size × dp_shard_size × sp_size sea igual al número total de GPU.

Por ejemplo, con 8 GPUs y sp_size=4, establezca 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óximos pasos

La paralelización de contexto con FSDP2 (cp_size) es independiente de Ulysses (sp_size). Consulte Evaluar con paralelización de contexto.
La guía de Sequence Parallelism de Accelerate aborda la implementación de Ulysses con mayor profundidad y muestra cómo escribir un bucle de entrenamiento personalizado.
El tutorial de DeepSpeed ALST cubre el sistema completo de ALST, incluido TiledMLP y la descarga de puntos de control de activación.
La guía de métodos de paralelismo muestra cómo combinar el paralelismo de secuencia con otras estrategias como ZeRO.
La publicación de blog Ulysses Sequence Parallelism: Training with Million-Token Contexts explica cómo funciona Ulysses y cómo se integra en Accelerate, Trainer y SFTTrainer.

Update on GitHub