Ring Attention

Ring Attention berechnet exakte Self-Attention über eine Sequenz, die für einen einzelnen Beschleuniger zu lang ist. Dazu wird die Sequenz in Blöcke aufgeteilt, die auf Geräten angeordnet sind, die einen Ring bilden. Jedes Gerät behält einen Abfrageblock und gibt Schlüssel- und Wertblöcke an sein Nachbargerät weiter. Diese Kommunikation läuft parallel zur Berechnung auf dem bereits vorhandenen Block. Dadurch steigt die maximale Sequenzlänge, die ein Modell verarbeiten kann, proportional zur Anzahl der Geräte. Hao Liu, Matei Zaharia und Pieter Abbeel von der UC Berkeley stellten das Verfahren im Oktober 2023 in Ring Attention with Blockwise Transformers for Near-Infinite Context vor. Es baut auf ihrem früheren blockweisen, parallelen Transformer auf. Die Autoren berichteten über Training und Inferenz mit Kontextfenstern von Millionen Tokens, ohne die Attention näherungsweise zu berechnen.

So funktioniert Ring Attention#

Eine Sequenz wird in N Blöcke aufgeteilt, einen pro GPU oder TPU. Jedes Gerät hält seinen Abfrageblock während der gesamten Berechnung vor. In jedem der N Schritte berechnet das Gerät die Attention zwischen seinen Abfragen und dem Schlüssel-Wert-Block, den es gerade hält. Anschließend sendet es diesen Block an das nächste Gerät und empfängt einen weiteren vom vorherigen Gerät:

Nach N Schritten hat jeder Abfrageblock jeden Schlüssel- und Wertblock verarbeitet. Kein Gerät speichert jemals die vollständige Attention-Matrix oder alle Schlüssel und Werte. Deshalb hängt der Speicherbedarf pro Gerät von der Blockgröße und nicht von der Sequenzlänge ab. Die Übertragung wird durch Berechnungen verdeckt, solange die Berechnung der Attention für einen Block länger dauert als dessen Übertragung. Das legt eine Mindestblockgröße für die jeweilige Hardware fest. Teilergebnisse werden mit derselben fortlaufenden Softmax kombiniert, die Flash Attention innerhalb einer GPU verwendet: Jedes Gerät hält eine laufende Summe der potenzierten Werte und der gewichteten Werte vor und dividiert am Ende die eine durch die andere. Die Simulation von vier Geräten auf einem Computer zeigt, dass das Ergebnis exakt ist:

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))

Ausgegeben wird Matches full attention: True: Jedes Gerät hielt zu jedem Zeitpunkt nur einen 8 × 8 großen Bewertungsblock vor, zusammen berechneten die Geräte jedoch die Attention über die vollständige 32 × 32-Matrix. Produktive Kernel ziehen vor dem Potenzieren außerdem ein laufendes Maximum ab, damit die Zahlen stabil bleiben. Bei einer kausalen Maske überspringen Geräte mit frühen Abfrageblöcken die meisten Schlüsselblöcke und bleiben untätig. Striped Attention löst das, indem jedes Gerät Tokens erhält, die gleichmäßig über die Sequenz verteilt sind. So steigt der Trainingsdurchsatz von Anfang bis Ende bei 256K Tokens um bis zu das 1,45-Fache.

Anwendungen#

Ring Attention ist überall dort wichtig, wo eine einzelne Eingabe für ein einzelnes Gerät zu lang ist. Videos sind das deutlichste Beispiel. Dieselbe Forschungsgruppe aus Berkeley nutzte das Verfahren zum Training des Large World Model (LWM), einer Familie von Bild-Sprach-Modellen (VLMs) mit 7 Milliarden Parametern. Das Kontextfenster wurde schrittweise von 4K auf 1 Million verschachtelter Text- und Video-Tokens erweitert. LWM kodiert jedes Frame als 256 Tokens. Eine Stunde Video, mit einem Frame pro Sekunde abgetastet, umfasst daher 921.600 Tokens. Das Modell beantwortete Fragen zu einstündigen YouTube-Zusammenstellungen – eine Aufgabe zum langfristigen Videoverständnis. Auch beim Training von großen Sprachmodellen (LLMs), die ganze Codebasen oder Dokumentensammlungen in einem Prompt verarbeiten, tritt dieselbe Begrenzung auf. Die context_parallel()-API von PyTorch ersetzt innerhalb ihres Kontexts Scaled-Dot-Product-Attention durch Ring Attention. Das TorchTitan-Team nutzte diese Funktion, um Llama 3 8B auf Sequenzen mit 1 Million Tokens auf 32 H100-GPUs zu trainieren.

Ring Attention und andere Parallelisierungsverfahren#

• Kontextparallelisierung: Die allgemeine Strategie, eine Sequenz auf mehrere Geräte aufzuteilen. Ring Attention ist der Algorithmus hinter der PyTorch-Implementierung. DeepSpeed Ulysses ist eine Alternative, die All-to-All-Austausch verwendet, damit jedes Gerät alle Tokens für eine Teilmenge der Attention-Heads erhält.
• Tensorparallelisierung und Pipeline-Parallelisierung: Dabei werden die Modellgewichte innerhalb oder zwischen den Schichten aufgeteilt. Das hilft, wenn die Parameter den Speicher eines einzelnen Geräts übersteigen.
• Flash Attention: Ein Kernel für ein einzelnes Gerät. Ring Attention führt auf jedem Gerät für dessen lokalen Block einen solchen blockweisen Kernel aus.

Datenparallelisierung teilt den Batch statt der Sequenz auf und ist der Modus, den Ultralytics YOLO26 für das Training mit mehreren GPUs verwendet: Die Angabe von device=[0, 1] startet PyTorch DistributedDataParallel, das auf jeder GPU eine vollständige Kopie des Modells vorhält, jeder GPU einen gleich großen Anteil an jedem Batch von Bildern zuweist und nach jedem Schritt die Gradienten mittelt.