Parallélisme de séquence Ulysses

Le parallélisme de séquence Ulysses (SP) entraîne sur des séquences très longues en les découpant sur plusieurs GPU. Pour calculer correctement l'attention, une opération collective all-to-all échange la dimension de partitionnement de la séquence vers les têtes d'attention. Chaque GPU dispose alors de la séquence complète et calcule l'attention localement sur un sous-ensemble de têtes. Un second all-to-all revient à la disposition partitionnée par séquence afin que le reste du passage avant continue localement sur chaque fragment.

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

Ce guide aborde le composant de parallélisme de séquence Ulysses de ALST (Arctic Long Sequence Training). Le système ALST complet inclut également TiledMLP et le déchargement du point de contrôle d'activation, qui ne sont pas disponibles dans Transformers. Consultez le tutoriel DeepSpeed ALST pour le système complet.

Configurer

Le parallélisme de séquences nécessite Accelerate v1.12.0 et au moins 2 GPU. Configurez le parallélisme de séquences dans ParallelismConfig d’Accelerate et transmettez-le à parallelism_config ou à un fichier de configuration Accelerate.

parallelism_config

Fichier de configuration d'Accelerate

Les champs suivants sont importants pour configurer le parallélisme de séquence.

Le Trainer gère automatiquement le fractionnement du DataLoader, la génération des position_ids, le décalage des étiquettes et l’agrégation de la perte entre les rangs SP. Si vous écrivez une boucle d’entraînement personnalisée, consultez à la place le guide Sequence Parallelism d’Accelerate.

sp_backend doit être défini sur "deepspeed" pour utiliser le parallélisme de séquence Ulysses.
sp_size est le nombre de GPU qui traitent une seule séquence en parallèle. Chaque rang SP reçoit un flux de données unique du DataLoader, contrairement au parallélisme de tenseurs où tous les rangs reçoivent des données identiques. La valeur effective de dp_world_size = world_size / sp_size, donc avec 4 GPU et sp_size=4, dp_world_size=1 pour les calculs de taille de lot. Les séquences doivent également être remplies jusqu'à un multiple de sp_size. Définissez pad_to_multiple_of en conséquence dans votre collateur de données.

Le nombre de têtes d'attention doit être divisible par sp_size. Un modèle avec 32 têtes prend en charge des valeurs de sp_size égales à 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 contrôle la gestion des longueurs de séquences variables. Définissez-le sur True (recommandé) pour des longueurs variables entre les lots. Définissez-le sur False lorsque toutes les séquences sont remplies jusqu'à une longueur fixe spécifiée par sp_seq_length.
sp_attn_implementation définit le backend d’attention. Les valeurs prises en charge sont "sdpa", "flash_attention_2" ou "flash_attention_3". FlashAttention est recommandé, en particulier lors de l’empilement de plusieurs échantillons dans un lot. SDPA peut traiter incorrectement l’attention au-delà des limites des échantillons lorsque ceux-ci sont empilés. L’attention Eager n’est pas prise en charge car son attention_mask 4D est ignoré pour des raisons de mémoire et de mise à l’échelle.

Combinaison avec le parallélisme de données

Le parallélisme de séquence et le parallélisme de données utilisent les mêmes GPU, et le SP n’exige pas de matériel supplémentaire. Pour exécuter les deux, définissez dp_replicate_size ou dp_shard_size de sorte que dp_replicate_size × dp_shard_size × sp_size soit égal au nombre total de GPU.

Par exemple, avec 8 GPU et sp_size=4, définissez 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",
    ),
)

Étapes suivantes

Le parallélisme de contexte avec FSDP2 (cp_size) est distinct d’Ulysses (sp_size). Consultez Évaluer avec le parallélisme de contexte.
Le guide Sequence Parallelism d’Accelerate aborde l’implémentation Ulysses en détail et montre comment écrire une boucle d’entraînement personnalisée.
Le tutoriel DeepSpeed ALST couvre le système ALST complet, y compris TiledMLP et le déchargement du point de contrôle d'activation.
Le guide sur les méthodes de parallélisme montre comment combiner le parallélisme de séquence avec d'autres stratégies comme ZeRO.
L’article de blog Ulysses Sequence Parallelism: Training with Million-Token Contexts explique le fonctionnement d’Ulysses et son intégration dans Accelerate, Trainer et SFTTrainer.

Update on GitHub