Context Parallelism

Kontextparallelität ist eine Technik des verteilten Rechnens, die eine lange Eingabesequenz auf mehrere Beschleuniger aufteilt. Jede GPU verarbeitet nur einen Teil der Sequenz und kooperiert dabei während der Attention mit den anderen. Dies reduziert den Aktivierungsspeicher pro Gerät, sodass ein transformer auf Eingaben trainiert werden kann, die den Speicher einer einzelnen GPU überschreiten, wie etwa sehr lange Dokumente, erweiterte Videos oder große Sammlungen von Bild-Patches.

Im Gegensatz zur bloßen Vergrößerung des context window eines Modells ändert die Kontextparallelität nicht, wie viele Informationen die Architektur theoretisch aufnehmen kann. Stattdessen macht sie die Verarbeitung dieses Kontexts durch die Verteilung der Sequenzdimension rechentechnisch praktikabel.

Wie Kontextparallelität funktioniert#

Angenommen, eine Sequenz enthält 32.000 Token und die Kontextparallelität verwendet vier GPUs. Jedes Gerät empfängt anfangs etwa 8.000 Token und speichert die entsprechenden intermediären Aktivierungen.

Die meisten Operationen, wie Normalisierungs- und Feed-Forward-Schichten, können diese lokalen Sequenzabschnitte unabhängig voneinander verarbeiten. Die Herausforderung ist der attention mechanism: Eine lokale Query muss möglicherweise auf Keys und Values zugreifen, die von jedem anderen Gerät gehalten werden.

Ausführungen tauschen daher Key-Value- bzw. KV-Blöcke zwischen GPUs aus. Beim ring attention berechnet jedes Gerät eine partielle Attention mithilfe seiner lokalen Daten, gibt einen KV-Blöcke an das nächste Gerät weiter und wiederholt dies, bis es die vollständige Sequenz verarbeitet hat. Das PyTorch context parallel tutorial demonstriert dieses Verhalten durch verteilte skalierte Punktprodukt-Attention, während das NVIDIA context parallel package All-Gather-, Reduce-Scatter- und ringbasierte Kommunikationsoptionen beschreibt.

Das Endergebnis ist mathematisch äquivalent zur Full-Sequence-Attention (vorbehaltlich normaler numerischer Abweichungen), aber keine einzelne GPU muss jede Sequenzaktivierung vorhalten.

Warum Kontextparallelität wichtig ist#

Lange Sequenzen erzeugen zwei große Skalierungsprobleme. Erstens verbrauchen gespeicherte Aktivierungen mit wachsender Sequenzlänge mehr Speicher. Zweitens vergleicht die Standard-Self-Attention Token über die gesamte Sequenz hinweg, was beträchtliche Berechnungen und temporäre Daten erzeugt.

Kontextparallelität löst das Speicherproblem durch die Aufteilung der Aktivierungen auf Geräte. Sie kann auch die Attention-Arbeit aufteilen, obwohl die Kommunikation neue Kosten verursacht. Effiziente Systeme überlappen KV-Transfers mit Berechnungen unter Verwendung von Operationen wie jenen, die im NCCL collective operations guide dokumentiert sind.

Die Technik ist besonders wertvoll, wenn die Sequenzlänge – und nicht die Modellgewichte oder die Batch-Größe – einen Out-of-Memory-Fehler verursacht. Sie gehört zum größeren Bereich des distributed training und wird häufig mit anderen Strategien kombiniert, um mehrere Dimensionen gleichzeitig zu skalieren.

Kontextparallelität vs. verwandte Techniken#

• Tensor parallelism unterteilt Operationen oder Gewichtsmatrizen innerhalb einzelner Schichten. Die Kontextparallelität unterteilt stattdessen Token entlang der Sequenzdimension.
• Pipeline parallelism weist verschiedene Gruppen von Modellschichten unterschiedlichen Geräten zu. Sie partitioniert die Modelltiefe anstelle der Sequenzlänge.
• Data parallelism repliziert das Modell und gibt jedem Replikat verschiedene Trainingsbeispiele. Die PyTorch distributed overview empfiehlt dies, wenn das vollständige Modell und jedes Beispiel auf eine GPU passen.
• Sequence parallelism shardet oft Aktivierungen für ausgewählte Operationen im Zusammenhang mit Tensorparallelität. Die Kontextparallelität wendet die Sequenzpartitionierung breiter über Netzwerkeingaben und -aktivierungen an.

Diese Ansätze ergänzen sich. Der NVIDIA parallelism strategies guide zeigt, wie Kontext-, Tensor-, Pipeline- und Datenparallelität ein mehrdimensionales Geräte-Layout bilden können.

Praxisanwendungen#

• Lange Dokumente KI: Ein Sprachmodell für Recht oder Medizin muss möglicherweise eine gesamte Fallakte, eine Patientenakte oder ein technisches Handbuch verarbeiten. Kontextparallelität verteilt die Tausenden von Dokument-Token auf Beschleuniger, reduziert den Druck auf den Aktivierungsspeicher und bewahrt gleichzeitig die Attention zwischen entfernten Abschnitten.

• Lange Videos und multimodales Verständnis: Video-Transformer und large vision models können Frames, Bild-Patches, Audiosegmente und Text als eine lange Tokensequenz darstellen. Die Verteilung dieser Sequenz hilft Modellen, erweiterte Aufnahmen zu analysieren, ohne die Frame-Anzahl oder die räumlichen Details aggressiv zu reduzieren. Die AWS Neuron context parallelism overview veranschaulicht, wie Beschleunigergruppen KV-Shards für diese Workloads mit langem Kontext austauschen können.

Für kompakte Computer-Vision-Architekturen wie Ultralytics YOLO26 ist Kontextparallelität meist unnötig. Die Standard-Multi-GPU-Datenparallelität über den Ultralytics model training workflow ist im Allgemeinen der geeignetere Weg, um das Training zu beschleunigen.

Praktische Anwendung und Kompromisse#

Das folgende minimale Beispiel verwendet PyTorchs experimentelle kontextparallele API und scaled dot-product attention controls. Speichern Sie es als cp_example.py und starten Sie es auf zwei GPUs mit torchrun --standalone --nproc-per-node=2 cp_example.py.

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, sodass jeder Prozess ein Sequenz-Shard erhält, während sich die Attention über das Geräte-Mesh hinweg koordiniert.

In der Praxis sollten Entwickler bestätigen, dass der reduzierte Speicher den Kommunikationsaufwand überwiegt, schnelle Verbindungen nutzen und repräsentative Sequenzlängen benchmarken. Für Standard-Vision-Projekte bietet Ultralytics Platform einfachere Cloud- und lokale Workflows für Datensatzannotation, Training, Bereitstellung und Monitoring, ohne dass eine manuelle kontextparallele Konfiguration erforderlich ist.