Context Parallelism

Kontextparallelität ist ein Verfahren des verteilten Rechnens, bei dem eine lange Eingabesequenz auf mehrere Beschleuniger aufgeteilt wird. Jede GPU verarbeitet nur einen Teil der Sequenz und arbeitet bei der Aufmerksamkeit mit den anderen zusammen. Das verringert den Aktivierungsspeicherbedarf pro Gerät und ermöglicht das Training eines Transformers mit Eingaben, die den Speicher einer einzelnen GPU übersteigen können, etwa sehr langen Dokumenten, langen Videos oder umfangreichen Sammlungen von Bildausschnitten.

Anders als eine bloße Vergrößerung des Kontextfensters eines Modells ändert Kontextparallelität nicht, wie viele Informationen die Architektur theoretisch aufnehmen kann. Stattdessen macht sie die Verarbeitung dieses Kontexts praktisch umsetzbar, indem sie die Sequenzdimension verteilt.

So funktioniert Kontextparallelität#

Angenommen, eine Sequenz enthält 32.000 Token und die Kontextparallelität nutzt vier GPUs. Jedes Gerät erhält zunächst ungefähr 8.000 Token und speichert die zugehörigen Zwischenaktivierungen.

Die meisten Operationen, etwa Normalisierung und vorwärtsgerichtete Schichten, können diese lokalen Sequenzabschnitte unabhängig voneinander verarbeiten. Die Herausforderung ist der Aufmerksamkeitsmechanismus: Eine lokale Abfrage muss möglicherweise auf Schlüssel und Werte zugreifen, die auf allen anderen Geräten liegen.

Deshalb tauschen Implementierungen Schlüssel-Wert- oder KV-Blöcke zwischen GPUs aus. Bei der Ring-Aufmerksamkeit berechnet jedes Gerät die partielle Aufmerksamkeit anhand seiner lokalen Daten, übergibt einen KV-Block an das nächste Gerät und wiederholt den Vorgang, bis die gesamte Sequenz verarbeitet wurde. Das PyTorch-Tutorial zur Kontextparallelität veranschaulicht dieses Verhalten mit verteilter skalierter Skalarproduktaufmerksamkeit. Das NVIDIA-Paket für Kontextparallelität beschreibt außerdem Optionen für All-Gather, Reduce-Scatter und ringbasierte Kommunikation.

Das Endergebnis ist, abgesehen von den üblichen numerischen Abweichungen, mathematisch gleichwertig mit der Aufmerksamkeit über die vollständige Sequenz. Keine einzelne GPU muss jedoch sämtliche Sequenzaktivierungen speichern.

Warum Kontextparallelität wichtig ist#

Lange Sequenzen verursachen zwei große Skalierungsprobleme. Erstens belegen gespeicherte Aktivierungen mit zunehmender Sequenzlänge mehr Speicher. Zweitens vergleicht standardmäßige Selbstaufmerksamkeit Token innerhalb der Sequenz und verursacht dadurch erheblichen Rechenaufwand sowie umfangreiche temporäre Daten.

Kontextparallelität löst das Speicherproblem, indem sie die Aktivierungen auf mehrere Geräte aufteilt. Sie kann auch die Aufmerksamkeitsberechnungen aufteilen, wobei durch die Kommunikation zusätzliche Kosten entstehen. Effiziente Systeme überlappen KV-Übertragungen mit Berechnungen und verwenden dafür Operationen wie jene, die im Leitfaden zu kollektiven NCCL-Operationen dokumentiert sind.

Das Verfahren ist besonders nützlich, wenn die Sequenzlänge und nicht Modellgewichte oder Stapelgröße einen Fehler wegen unzureichenden Speichers verursacht. Es gehört zum umfassenderen Bereich des verteilten Trainings und wird häufig mit anderen Verfahren kombiniert, um mehrere Dimensionen gleichzeitig zu skalieren.

Kontextparallelität im Vergleich zu verwandten Verfahren#

• Tensorparallelität teilt Operationen oder Gewichtsmatrizen innerhalb einzelner Schichten auf. Kontextparallelität teilt dagegen Token entlang der Sequenzdimension auf.
• Pipeline-Parallelität weist verschiedene Gruppen von Modellschichten unterschiedlichen Geräten zu. Sie unterteilt die Tiefe des Modells und nicht die Sequenzlänge.
• Datenparallelität repliziert das Modell und weist jeder Replik eine andere Auswahl von Trainingsbeispielen zu. Die Übersicht zum verteilten Training in PyTorch empfiehlt dieses Verfahren, wenn das vollständige Modell und jedes Beispiel auf eine GPU passen.
• Sequenzparallelität verteilt häufig Aktivierungen für ausgewählte Operationen, die mit Tensorparallelität zusammenhängen. Kontextparallelität wendet die Aufteilung von Sequenzen umfassender auf Netzwerkeingaben und Aktivierungen an.

Diese Ansätze ergänzen einander. Der Leitfaden von NVIDIA zu Parallelisierungsstrategien zeigt, wie Kontext-, Tensor-, Pipeline- und Datenparallelität ein mehrdimensionales Geräte-Layout bilden können.

Anwendungen in der Praxis#

• KI für lange Dokumente: Ein Sprachmodell für juristische oder medizinische Anwendungen muss möglicherweise eine vollständige Fallakte, Krankengeschichte oder technische Anleitung verarbeiten. Kontextparallelität verteilt Tausende Dokument-Token auf Beschleuniger und verringert so den Speicherbedarf für Aktivierungen, während die Aufmerksamkeit zwischen weit auseinanderliegenden Abschnitten erhalten bleibt.

• Verständnis langer Videos und multimodaler Daten: Videotransformer und große Bildverarbeitungsmodelle können Einzelbilder, Bildausschnitte, Audiosegmente und Text als eine lange Token-Sequenz darstellen. Durch deren Verteilung können Modelle längere Aufnahmen analysieren, ohne die Zahl der Einzelbilder oder die räumlichen Details stark reduzieren zu müssen. Die Übersicht von AWS Neuron zur Kontextparallelität veranschaulicht, wie Beschleunigergruppen KV-Teilstücke für diese Aufgaben mit langen Kontexten austauschen können.

Für kompakte Bildverarbeitungsarchitekturen wie Ultralytics YOLO26 ist Kontextparallelität normalerweise nicht erforderlich. Standardmäßige Datenparallelität mit mehreren GPUs über den Ultralytics-Ablauf zum Modelltraining ist im Allgemeinen die passendere Methode, um das Training zu beschleunigen.

Praktische Anwendung und Kompromisse#

Das folgende minimale Beispiel verwendet die experimentelle Kontextparallelitäts-API von PyTorch sowie Steuerungsmöglichkeiten für skalierte Skalarproduktaufmerksamkeit. Speichere es unter cp_example.py und starte es mit torchrun --standalone --nproc-per-node=2 cp_example.py auf zwei GPUs.

import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(rank)
torch.cuda.manual_seed(0)
dist.init_process_group("nccl")

mesh = init_device_mesh("cuda", (world_size,))
qkv = [torch.randn(1, 4, 4096, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

with sdpa_kernel(SDPBackend.FLASH_ATTENTION), context_parallel(mesh, buffers=tuple(qkv), buffer_seq_dims=(2, 2, 2)):
    output = F.scaled_dot_product_attention(*qkv, is_causal=True)

output.float().square().mean().backward()
dist.destroy_process_group()

Hier ist Dimension 2 die Sequenzdimension. Daher erhält jeder Prozess einen Sequenzabschnitt, während die Aufmerksamkeit über das Gerätegitter hinweg koordiniert wird.

In der Praxis sollten Entwickler bestätigen, dass die Speicherersparnis den Kommunikationsaufwand überwiegt, schnelle Verbindungen verwenden und repräsentative Sequenzlängen messen. Für gewöhnliche Bildverarbeitungsprojekte bietet die Ultralytics Platform einfachere Cloud- und lokale Abläufe für Datensatzannotation, Training, Bereitstellung und Überwachung, ohne dass eine manuelle Konfiguration der Kontextparallelität erforderlich ist.