Das Dilemma des Routers: Eine hardwarezentrierte Tiefenanalyse der Expert Parallelism

Wir haben Tensor Parallelism (TP) untersucht, das einzelne Operationen aufteilt, um massive Schichten in den Speicher zu passen, und Pipeline Parallelism (PP), das Schichten vertikal aufteilt, um über mehrere Knoten zu skalieren. Doch während Modelle nicht nur in Tiefe oder Breite, sondern auch in Vielfalt wachsen – insbesondere mit Mixture-of-Experts (MoE)-Architekturen wie Mixtral 8x7B oder DeepSeek-V3 – stoßen wir auf einen neuen Engpass.

In einem dichten Modell wird jeder Parameter für jedes Token verwendet. In einem MoE-Modell ist nur ein Bruchteil der Parameter (Experten) pro Token aktiv. Das Replizieren der massiven, meist inaktiven Expertengewichte auf jeder GPU (Data Parallelism) ist eine Verschwendung von VRAM. Das Aufteilen mit TP ist ineffizient, da die Rechenintensität pro Experte zu gering ist, um die feinkörnige Synchronisation zu rechtfertigen.

Hier kommt Expert Parallelism (EP) ins Spiel. In diesem Paradigma verteilen wir die Experten selbst auf verschiedene GPUs. GPU 0 hält Experte A, GPU 1 hält Experte B. Die Herausforderung verlagert sich vom Aufteilen von Matrizen zum Routing der Tokens: Wir müssen Daten (Tokens) physisch zu der GPU verschieben, die die benötigten Parameter zur Verarbeitung enthält.

Dieser Artikel analysiert die Hardware-Mechanik von EP, das "All-to-All"-Kommunikationsprinzip, das sie definiert, und die entscheidenden Kompromisse, die dabei eine Rolle spielen.

1. Die Hardware-Mechanik: Token-Routing

Bei Expert Parallelism ist das Modell physisch fragmentiert. Das entscheidende Hardware-Ereignis ist nicht mehr eine Matrixmultiplikations-Synchronisation (wie bei TP) oder ein Pipeline-Flush (wie bei PP); es ist ein massives Umschichten von Daten über die Interconnect.

Der Vier-Stufen-Lebenszyklus

Für eine einzelne MoE-Schicht sieht der Hardware-Ausführungsablauf folgendermaßen aus:

Route (Lokale Berechnung): Jede GPU verarbeitet ihr Token-Batch durch ein kleines "Gating Network" (in der Regel eine einzelne Linear-Schicht). Dieser Router berechnet, welchen Experten (z. B. Expert 0, 1, ... 7) jedes Token aufsuchen muss.
- Hardware-Status: Tokens befinden sich derzeit auf der "Quell"-GPU.
Dispatch (Der Kommunikations-Engpass): Dies ist der entscheidende Moment von EP. GPU 0 könnte 100 Tokens haben, die zu Expert 5 (auf GPU 5) müssen. GPU 1 könnte 50 Tokens für Expert 5 haben. Jede GPU muss bestimmte Teilmengen ihrer Daten gleichzeitig an jede andere GPU senden. Dies erfordert eine All-to-All-Kollektivoperation.
- Hardware-Auswirkung: Dies belastet die Bisection Bandbreite des Clusters. Im Gegensatz zu den ringbasierten Mustern von TP erzeugt All-to-All ein dichtes Netz an Datenverkehr.
Expert Compute (Lokale Berechnung): Sobald die Tokens bei ihren Ziel-GPUs (dem "Zuhause" ihrer zugewiesenen Experten) angekommen sind, führt die GPU die Standard-Feed-Forward-Network (FFN)-Berechnung durch.
- Hardware-Vorteil: Die arithmetische Intensität ist hier hoch (große Batch-Größe pro Experte), und während der aufwendigen Berechnungen findet keine Kommunikation statt.
Combine (Kommunikation): Die verarbeiteten Tokens müssen zu ihren ursprünglichen Quell-GPUs zurückkehren, um die Reihenfolge der Sequenz für die nächste Schicht (z. B. Attention) zu erhalten. Dies erfordert eine zweite All-to-All-Operation, um die Ergebnisse zurückzusenden.

Die Hardware-Realität: All-to-All

Das All-to-All-Primitive unterscheidet sich von All-Reduce.

All-Reduce (TP/DP): Jede GPU verfügt über einen Puffer gleicher Größe. Wir summieren sie auf. Das Volumen ist fest.
All-to-All (EP): GPU $i$ sendet einen personalisierten Puffer an GPU $j$ . Das Volumen kann dynamisch und ungleichmäßig sein (Lastungleichgewicht), was zu Netzwerküberlastung führen kann, wenn eine GPU deutlich mehr Daten als andere empfängt.

2. Bare Metal Implementierung: Reines PyTorch

Um die Routing-Mechanik zu verstehen, implementieren wir eine vereinfachte Expert Parallel-Schicht. Wir verwenden torch.distributed.all_to_all_single, das einen flachen Token-Puffer basierend auf Splits durchmischt.

import torch
import torch.nn as nn
import torch.distributed as dist
 
class ExpertParallelLayer(nn.Module):
    def __init__(self, num_experts, input_dim, hidden_dim, world_size, my_rank):
        super().__init__()
        self.world_size = world_size
        self.my_rank = my_rank
        
        # 1. Assign Experts to Ranks
        # Assume we have N experts and N GPUs, so 1 expert per GPU.
        # In this setup, I strictly hold the weights for the expert matching my_rank.
        self.local_expert = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )
        
        # A tiny router to decide where tokens go
        self.router = nn.Linear(input_dim, num_experts)
 
    def forward(self, x):
        # x shape: [batch_size, seq_len, input_dim]
        # Flatten batch and sequence for routing
        tokens = x.view(-1, x.size(-1)) 
        
        # 2. Routing Decision
        # logits: [num_tokens, num_experts]
        router_logits = self.router(tokens)
        # expert_indices: [num_tokens] -> which expert (rank) does each token need?
        expert_indices = torch.argmax(router_logits, dim=-1)
 
        # 3. Prepare for Dispatch (All-to-All)
        # We need to group tokens by their destination rank.
        # Calculate how many tokens we are sending to each rank.
        # send_counts: list of length world_size
        send_counts = [0] * self.world_size
        for dest_rank in expert_indices:
            send_counts[dest_rank.item()] += 1
            
        # Sort tokens so they are grouped by destination rank
        # (Skipping complex sort/index tracking code for brevity)
        sorted_tokens = self._sort_tokens_by_dest(tokens, expert_indices)
        
        # We need to tell other GPUs how many tokens to expect from us
        # (This usually requires a small all-to-all of simple integers first)
        recv_counts = list(torch.zeros(self.world_size, dtype=torch.long))
        dist.all_to_all(recv_counts, torch.tensor(send_counts)) 
        
        # 4. The Big Shuffle: All-to-All Dispatch
        # send_buffer: My sorted tokens
        # recv_buffer: The tokens coming to ME from everyone else
        recv_total = sum(recv_counts)
        recv_buffer = torch.empty(recv_total, tokens.size(-1), device=x.device)
        
        # This moves the actual high-bandwidth tensor data across the interconnect
        dist.all_to_all_single(
            recv_buffer, 
            sorted_tokens, 
            output_split_sizes=recv_counts, 
            input_split_sizes=send_counts
        )
        
        # 5. Local Expert Compute
        # I now possess all tokens that "chose" me. I process them.
        # Note: This is standard high-intensity compute!
        processed_tokens = self.local_expert(recv_buffer)
        
        # 6. The Big Return: All-to-All Combine
        # Send the processed tokens back to their owners.
        # We simply reverse the send/recv counts.
        return_buffer = torch.empty_like(tokens)
        
        dist.all_to_all_single(
            return_buffer,
            processed_tokens,
            output_split_sizes=send_counts,  # I expect back what I sent
            input_split_sizes=recv_counts    # I send back what I received
        )
        
        # (Unsort tokens back to original order and return)
        return self._unsort_tokens(return_buffer)
 
    def _sort_tokens_by_dest(self, tokens, indices):
        # Implementation of sorting/grouping logic
        pass

Hinweis: Reale Implementierungen wie Megatron-LM oder DeepSpeed verwenden komplexes Padding und Index-Sortierung, um dynamische Token-Anzahlen effizient zu verarbeiten. Dieses Beispiel vereinfacht die Sortierlogik zur besseren Verständlichkeit.

In dieser Implementierung tauschen die GPUs effektiv ihre Rollen: Sie sind nicht länger "Eigentümer eines Batch-Slices", sondern werden vorübergehend zu "Compute-Servern für bestimmte Operationen" und verarbeiten Tokens aus dem gesamten Cluster.

3. Hardware-Einschränkungen & Kompromisse

Der All-to-All-Latenzsprung

Im Gegensatz zu TP (das über ultraschnelles NVLink läuft) oder PP (das sehr wenige Daten sendet), überträgt EP riesige Datenmengen (Aktivierungen) über das Netzwerk.

Inter-Node-Flaschenhals: Wenn Ihre Experten auf verschiedene Serverknoten verteilt sind, muss die All-to-All-Operation das Ethernet/InfiniBand-Netzwerk durchqueren. Dies ist deutlich langsamer als die intra-node NVLink-Verbindungen.
Latenzempfindlichkeit: Bei kleinen Batch-Größen (Inference) kann die Latenz des All-to-All-Schritts die Ausführungszeit dominieren, wodurch EP langsamer als einfache Datenparallelität wird.

Der Albtraum des Lastenausgleichs

Hardware mag keine Unregelmäßigkeiten. GPUs sind am schnellsten, wenn sie statische, gleichmäßige Strukturen verarbeiten.

Das Problem: Wenn der Router entscheidet, dass 90% der Tokens zu Expert 0 (GPU 0) gehen müssen, wird GPU 0 zum Flaschenhals, während GPUs 1–7 untätig bleiben. Dies ist Rechen-Skew.
Die Konsequenz: Die All-to-All-Operation wird blockiert, da alle darauf warten, Daten an die überlastete GPU 0 zu senden.
Die Lösung (Hardware-Ebene): Häufig verwerfen wir Tokens. Wenn der Puffer von GPU 0 voll ist, verwerfen wir einfach die zusätzlichen Tokens (oder verarbeiten sie mit einem Dummy-Durchlauf), um die Synchronisation zu erhalten und tauschen Genauigkeit gegen Durchsatz.

Zusammenfassung der Vor- und Nachteile

Vorteile

Massiver Parameterskalierung: Sie können Modelle mit Billionen von Parametern trainieren (wie GPT-4 oder Switch Transformer), da Sie nur einen Bruchteil der Gewichte ( $1/N$ ) auf jeder GPU speichern.
Recheneffizienz: Anders als bei dichten Modellen, bei denen die FLOPs linear mit der Parameteranzahl skalieren, ermöglicht MoE/EP die Skalierung der Parameter ohne die Rechenkosten pro Token zu erhöhen. Sie erhalten "intelligentere" Modelle bei gleichen Inferenzkosten.
Asynchrones Potenzial: Fortschrittliche Implementierungen können die Dispatch-Kommunikation mit der Router-Berechnung überlappen.

Nachteile

Kommunikationsaufwand: Die All-to-All-Primitve ist aufwendig. Sie erfordert eine hohe Bisection-Bandbreite über den gesamten Cluster.
Lastverteilung: Erfordert während des Trainings komplexe Hilfsverluste (Lastverteilungsverlust), um sicherzustellen, dass der Router die Arbeit gleichmäßig verteilt. Ohne dies bricht die Hardwareauslastung zusammen.
Speicherfragmentierung: Die dynamische Natur des Token-Routings (unterschiedliche Puffergrößen bei jedem Schritt) zwingt Speicherallokatoren zu erhöhter Arbeit, was häufig zu Fragmentierung oder dem Bedarf an komplexen Padding-Strategien führt.