Parallélisme de séquence dans 🤗 accelerate

Ce guide abordera les bases de l'utilisation du parallélisme de séquences dans 🤗accelerate.

Voir également le très lié Context Parallelism.

Pourquoi le parallélisme de séquence ?

Avec l’avènement des grands modèles de langage, et récemment des modèles de raisonnement, la longueur des séquences a augmenté rapidement. Cela, combiné à la complexité mémoire quadratique de l’attention, a engendré le besoin de méthodes plus efficaces pour entraîner des modèles avec de longues séquences. Avec une longueur de séquence de 128k, l’exigence mémoire de la matrice d’attention est de 128k * 128k * 2 bytes * num_heads = ~32 GB * num_heads pour une précision bf16, en supposant une implémentation d’attention standard. Certes, avec l’utilisation de flash attention ou de SDPA, qui ne matérialisent pas ces poids d’attention, cela diminue considérablement, mais la croissance des exigences mémoire reste importante.

Le parallélisme de séquence Ulysses permet de fractionner les entrées du calcul d'attention le long de la dimension de la séquence et de calculer l'attention normalement, mais en utilisant uniquement une tranche de têtes d'attention sur chaque GPU. Grâce à cela, nous pouvons entraîner des modèles avec de longues séquences, avec quelques outils supplémentaires, en passant à l'échelle de plus de 15 millions de longueur de séquence. Pour voir comment compléter Ulysses SP avec TiledMLP, Liger-Kernel, le déchargement du point de contrôle d'activation vers le CPU et quelques autres astuces, veuillez vous référer à l'article : Arctic Long Sequence Training: Scalable And Efficient Training For Multi-Million Token Sequences.

En quoi Ulysses SP diffère-t-il de FSDP CP

Dans le document Context Parallelism, vous pouvez découvrir le déploiement d’une autre technologie appelée Context Parallelism, qui elle aussi découpe sur la dimension de la séquence, mais utilise Ring Attention au lieu de découper sur la dimension des têtes.

Les articles suivants proposent une explication très détaillée des différences entre les deux technologies :

• https://insujang.github.io/2024-01-11/tensor-parallelism-and-sequence-parallelism-detailed-analysis/
• https://huggingface.co/blog/exploding-gradients/ulysses-ring-attention

Un résumé rapide adapté d'un des articles :

• Ulysses SP présente une surcharge de communication relativement faible, mais il est limité par le nombre de têtes d'attention et impose donc certaines exigences en matière de topologie réseau (le nombre de têtes d'attention doit être divisible par le nombre de GPU participant à une seule réplique). La communication all-to-all est sensible à la latence et nécessite DeepSpeed.
• La communication en anneau P2P de l’attention en anneau CP de FSDP n’est pas soumise aux exigences de divisibilité mentionnées précédemment, mais elle implique un volume de communication plus élevé.

Enfin, il devrait être possible de combiner SP + CP, comme expliqué dans l'article USP: A Unified Sequence Parallelism Approach for Long Context Generative AI, afin de prendre en charge des longueurs de séquence encore plus importantes, bien que cela ne soit pas encore intégré dans 🤗accelerate.

Backends de parallélisme de séquences pris en charge

Actuellement, le seul backend de parallélisme de séquences est deepspeed, issu de la version modernisée d’Ulysses SP, qui fait partie de la technologie Arctic Long Sequence Training. Un tutoriel est également disponible si vous souhaitez l’intégrer directement dans votre propre code.

Comment utiliser le parallélisme de séquence ?

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

+# Example: 4 GPUs with sp_size=4, dp_shard_size=1
+# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
+     sp_backend="deepspeed",
+     sp_size=4,
+     dp_shard_size=1,  # Explicit: no data parallelism
+     sp_handler=DeepSpeedSequenceParallelConfig(
+         sp_seq_length_is_variable: true,
+         sp_attn_implementation="sdpa",
+     ),
+ )

accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

Comme pour toute autre fonctionnalité de 🤗accelerate, vous pouvez également activer le parallélisme de séquence en passant les drapeaux correspondants à accelerate launch. Dans ce cas, rien ne change :

accelerate launch --parallelism-config-sp-size 8  ...

Vous pouvez également définir sp_size et d’autres paramètres de configuration dans la commande accelerate config, qui les enregistre dans votre fichier de configuration accelerate, afin de ne plus avoir à les passer à chaque lancement de votre script.

Le parallélisme de séquences se combine avec le parallélisme de données. Il n’exige pas de GPU supplémentaires. Ainsi, si vous disposez de 8 GPU, vous pouvez exécuter : --parallelism-config-dp-shard-size 8 --parallelism-config-sp-size 8. Vous pouvez également utiliser la classe ParallelismConfig pour les définir de manière programmatique.

Important : vous devez vous assurer que dp_replicate_size × dp_shard_size × sp_size = num_processes. Par exemple, avec 8 GPU et sp_size=8, il faut dp_shard_size=1 (puisque 1 × 1 × 8 = 8). Avec 4 GPU et sp_size=2, vous pouvez utiliser dp_shard_size=2 (puisque 1 × 2 × 2 = 4) pour un parallélisme 2D.

Configuration du backend ALST/Ulysses SP

ALST/UlyssesSP implémente le parallélisme de séquence en utilisant le parallélisme des têtes d'attention, comme expliqué dans ce papier. Pour simplifier, nous réutilisons le concept et la configuration du parallélisme de séquence, qui, du point de vue de l'utilisateur, est identique : plusieurs GPU sont utilisés pour traiter un seul lot.

Pour donner une idée de ce que ALST a rendu possible – il nous a permis d’entraîner en bf16 avec 500K tokens sur un seul GPU H100, 3,7M sur un seul nœud, et 15M sur Llama-8B en utilisant seulement quatre nœuds. Cette fonctionnalité de HF Accelerate n’active que 1 des 3 composants d’ALST, donc la longueur de séquence atteignable sera plus petite. Vous voudriez activer TiledMLP, Activation checkpoint offload to CPU, et quelques autres choses pour obtenir toute la puissance d’ALST. Pour plus de détails, veuillez consulter ce tutoriel.

Pour configurer le backend deepspeed :

# Example: 4 GPUs with sp_size=4, dp_shard_size=1
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 1 × 1 × 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_shard_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length=256,
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)
accelerator = Accelerator(
    ...,
    parallelism_config=parallelism_config,
)

• sp_backend : défini sur deepspeed ici
• sp_size est le degré du parallélisme de séquence — dans l’exemple ci-dessus, il vaut 4, donc 4 GPU seront utilisés pour traiter un seul lot (tout en effectuant DP=4 sur les mêmes GPU)
• sp_seq_length et sp_seq_length_is_variable sont utilisés pour gérer les longueurs de séquence. Si sp_seq_length_is_variable=True, le backend fonctionnera avec une longueur de séquence susceptible de varier d’un lot à l’autre, auquel cas la valeur de sp_seq_length peut être définie à n’importe quelle valeur divisible par le degré de parallélisme de séquence, ou laissée non définie. Dans ce cas, à chaque forward, les variables de séquence seront dérivées de l’entrée. Si False, alors seq_length doit correspondre à la dimension de la longueur de séquence du lot, qui devra alors être complétée (padding) pour rester toujours identique. La valeur par défaut est True.
• sp_attn_implementation est l’un de sdpa, flash_attention_2 ou flash_attention_3. Cette implémentation parallèle de séquences utilise position_ids au lieu de attention_mask ; par conséquent, eager ne peut pas fonctionner ici tant qu’il ne prend pas en charge le travail avec position_ids. Veuillez également noter que sdpa ne gère pas correctement plusieurs échantillons combinés en un seul ; il prendra en compte l’ensemble de l’échantillon comme un seul. Si les échantillons ne sont pas combinés, sdpa fonctionnera correctement. Par conséquent, Flash Attention devrait être le choix idéal, car il fonctionne toujours.

Au lieu de définir ces valeurs dans l’objet DeepSpeedSequenceParallelConfig, vous pouvez également utiliser les variables d’environnement pour obtenir le même résultat – voici leur correspondance avec la fin de la liste ci-dessus.

• PARALLELISM_CONFIG_SP_BACKEND
• PARALLELISM_CONFIG_SP_SEQ_LENGTH
• PARALLELISM_CONFIG_SP_SEQ_LENGTH_IS_VARIABLE
• PARALLELISM_CONFIG_SP_ATTN_IMPLEMENTATION

Si sp_size n'est pas passé dans le code, il peut être défini via l'argument CLI --parallelism_config_sp_size. Idem pour les autres arguments. Vous pouvez également utiliser la configuration de type fichier accelerate config, par exemple pour 2 GPU :

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 2
parallelism_config:
  parallelism_config_dp_replicate_size: 1
  parallelism_config_dp_shard_size: 1  # Must satisfy: 1 × 1 × 2 = 2 num_processes
  parallelism_config_sp_size: 2
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

Comme mentionné précédemment, le parallélisme de séquence Ulysses est normalement superposé au parallélisme de données - les mêmes rangs sont utilisés pour alimenter des flux de données uniques et effectuent également le parallélisme de séquence Ulysses. Mais vous pouvez également créer des répliques comme suit :

# Example: 4 GPUs with 2D parallelism (SP=2, DP=2)
# Ensure: dp_replicate_size × dp_shard_size × sp_size = 2 × 1 × 2 = 4 GPUs
parallelism_config = ParallelismConfig(
    dp_replicate_size=2,
    dp_shard_size=1,  # Explicit: no sharding within replicas
    sp_size=2,
    sp_backend="deepspeed",
    sp_handler=DeepSpeedSequenceParallelConfig(...),
)

Ici, nous utilisons 4 gpus, avec 2 répliques de parallélisme de séquence. Deepspeed-ZeRO est ce qui pilote le parallélisme de données ici.

Veuillez noter qu’une grande partie de la logique est cachée à l’intérieur de UlyssesSPDataLoaderAdapter. Il est utilisé en arrière-plan, en encapsulant votre objet DataLoader d’origine, mais vous devez en avoir conscience au cas où vous rencontreriez des problèmes. Il injecte également automatiquement le shift_labels correct dans le dictionnaire de lot, avant que le lot ne soit réparti entre les rangs participants.

Il ne reste plus qu'une étape pour commencer à utiliser ALST/UlyssesSP : agréger la perte entre les rangs à l'aide d'un all_gather différentiable afin d'obtenir les bons gradients. Le code ci-dessous le fait, tout en excluant les tokens masqués avec -100, pour obtenir la moyenne correcte :

sp_size = parallelism_config.sp_size if parallelism_config is not None else 1
if sp_size > 1:
    sp_group = accelerator.torch_device_mesh["sp"].get_group()
    sp_world_size = parallelism_config.sp_size

# Normal training loop
for iter, batch in enumerate(dl):
    optimizer.zero_grad()

    batch = move_to_device(batch, model.device)

    # The model automatically receives shift_labels via **kwargs and uses it for loss computation.
    # Both standard transformers models and Liger-patched models handle this correctly.
    outputs = model(**batch)
    loss = outputs.loss
    shift_labels = batch["shift_labels"]

    if sp_size > 1:
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (shift_labels != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(
            good_tokens, group=sp_group
        )
        # Skip ranks with zero valid tokens to avoid NaN contamination (NaN * 0 = NaN)
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

    if rank == 0: accelerator.print(f"{iter}: {loss=}")
    accelerator.log(dict(train_loss=loss, step=iter))

    accelerator.backward(loss)
    optimizer.step()

Notez que les modèles gèrent automatiquement shift_labels lorsqu’il est présent dans le lot. Le passage avant du modèle reçoit shift_labels via **kwargs et le transmet à la fonction de perte, qui calcule correctement la perte pour le parallélisme de séquence. Si vous utilisez Liger Kernel, il gère également shift_labels de manière transparente et calcule la perte de façon très économe en mémoire. Liger est fortement recommandé pour les longueurs de séquence importantes, car il libère la mémoire GPU en utilisant des opérations fusionnées (par exemple, le calcul fusionné logit-perte qui ne matérialise jamais le tenseur complet des logits en mémoire).

Si vous souhaitez voir ce que HF Accelerate a fait en coulisses, lisez ce tutoriel d’intégration complet.

Pour un exemple de boucle d’entraînement Accelerate avec ALST/UlyssesSP activés, voir examples/alst_ulysses_sequence_parallelism.

[!Avertissement]

Cette API est assez nouvelle et se trouve encore à un stade expérimental. Bien que nous nous efforcions de fournir une API stable, certaines petites parties de l'API publique peuvent évoluer à l'avenir.

Étant donné qu'il s'agit d'un backend Deepspeed, la configuration Deepspeed habituelle s'applique, vous permettant ainsi de combiner le parallélisme de séquence avec le déchargement des états de l'optimiseur et/ou des poids afin de libérer davantage de mémoire GPU et d'activer une longueur de séquence encore plus importante. Cette technologie a été testée et fonctionne avec DeepSpeed ZeRO stage 2 et 3.