Ring Attention

L'attention en anneau calcule une auto-attention exacte sur une séquence trop longue pour un seul accélérateur en la divisant en blocs répartis entre des appareils organisés en anneau. Chaque appareil conserve un bloc de requêtes et transmet les blocs de clés et de valeurs à son voisin, en faisant chevaucher cette communication avec le calcul sur le bloc qu'il détient déjà ; la longueur maximale de séquence que le modèle peut gérer augmente donc proportionnellement au nombre d'appareils. Hao Liu, Matei Zaharia et Pieter Abbeel, de l'UC Berkeley, l'ont présentée en octobre 2023 dans Ring Attention with Blockwise Transformers for Near-Infinite Context, dans le prolongement de leurs travaux antérieurs sur le transformeur parallèle par blocs ; ils ont rapporté des entraînements et des inférences utilisant des fenêtres de contexte de plusieurs millions de jetons, sans approximation de l'attention.

Fonctionnement de l'attention en anneau#

Une séquence est divisée en N blocs, un par GPU ou TPU. Chaque appareil conserve son bloc de requêtes pendant toute l'opération. À chacune des N étapes, il calcule l'attention entre ses requêtes et le bloc de clés-valeurs qu'il détient, transmet ce bloc à l'appareil suivant et en reçoit un de l'appareil précédent :

Après N étapes, chaque bloc de requêtes a été mis en relation avec tous les blocs de clés et de valeurs. Aucun appareil ne stocke jamais la matrice d'attention complète ni l'ensemble complet des clés et des valeurs ; la mémoire utilisée par appareil dépend donc de la taille des blocs plutôt que de la longueur de la séquence. Le transfert est masqué par le calcul tant que le calcul de l'attention sur un bloc prend plus de temps que son envoi, ce qui définit une taille minimale de bloc pour le matériel considéré. Les résultats partiels se combinent à l'aide de la même softmax cumulative que l'attention flash utilise sur un seul GPU : chaque appareil conserve une somme cumulative des scores exponentiés et des valeurs pondérées, puis divise l'une par l'autre à la fin. La simulation de quatre appareils sur une seule machine montre que le résultat est exact :

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

Le programme affiche Matches full attention: True : chaque appareil n'a jamais détenu qu'un bloc de scores de 8 × 8, et ensemble, les appareils ont calculé l'attention sur toute la matrice 32 × 32. Les noyaux de production soustraient également un maximum cumulatif avant l'exponentiation afin de stabiliser les nombres. Avec un masque causal, les appareils qui détiennent les premiers blocs de requêtes ignorent la plupart des blocs de clés et restent inactifs ; Striped Attention corrige ce problème en répartissant uniformément les jetons de la séquence entre les appareils, ce qui augmente le débit d'entraînement de bout en bout jusqu'à 1,45 fois avec 256K jetons.

Applications#

L’attention en anneau est utile chaque fois qu’une entrée est trop longue pour être traitée sur un seul appareil, et la vidéo en est le cas le plus évident. Le même groupe de Berkeley s’en est servi pour entraîner le Grand modèle du monde (LWM), une famille de modèles vision-langage (VLMs) de 7 milliards de paramètres chacun, dont la fenêtre de contexte s’est progressivement étendue de 4K à 1 million de jetons de texte et de vidéo entrelacés. LWM encode chaque image en 256 jetons, de sorte qu’une heure de vidéo échantillonnée à raison d’une image par seconde représente 921 600 jetons ; le modèle a répondu à des questions sur des compilations YouTube d’une heure, dans le cadre d’une tâche de compréhension vidéo portant sur des séquences de longue durée. L’entraînement de grands modèles de langage (LLMs) qui lisent des bases de code entières ou des collections de documents dans une seule invite se heurte à la même limite. L’API context_parallel() de PyTorch remplace l’attention à produit scalaire avec mise à l’échelle par l’attention en anneau au sein de son contexte, et l’équipe TorchTitan a utilisé l’attention en anneau pour entraîner Llama 3 8B sur des séquences d’un million de jetons réparties sur 32 GPU H100.

L'attention en anneau et les autres formes de parallélisme#

• Parallélisme de contexte : la stratégie générale qui consiste à répartir une séquence entre plusieurs appareils. L'attention en anneau est l'algorithme qui sous-tend l'implémentation de PyTorch ; DeepSpeed Ulysses est une autre solution, qui utilise des échanges tous-vers-tous pour attribuer à chaque appareil tous les jetons d'un sous-ensemble de têtes d'attention.
• Parallélisme tensoriel et parallélisme de pipeline : les poids du modèle sont répartis au sein des couches ou entre elles, ce qui est utile lorsque les paramètres dépassent la mémoire d'un seul appareil.
• Attention flash : un noyau mono-appareil. L'attention en anneau exécute un noyau par blocs de ce type sur chaque appareil, pour son bloc local.

Le parallélisme de données répartit le lot plutôt que la séquence ; c'est le mode utilisé par Ultralytics YOLO26 pour l'entraînement multi-GPU : transmettre device=[0, 1] lance DistributedDataParallel de PyTorch, qui conserve une copie complète du modèle sur chaque GPU, attribue à chaque GPU une part égale de chaque lot d'images et calcule la moyenne des gradients après chaque étape.