Tensor Parallelism: Silizium zerschneiden

Das effektive Training von Large Language Models (LLMs) erfordert mehr als einfache Datenparallelität. Wenn die Gewichtsmatrizen eines Modells zu groß sind, um in den VRAM einer einzelnen GPU zu passen – oder wenn der Speicherbedarf für Zwischenaktivierungen explodiert – müssen wir das Modell selbst aufteilen.

Tensor Parallelism (TP) ist eine Technik der Modellparallelität innerhalb einer Schicht. Im Gegensatz zur Pipeline Parallelism (bei der das Modell vertikal nach Schichten aufgeteilt wird) oder Data Parallelism (bei der das Modell repliziert wird), teilt TP das Modell horizontal, indem einzelne Tensoren (Gewichtsmatrizen) auf mehrere Geräte verteilt werden. Dadurch kann ein Cluster von GPUs wie ein einziger, massiver Beschleuniger agieren und eine einzelne Operation gleichzeitig verarbeiten.

Dieser Artikel erklärt die Mechanik von TP mit einer PyTorch-Illustration nur für den Vorwärtspfad und diskutiert die Hardware-Abwägungen.

1. Die Hardware-Mechanik: Die Matrix aufteilen

Auf Registerebene wird das Training neuronaler Netzwerke von Matrixmultiplikation (MatMul) dominiert: $Y = X \cdot W$ . Bei Tensor Parallelism nutzen wir die Eigenschaften der linearen Algebra, um genau diese Operation auf $N$ GPUs zu verteilen.

Es gibt zwei Hauptmethoden, diese Berechnung aufzuteilen, wobei jede spezifische Kommunikationsprimitive erfordert, um die Ergebnisse zu synchronisieren:

Strategie A: Spalten-lineare Parallelität

In diesem Schema teilen wir die Gewichtsmatrix $W$ entlang ihrer Spalten auf.

Partitionierung: Wenn wir 2 GPUs haben, teilen wir $W$ in $[W_1, W_2]$ auf.
Ausführung: Wir replizieren die Eingabe $X$ auf beiden GPUs. GPU 1 berechnet $Y_1 = X \cdot W_1$ und GPU 2 berechnet $Y_2 = X \cdot W_2$ .
Ergebnis: Jede GPU hält eine partielle Breite des Ausgaberesultats (z. B. die erste Hälfte der Ausgabe-Features).
Kommunikation: Wenn die nächste Operation die vollständige Ausgabe $Y$ benötigt, kann All-Gather die Teilergebnisse verketten. Eine nachfolgende row-parallele Schicht kann die Shards direkt verarbeiten.

Strategie B: Zeilen-lineare Parallelisierung

Hier teilen wir die Gewichtsmatrix $W$ entlang ihrer Zeilen auf.

Partitionierung: $W$ wird in zwei Zeilenblöcke $W_1$ und $W_2$ (vertikal gestapelt) aufgeteilt.
Ausführung: Damit die Mathematik funktioniert, muss die Eingabe $X$ ebenfalls entlang ihrer letzten Dimension (Spalten) in $[X_1, X_2]$ aufgeteilt werden. GPU 1 berechnet $Y_1 = X_1 \cdot W_1$ .
Ergebnis: Jede GPU hält eine Teilsumme des Endergebnisses. $Y = Y_1 + Y_2$ .
Kommunikation: Um die gültige endgültige Ausgabe zu erhalten, müssen die Ergebnisse aller GPUs summiert werden. Dies erfordert eine All-Reduce-Operation.

Die "Megatron-LM"-Optimierung

Die Effizienz im TP ergibt sich aus der Kombination dieser beiden Strategien, um die Kommunikation zu minimieren. In einem Standard-Transformer-MLP-Block (Linear $\rightarrow$ GeLU $\rightarrow$ Linear) können wir die Aufteilungen so anordnen, dass eine Synchronisation in der Mitte vermieden wird.

Layer 1 (Column Parallel): Die Gewichtsmatrix wird spaltenweise aufgeteilt. Die Ausgabe sind auf jeder GPU aufgeteilte Aktivierungen.
Nichtlinearität (GeLU): Da GeLU eine elementweise Operation ist ( $f(x)$ ), kann sie unabhängig auf die Teilergebnisse jeder GPU angewendet werden. Hier ist keine Kommunikation erforderlich.
Layer 2 (Row Parallel): Die zweite Gewichtsmatrix wird zeilenweise aufgeteilt. Sie nimmt die aufgeteilte Ausgabe von Layer 1 direkt als aufgeteilten Input entgegen.
Finale Synchronisation: Erst nach der zweiten Schicht führen wir ein All-Reduce durch, um die Teilergebnisse zu summieren.

Bei diesem MLP aus zwei linearen Schichten entfällt das All-Gather dazwischen; im Vorwärtspfad bleibt die Reduktion nach der row-parallelen Schicht. Attention und Rückwärtspfad benötigen weitere Kollektivoperationen.

2. Bare Metal Implementierung: Reines PyTorch

Das folgende Beispiel nur für den Vorwärtspfad zeigt eine row-parallele lineare Schicht mit torch.distributed. Auf allen Ranks müssen eine Prozessgruppe und dieselbe replizierte Eingabe vorliegen. Für korrektes Training wären autogradfähige Kollektivoperationen nötig.

python

import torch
import torch.nn as nn
import torch.distributed as dist
 
class RowParallelLinear(nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        # 1. Setup World Info
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        
        # 2. Calculate Shard Size
        # We split the input dimension (rows of the weight matrix) across GPUs
        # Note: We assume input_size is divisible by world_size for simplicity
        assert input_size % self.world_size == 0
        self.input_shard_size = input_size // self.world_size
        
        # 3. Initialize Sharded Weights
        # Each GPU only holds a fraction of the total weights!
        # Shape: [output_size, input_shard_size]
        # Note: PyTorch Linear weights are typically (out_features, in_features)
        self.weight = nn.Parameter(torch.randn(output_size, self.input_shard_size))
        
        # Bias is usually handled by one rank or replicated and reduced.
        # Simplified here: no bias for clarity.
 
    @torch.no_grad()  # Forward-only illustration; training needs autograd-aware collectives.
    def forward(self, x):
        # x input shape: [batch_size, input_size]
        
        # 4. Scatter Input (Simulating the split)
        # In a real Transformer, x might already be sharded from the previous 
        # Column-Parallel layer. Here we manually shard it to simulate the state.
        # Split x along the last dimension (columns)
        input_shards = list(x.chunk(self.world_size, dim=-1))
        local_input = input_shards[self.rank].contiguous()
        
        # 5. Local MatMul
        # Compute the partial result for this GPU's slice of data
        # local_output shape: [batch_size, output_size]
        # We use .t() because PyTorch Linear weights are stored transposed
        local_output = torch.matmul(local_input, self.weight.t())
        
        # 6. All-Reduce (The Communication Bottleneck)
        # Sum the partial results from all GPUs into every GPU
        # The next dependent operation waits for the reduction result.
        dist.all_reduce(local_output, op=dist.ReduceOp.SUM)
        
        return local_output
 
# Usage pseudo-code:
# dist.init_process_group(backend='nccl')
# layer = RowParallelLinear(1024, 512).cuda()
# y = layer(x)

In diesem Code ist die Zeile dist.all_reduce der kritische Pfad. Die nächste abhängige Operation wartet auf die Reduktion. Kollektivalgorithmus und Überlappung hängen von Runtime, Nachrichtengröße und Topologie ab.

3. Vor- und Nachteile: Eine Hardware-Perspektive

Vorteile

Speicherreduktion: Bei teilbaren Schichten kann TP deren Parameter, Gradienten und zugehörige Optimiererzustände pro Rank in Richtung $1/N$ senken. Replizierte Gewichte, Aktivierungen und Kommunikationspuffer halten den Gesamtspeicher über diesem Idealwert.
Zugriff auf riesige Modelle: Es ermöglicht das Training von Modellen, bei denen die Gewichte einer einzelnen Schicht einfach zu groß sind, um in den VRAM einer einzelnen GPU zu passen.
Parallele Schichtarbeit: TP-Ranks berechnen Teile derselben Schicht gleichzeitig. Ob dies die Latenz gegenüber Pipeline-Parallelität verbessert, hängt von Batch, Modell und Kommunikation ab.

Nachteile

Kommunikation auf dem kritischen Pfad: Aufgeteilte Schichten benötigen im Vorwärts- und Rückwärtspfad oft Kollektivoperationen. Eine abhängige Operation wartet auf das Ergebnis; optimierte Zeitpläne können manche Kommunikation mit unabhängiger Arbeit überlappen.
Bandbreitenlimit: Aufgrund der hohen Synchronisationsfrequenz (oft mehrfach pro Transformer-Block) ist TP stark von der Bandbreite der Verbindungen abhängig.
- Intra-Node (NVLink): TP profitiert häufig von schnellen Verbindungen innerhalb eines Knotens. H100 NVLink 4 nennt 900 GB/s bidirektional pro GPU (etwa 450 GB/s je Richtung vor Protokoll- und Kollektivaufwand).
- Inter-Node (Ethernet/InfiniBand): Knotenübergreifendes TP kann langsamer sein, wenn Kollektivverkehr dominiert; es ist nicht grundsätzlich unbrauchbar. In einem berichteten DeepSeek-V3-Experiment mit 16 H100 erreicht SiPipes pipeline-basiertes Design bis zu 4,5× den Durchsatz der reinen TP-Baseline des Papers ¹. Dieses Ergebnis gilt für den untersuchten Workload.
Implementierungskomplexität: TP erfordert Shard-fähige Schichten und Kollektivoperationen, die man manuell oder mit einem Framework implementieren kann.

Übersichtstabelle: Wann sollte TP verwendet werden?

Szenario	Empfehlung	Hardware-Begründung
Einzelner Knoten (<= 8 GPUs)	Prüfen	Schnelle Verbindungen helfen; Modellgröße und Batch bestimmen den Nutzen.
Mehrere Knoten (> 8 GPUs)	Messen	Netzwerk-Kollektivoperationen können dominieren; TP, PP und Mischformen vergleichen.
Riesige Gewichtsmatrizen	Eine Möglichkeit	TP teilt eine Schicht; andere Partitionierungs- und Offload-Verfahren haben andere Abwägungen.

Die Rolle der 3D-Parallelisierung

Große Trainingssysteme kombinieren TP oft in einer schnell verbundenen GPU-Gruppe mit PP und DP; Platzierung und Parallelitätsgrade hängen von Modell und Cluster ab.

Referenzen

He, Zhao & Cao (ScitiX AI), SiPipe: Bridging the CPU–GPU Utilization Gap for Efficient Pipeline-Parallel LLM Inference, arXiv:2506.22033 (June 2025). https://arxiv.org/abs/2506.22033 ↩