Effiziente und portable Mixture-of-Experts-Kommunikation

Wir präsentieren eine hochleistungsfähige, portable Open-Source-Bibliothek für Mixture-of-Experts (MoE)-Kommunikation, die eine 10-mal schnellere Leistung im Vergleich zu Standard-All-to-All-Kommunikationsprimitiven erzielt.

Unsere Implementierung umfasst mehrere wichtige technische Innovationen, die eine überlegene Effizienz der MoE-Kommunikation liefern:

GPU-initiierte Kommunikation (IBGDA): Unterstützt die direkte GPU-zu-NIC-Kommunikation und reduziert die Latenz erheblich, indem die CPU-Einbindung umgangen wird.

Überlappung von Kommunikation und Berechnung: Eine Split-Kernel-Architektur mit getrennten Sende- und Empfangsphasen ermöglicht es, dass Berechnungen gleichzeitig mit Netzwerktransfers ablaufen.

Schnellste Single-Node-Leistung: 2,5-mal geringere Latenz als bei der bisher schnellsten Implementierung in Single-Node-Konfigurationen.

Effiziente und portable Multi-Node-Leistung: Unsere Implementierung erreicht Geschwindigkeiten, die bis zu 10-mal schneller sind als bei der Standard-All-to-All-Kommunikation. Obwohl unser Ansatz etwa 2-mal langsamer als hochspezialisierte Implementierungen ist, bietet er eine bessere Portabilität über NVSHMEM-Versionen und Netzwerkumgebungen (NVLink, CX-7 und EFA) hinweg.

Die Bibliothek ist vollständig Open-Source und verfügbar unter https://github.com/ppl-ai/pplx-kernels.

In diesem Artikel untersuchen wir die Herausforderungen der Experten-Parallelisierung in MoE-Modellen im großen Maßstab und beschreiben unseren Ansatz für eine effiziente Token-Verteilung und -Kombination in verteilten GPU-Umgebungen.

Einführung

Mixture-of-Experts (MoE)-Modelle, wie DeepSeek R1 und Mixtral 8x7B, verbessern dense Modelle, indem sie die Anzahl der Gewichte begrenzen, die für jedes Token aktiviert werden. Anstatt hohe Parameterzahlen durch eine Erhöhung der Anzahl der Schichten oder eine Vergrößerung der linearen Schichten im Multi-Layer Perceptron (MLP) jeder Decoder-Schicht zu erreichen, ersetzen MoE-Modelle das traditionelle MLP durch mehrere Experten und einen Router. Beispielsweise werden von den 671B Parametern von DeepSeek R1 während der Inferenz nur 37B mit einem gegebenen Token multipliziert. Dies reduziert den für die Dekodierung eines Tokens erforderlichen Rechenaufwand erheblich und verbessert die Latenz im Vergleich zu einem dense Modell.

MoE-Modelle stellen im Vergleich zu dense Modellen zusätzliche Herausforderungen an die Inferenz. Während die Experten selbst kleine MLP-Schichten sind, enthält jede Decoder-Schicht einen Router, der entscheidet, an welche Experten ein Token verteilt wird, wobei jedes Token an mehrere Experten verteilt wird. Der Router ist typischerweise eine kleine lineare Schicht, die eine Wahrscheinlichkeitsverteilung erzeugt. Üblicherweise werden die Experten mit den Top-K-Scores ausgewählt und die endgültige Aktivierung als gewichteter Durchschnitt berechnet, wobei die Expertenausgaben mit einem aus der Wahrscheinlichkeitsverteilung des Routers abgeleiteten Gewicht multipliziert werden.

Um die Latenz in verteilten MoE-Systemen zu minimieren, kann Parallelisierung über mehrere Geräte genutzt werden, was jedoch Kommunikationsherausforderungen mit sich bringt. Modelle wie Mixtral 8x7B oder Llama-70B passen in 8 Geräte über einen einzelnen Node, wobei sie von schnellen NVLink-Verbindungen mit geringer Latenz (bis zu 900 Gbps) profitieren. Größere Modelle erfordern jedoch ein Sharding der Experten über mehrere Nodes mittels InfiniBand (mit Spitzenwerten bei 400 Gbps), was zusätzliche Latenzherausforderungen mit sich bringt.

Vergrößern

In diesem Artikel untersuchen wir das Problem der Experten-Parallelisierung und beschreiben unsere Implementierung von GPU-Kerneln, die Token effizient an Geräte verteilen, die die zugewiesenen Experten halten, und diese durch Berechnung der gewichteten Summe der Expertenausgaben wieder einsammeln. Während aggressiv optimierte Implementierungen für spezifische Netzwerkhardware existieren, stellen wir eine hochgradig portable Implementierung vor, die auf einer minimalen Menge an NVSHMEM-Primitiven basiert, aber dennoch eine 10-mal schnellere Leistung als die Standard-All-to-All-Kommunikation liefert. Unsere Implementierung erzielt eine erstklassige Leistung bei Single-Node-Konfigurationen bei gleichzeitig exzellenter Portabilität über verschiedene Netzwerkumgebungen hinweg.

Sharding und Parallelität

Für eine effiziente Inferenz müssen die Gewichte eines Modells im Gerätespeicher gehalten werden, wobei gleichzeitig ausreichend Platz für Aktivierungen, KV-Caches und andere Puffer verbleiben muss, die für den Vorwärtsdurchlauf durch das Modell erforderlich sind. Die leistungsfähigsten Modelle übersteigen die Kapazität selbst der leistungsfähigsten GPUs, daher muss die Inferenz auf mehrere Geräte verteilt werden, die gemeinsam die Gewichte des Modells speichern können. Basierend auf den Sharding-Schemata der Gewichte müssen unterschiedliche Kommunikations- und Berechnungspläne verwendet werden, um die Berechnungen auszulagern und die Geräte zu synchronisieren.

Vergrößern

Expert Parallelism (EP), wie in der ersten Abbildung dargestellt, parallelisiert nur die Expertenberechnung. Verschiedenen Geräten werden unterschiedliche Experten zugewiesen, die deren Gewichte halten. Nach dem Routing werden Token an das entsprechende Gerät gesendet, wobei die Ergebnisse anschließend gesammelt und akkumuliert werden. Die Komplexität des Routings hängt vom Grad der Parallelität in den anderen Teilen des Modells ab: Die Replikation anderer Schichten könnte die Notwendigkeit für das Routing insgesamt eliminieren, da jeder Rang die Token aus einer lokal replizierten Routing-Tabelle auswählen kann. Wenn jedoch nur einer der Ränge das Routing ausführt, ist ein Broadcast erforderlich, um Token, Indizes und Gewichte an ihre jeweiligen Experten zu verteilen. Abschließend synchronisiert ein All-Gather oder ein All-to-All-Broadcast die Ausgabe-Token mit dem Rang, der die Ausführung des Modells fortsetzt. Eine solche Implementierung ist relativ einfach, da torch bereits die erforderlichen Primitiven bereitstellt, wenngleich einige Vorteile durch die Fusion der Akkumulation mit der Gather-Operation zum Sammeln der Token von den Experten erzielt werden könnten.

Vergrößern

Expert-Only-Parallelität skaliert nicht ideal, da Nodes in einem Cluster im Leerlauf sein könnten, während das Modell Nicht-Experten-Schichten wie Attention, Norm und Sampling ausführt. Die Berechnung dieser Schichten, primär Attention, kann jedoch ebenfalls von Tensor Parallelism (TP) profitieren. Die meisten Modelle basieren auf Multi-Head-Attention, was bedeutet, dass die Attention-Heads und ihre entsprechenden Q-, K- und V-Projektionen ebenfalls über Geräte hinweg geshardet, repliziert oder zwischen verschiedenen Schichten gesammelt werden können. Wenn Attention über alle Geräte verteilt ist, kann ein All-Gather die Aktivierungen synchronisieren, was eine Replikation des Routings ermöglicht und für eine effiziente Implementierung Synchronisationsprimitive erfordert, die dem Fall der Experten-Parallelisierung ähneln. Es gibt jedoch Grenzen für die Parallelität auf dieser Ebene, da eine Verringerung der Anzahl der Attention-Heads unter einen bestimmten Schwellenwert zu abnehmenden Erträgen führt.

Vergrößern

Um alle Geräte optimal zu nutzen und einen sehr hohen Grad an Experten-Parallelität von bis zu 128 oder 256 GPUs zu unterstützen, ist Data Parallelism (DP) erforderlich. Bei diesem Schema werden die Geräte, auf die ein Modell aufgeteilt ist, gruppiert, um Anfragen gleichzeitig zu bearbeiten, wobei Attention berechnet und KV-Caches über ihre lokale Gruppe hinweg geshardet werden. Mehrere Instanzen dieser parallelen Gruppen arbeiten bei der Expertenbewertung zusammen, wobei jede eine unterschiedliche Teilmenge der Expertengewichte hostet. Basierend auf der Anzahl der Attention-Heads kann eine Gruppe typischerweise bis zur Größe eines gesamten Nodes umfassen, da die Intra-Node-Kommunikation schneller ist. In der Abbildung oben bedient beispielsweise ein DP-Rang unabhängig zwei Anfragen, A und B, und wickelt Attention, Norm und jede andere Buchhaltung für die Anfragen ab. Der andere DP-Rang verarbeitet eine separate Anfrage C. Da jedoch der erste Node die Hälfte der Experten hostet, während der andere Node die andere Hälfte hostet, können nach dem Routing Token von A und B an den zweiten Node gesendet werden und umgekehrt. Dies führt zu einem spärlichen Kommunikationsproblem: Jedes Gerät sendet möglicherweise eine unterschiedliche Anzahl von Token an einen anderen Zielrang. Bestehende Primitiven von torch, primär all_to_all, sind nicht besonders gut geeignet, da sie möglicherweise eine Form von Padding oder GPU-zu-CPU-Synchronisation und Metadaten-Broadcast erfordern. Um die Kommunikation effektiv zu implementieren, sind benutzerdefinierte Kernel erforderlich, um die Routing-Informationen zu verarbeiten und die Kommunikation von den Geräten ohne CPU-Eingriff zu initiieren. Nach dem Routing muss ein Dispatch-Kernel Token an die Ränge senden, an die sie geroutet wurden, während auf der Kombinierungsseite die Aktivierungen, die zu den Anfragen in der aktuellen DP-Gruppe gehören, gesammelt werden müssen. Zusätzlich muss die Arbeit innerhalb einer DP-Gruppe, die mehrere Geräte haben kann, ausgeglichen werden, um Broadcasts zu minimieren und die Token auf allen Rängen für die nachfolgenden Schichten korrekt zu replizieren.

NVSHMEM

NVSHMEM ist eine NVIDIA-spezifische OpenSHMEM-Implementierung, die portable Inter-Device-Kommunikationseinrichtungen bereitstellt, welche die Komplexität der zugrunde liegenden Hardware abstrahieren. Die API kann Device-to-Device-Lese- und Schreibvorgänge ausdrücken, die auf die Primitiven der individuellen Transportschichten abgebildet werden. Sowohl NVLink als auch RDMA werden unterstützt, was einen Grad an Portabilität gewährt und es NVSHMEM-Kerneln ermöglicht, innerhalb eines Nodes oder über einen Cluster von mehreren Nodes hinweg zu arbeiten. In unseren Kerneln verwenden wir Remote-Speicherschreibvorgänge und Signalisierungsoperationen, um Daten zu übertragen und die teilnehmenden Geräte zu synchronisieren.

NVSHMEM-Operationen basieren auf dem Konzept des symmetrischen Speichers: Sie operieren auf Puffern, die auf allen Geräten zugewiesen wurden, die an der Inter-Device-Kommunikation teilnehmen. Ein lokales Gerät kann eine innerhalb seines eigenen lokalen Puffers abgeleitete Adresse zusammen mit dem Index des Zielrangs verwenden, um das Ziel einer Operation festzulegen. Die folgende Abbildung veranschaulicht dieses Konzept: Beide GPUs weisen symmetrische Puffer der gleichen Größe zu und behalten src-Pointer auf diese. Der erste Rang möchte 3 Ganzzahlen an den zweiten senden und platziert sie am Anfang des Puffers: nvshmem_int_put_nbi leitet die Startadresse vom lokalen Puffer ab und spezifiziert das Zielgerät. Der zweite Rang leitet einen Offset von seinem eigenen Puffer ab und sendet ein Element an das erste Gerät, wobei er um eins versetzt ist. Während Zieladressen immer symmetrische Puffer sein müssen, die mittels nvshmem_alloc zugewiesen wurden, können Quellpuffer beliebige Bereiche des Gerätespeichers sein, sofern sie vorregistriert sind mit NVSHMEM.

Während NVSHMEM eine breite Palette an Primitiven bietet, verlassen sich unsere Kernel nur auf 3 Funktionen und bauen darauf die gesamte Synchronisation und das Fencing auf.

Vergrößern

nvshmemx_putmem_signal_nbi_warp: Überträgt einen Datenblock von einem Gerät auf ein anderes, während gleichzeitig ein Flag auf dem Remote-Gerät gesetzt wird. Die Operation setzt (NVSHMEM_SIGNAL_SET) oder inkrementiert (NVSHMEM_SIGNAL_ADD) einen 64-Bit-Speicherort. Das Flag wird aktualisiert, nachdem der gesamte Speicherblock übertragen wurde. Wenn das Remote-Gerät eine Änderung des Flags beobachtet hat, kann es sicher auf den Puffer in seinem eigenen Speicher zugreifen. Diese Funktion ist nützlich für die Zusammenführung von Datentransfer und Synchronisation.

nvshmemx_signal_op operiert auf einem einzelnen Speicherort, typischerweise einem 64-Bit-Flag, und setzt oder inkrementiert diesen atomar. Es ist nützlich beim Senden von Metadaten und bei der Synchronisation von Geräten.

nvshmem_uint64_wait_until wird am empfangenden Ende eines Signals verwendet, um ein Flag abzufragen, bis das Remote-Gerät es aktualisiert.

Die Sicherstellung irgendeiner Form von Reihenfolge zwischen Operationen durch nvshmem_quiet oder nvshmem_fence ist kostspielig, daher vermeidet unsere Implementierung der Kernel die Verwendung anderer NVSHMEM-Primitive und minimiert alle Abhängigkeiten außer individuellen Sende-Empfangs-Paaren. Beim Senden von Daten verwenden wir immer die nicht-blockierende Version der Funktionen, ohne darauf zu warten, dass die Daten vom lokalen Rang überhaupt gesendet wurden. Andere, implizite Synchronisationen über die Gruppe der Kernel stellen sicher, dass diese Puffer nicht destruktiv überschrieben werden, während die NICs Zugriff auf sie benötigen.

NVSHMEM verwendet GPUDirect RDMA, sodass die NIC direkt GPU-Speicher für den Datenaustausch lesen und schreiben kann, ohne CPUs einzubeziehen. Darüber hinaus unterstützt NVSHMEM auf ConnectX-NICs die GPU-initiierte Kommunikation (auch bekannt als Infiniband GPUDirect Async oder IBGDA), was es der GPU ermöglicht, RDMA-Operationen direkt an die NIC zu posten und die Completion Queue der NIC direkt abzufragen. Auf Plattformen, die keine GPU-initiierte Kommunikation unterstützen, startet NVSHMEM einen „Proxy“-Thread auf der Host-CPU, um die Kommunikation im Namen der GPU zu initiieren und die Fertigstellung abzufragen. Ein NVSHMEM-Programm ist portabel, unabhängig davon, ob GPU-initiierte Kommunikation unterstützt wird oder nicht. Die GPU-initiierte Kommunikation reduziert jedoch die Latenz erheblich, da sie den Umweg über die CPU vollständig umgeht.

Portable Kernel

Wir implementieren MoE-Kommunikation durch ein Paar von Dispatch- und Combine-Kerneln. Die Dispatch-Kernel sind dafür verantwortlich, Token und Routing-Informationen auf jedem Rang zu lesen und sie an die entsprechenden Experten zu verteilen. Die Combine-Kernel sammeln die von den Experten produzierten Aktivierungen und senden sie an ihre Ursprungsränge zurück, während sie gleichzeitig den gewichteten Durchschnitt der ausgewählten Experten basierend auf den vom Router berechneten Gewichten berechnen. Die Kernel sind weiter in eine Sende- und Empfangskomponente unterteilt, um es zu ermöglichen, dass Datentransfers mit Berechnungen überlappt werden. Die Sendekernel sind alle nicht-blockierend und nicht-synchronisierend: Sie leiten einfach alle Schreibvorgänge an die Remotes weiter. Auf der anderen Seite lesen die Combine-Kernel nur aus dem Speicher und warten, bis alle erforderlichen Daten übertragen wurden. Nachdem die Arbeit an die NICs verteilt wurde, können die GPUs während der asynchronen Datenübertragung über das Kabel andere nützliche Arbeiten lokal ausführen, wie etwa das Anwenden geteilter Experten oder das Berechnen von Attention.

Jedes Paar von Kerneln hat seine eigenen symmetrischen Speicherpuffer, wobei jeder Rang privaten Speicher zuweist, um Daten von allen anderen Rängen zu empfangen. Die maximale Anzahl von Token, die von einer DP-Gruppe gesendet werden können, ist über die Ränge hinweg fixiert, was auch die Obergrenze festlegt, die jeder Rang für jeden lokalen Experten von jedem DP-Rang empfangen kann. Dies ermöglicht es Sender-Rängen, eine eindeutige Adresse auf dem Zielrang abzuleiten, an die geschrieben werden soll, ohne dass eine Synchronisation zwischen ihnen erforderlich ist. Nach dem Datenempfang mischen die Kernel die Token, um sie im Speicher in dem optimalen Format anzuordnen, das von den empfangenden Kerneln benötigt wird. Obwohl die Puffer eine beträchtliche Dimension haben, werden sie über alle sequenziellen Schichten eines Modells hinweg wiederverwendet.

Vergrößern

Die einzige Form der globalen Synchronisation, wie in der Abbildung oben illustriert, ist im Combine-Receive-Kernel implementiert. Sobald Daten vom Senden verteilt werden, kehrt ein Rang zurück und setzt die Ausführung fort, wenn er alle Token empfängt, die er von den anderen Rängen benötigt, und tritt in die Combine-Send-Kernel ein, sobald diese verarbeitet sind. Der Barrier im Combine-Receive stellt sicher, dass kein Rang vorauslaufen und Dispatch-Send starten kann, während irgendein anderer Rang noch auf den Empfang von Daten wartet, was potenziell destruktive Überlappungen verursachen könnte. Die Synchronisation erfolgt aus Gründen der Einfachheit im Combine-Kernel. Als Optimierung könnte derselbe Barrier zwischen dem Dispatch-Receive und den Combine-Send-Kerneln implementiert werden, was die Latenz weiter reduziert, indem den NICs mehr Zeit gelassen wird, die Transfers abzuschließen, während die GPU andere Kernel ausführt.

Beide Kernel sind über alle verfügbaren SMs der Geräte verteilt: Während die Dispatch-Send- und Combine-Receive-Kernel über eine pro-Rang maximale Tokenanzahl (max_m) parallelisieren müssen, müssen die Dispatch-Receive- und Combine-Send-Kernel eine entsprechende Anzahl von Token von jedem Rang und lokalen Experten verarbeiten, für ein Maximum von max_m * (num_experts // EP) * (EP // TP).

Dispatch

Der Dispatch-Kernel ist dafür verantwortlich, Token und Block-Skalierungsfaktoren an Peers zu senden, die die Experten hosten, an die die Token geroutet wurden. Die Senderseite verlässt sich auf Warp-Spezialisierung, um zwei Aufgaben zu parallelisieren: das Aggregieren von Routing-Informationen zur Bestimmung, wie viele Token von jedem Rang gesendet werden, und das Packen der Token in Nachrichten, die an den Remote-Empfänger gesendet werden sollen. Die Empfängerseite wartet zuerst auf den Empfang aller Token-Zahlen, dann liest und entpackt sie die Payloads aus dem Shared Memory in die Tensoren, die an die nachgelagerten Grouped GEMM-Kernel weitergegeben werden. Sie speichert auch Informationen in Puffern, die mit den Combine-Kerneln geteilt werden, um anzuzeigen, wohin jedes Token zurückgeschickt werden soll. Diese Information ist erforderlich, da der Empfangskernel Token in zusammenhängenden Puffern in einer nicht-deterministischen Reihenfolge verschiebt.

Vergrößern

Auf der Senderseite ist ein einzelner Warp von 32 Threads pro Block dafür verantwortlich, die Routing-Informationen zu lesen und die Anzahl der Token zu zählen, die jeder Ziel-Experte von diesem Rang voraussichtlich empfangen wird, wobei die Anzahl plus eins mittels nvshmem_signal_op gesendet wird. Die Anzahl wird um eins erhöht, da der Übergang von Null auf Nicht-Null auf der Remote-Seite den Empfang der Zahlen von einem Quellrang signalisiert. Parallel dazu arbeiten die verbleibenden Warps zusammen, um Token in den symmetrischen Speicher über alle Blöcke hinweg parallel zu kopieren, wobei Aktivierungen, Skalierungsfaktoren und ihr Index auf dem lokalen Rang zu einem zusammenhängenden Block gepackt werden. Der Index wird vom Combine-Sender benötigt, um die Adresse zu bestimmen, an die das Token geschrieben wird. Danach, nachdem sichergestellt wurde, dass alle Daten durch einen Barrier über die Warp-Gruppen hinweg kopiert wurden, operieren die Warps erneut unabhängig voneinander und senden jeweils denselben Puffer parallel an einen anderen Experten. Die Token werden mittels nvshmemx_putmem_signal_nbi_warp gesendet, was auch die Anzahl der vom lokalen Rang gesendeten Token auf dem Remote-Gerät atomar erhöht. Innerhalb einer DP-Gruppe, da jeder Rang ein Replikat des zu sendenden Tokens besitzt, wird das Dispatching gleichmäßig ausgeglichen, wobei jedes Gerät eine Teilmenge der Token aussendet.

Auf der Empfängerseite synchronisieren sich alle Kernel zuerst mit den Sender-Rängen, indem sie darauf warten, dass die gesamten Token-Zahlen gesendet werden. Danach warten sie alle darauf, dass sich die atomar inkrementierten gesendeten Token-Zahlen auf die Gesamtzahl einpendeln, was dank der Semantik des putmem-Aufrufs anzeigt, dass auch alle Payloads von den Quellrängen gesendet wurden. Die Kernel fragen die Zahlen mittels nvshmem_uint64_wait_until ab und parallelisieren die Operation über alle Blöcke und Threads. Anschließend stellt ein Cross-Block-Barrier sicher, dass kein Block die Puffer liest, es sei denn, alle Daten wurden korrekt empfangen. Über Blöcke verteilt und über einen atomar inkrementierten Zähler synchronisiert, werden die Token aus dem symmetrischen Speicher in reguläre zusammenhängende Puffer kopiert, die zu den Eingaben für die nachgelagerten Kernel werden, die die Experten implementieren. Der Quellrang, der Experten-Index und der Token-Index werden separat gespeichert und lokalisieren genau die Stelle, an die der Combine-Kernel die Aktivierungen senden muss. Auch wenn Token innerhalb einer DP-Gruppe von verschiedenen Geräten gesendet werden, werden sie alle zusammen gruppiert, um an den entsprechenden Experten weitergegeben zu werden.

Combine

Die Combine-Kernel sind wiederum in Sende- und Empfängerhälften unterteilt: Die Sender kopieren die unquantisierten 16-Bit-Aktivierungen auf die Remote-Geräte, wobei die Empfangskernel auf das Senden der Daten warten und den gewichteten Durchschnitt aller Expertenbeiträge lokal berechnen. Zusätzlich fungieren sie als Barrier, um die Dispatch-Combine-Sequenz zu synchronisieren: Jeder Rang setzt beim Eintritt in die Sendekernel ein Flag auf jedem Peer, wobei es den Empfangskerneln nicht erlaubt ist zurückzukehren, es sei denn, sie beobachten, dass das Flag gesetzt ist. Die Latenz der Synchronisation ist minimal, da sie sich mit der eigentlichen Kommunikation und Berechnung überlappt.

Vergrößern

Auf der Senderseite durchlaufen die Kernel die Liste der allen lokalen Experten zugewiesenen Token parallel und schreiben sie in einen Puffer auf dem Zielrang. Der Zielrang, der Experten-Index und der Token-Index werden aus den pro-Token-Puffern gelesen, die von den Scatter-Kerneln nach dem Empfang der Token gefüllt wurden. Jeder Sender hat für jeden Experten einen eigenen privaten Speicherbereich, in den geschrieben wird, wie in der Abbildung oben gezeigt, wodurch eine Synchronisation vermieden wird. Ähnlich wie beim Dispatch inkrementiert Combine atomar pro-Token-Zähler auf dem Zielrang, um den Empfang der Daten anzuzeigen: Wenn der Zähler mit der Anzahl der Experten übereinstimmt, an die ein Token verteilt wurde, kann auf den Token-Inhalt zugegriffen werden.

Im Empfangskernel wird die Liste der Token parallel über mehrere Blöcke hinweg durchlaufen, wobei auf deren Ankunft gewartet wird, indem das durch die Signalisierungsoperation gesetzte Flag abgefragt wird. Nach der Ankunft werden die Payloads aus den privaten Puffern gelesen, wobei die Routing-Tabelle anzeigt, aus welchem Puffer gelesen werden soll und welches Gewicht jedem Experten zuzuweisen ist. Die Ergebnisse werden dann in extern zugewiesene Tensoren geschrieben, wobei der Kernel die Ausführung beendet, sobald alle Geräte den Barrier passiert haben. Über eine DP-Gruppe hinweg empfangen alle Ränge eine Kopie jeder Experten-Aktivierung, um ihre eigenen Replikate zu berechnen.

Benchmarks

Wir evaluieren unsere Kernel auf einem Cluster von 128 GPUs, verteilt über 16 Nodes, die über InfiniBand und CX-7-NICs verbunden sind. Wir testen sowohl GPUDirect Async (IBGDA) als auch Reliable Connection (RC) mit einem CPU-Proxy. Wir vergleichen sie mit der hochoptimierten DeepSeek-Implementierung sowie den von PyTorch (via NCCL) oder NVSHMEM bereitgestellten Dense-Primitiven.

Vergrößern

Obwohl die Leistung von IBRC und IBGDA bei der Dense-NVSHMEM-All-to-All-Operation ähnlich ist (6378 µs gegenüber 6180 µs), ist IBGDA bei den Sparse-Kerneln deutlich schneller. Addiert man Dispatch und Combine, benötigt IBGDA 902 µs, während IBRC 3223 µs benötigt – eine 3,6-fache Verbesserung der Latenz. Während All-to-All bandbreitengebunden ist, übertragen die Sparse-Kernel um Größenordnungen weniger Daten und sind latenzgebunden. Durch das Auslösen von Netzwerktransfers direkt von einer GPU, ohne einen CPU-Proxy zur Koordination von GPU und NIC zu benötigen, wird die End-to-End-Latenz signifikant reduziert.

Obwohl unsere portablen Kernel etwa 2-mal langsamer sind als die hochoptimierten DeepSeek-Kernel, verbessern sie die Latenz um das 10-fache im Vergleich zu den Dense-Kerneln (902 µs gegenüber 9944 µs).

Zusätzlich ermöglicht die Aufteilung in Sender-Empfänger-Komponenten es auch, einen Teil der Latenz zu verbergen, anders als bei den Bibliotheksprimitiven.

Vergrößern

Auf einem Single-Node (EP8) nutzt NVSHMEM NVLINK für den Transport und liefert eine geringere Latenz und höheren Durchsatz als die Inter-Node-Vernetzung. Unsere portablen Kernel sind etwa 2,5-mal schneller als DeepEP auf einem Single-Node (186 µs gegenüber 481 µs).

Weitere Verbesserungen

Die hier beschriebenen Kernel übertreffen die integrierten Primitiven von ML-Frameworks und bieten eine ordentliche Leistung, ohne sich zu sehr auf bestimmte Inter-Device-Transporte wie NVLink oder InfiniBand zu spezialisieren. Neben den bereits erwähnten Möglichkeiten sind weitere Leistungssteigerungen durch den Ersatz der Kommunikationsprimitive durch spezialisiertere Versionen erreichbar. Zum Beispiel könnte über NVLink die Verwendung von symmetrischem Speicher durch Puffer ersetzt werden, die zwischen den Geräten geteilt werden, was einige Kopiervorgänge eliminieren und feinere Synchronisationsschemata ermöglicht, über individuelle Token statt über Token-Batches hinweg. Über InfiniBand könnte eine Implementierung direkt auf die zugrunde liegenden Queue-Pairs zugreifen, wodurch die Notwendigkeit entfällt, sowohl die erwarteten als auch die gesendeten Token-Zahlen zu übertragen und Token aus den Queues in Ausgabepuffer zu flushen, sobald sie ankommen. Solche Implementierungen gehen jedoch zu Lasten der Portabilität, während unsere Implementierungen es ermöglichen, Cluster mit spezifischen Konfigurationen zu evaluieren, bevor Entwicklungsaufwand investiert wird, um sie für bestimmte Hardware zu optimieren.

Fazit

Wir haben eine leistungsfähige, portable Bibliothek für MoE-Kommunikation vorgestellt, die eine 10-mal schnellere Leistung gegenüber der Standard-All-to-All-Kommunikation erreicht und gleichzeitig die Kompatibilität über diverse Hardwarekonfigurationen hinweg beibehält. Bei Single-Node-Bereitstellungen mit NVLink zeigt unsere Lösung eine 2,5-mal geringere Latenz als bisherige Implementierungen.

Unser Ansatz gleicht Leistung mit Portabilität durch wichtige Innovationen aus, einschließlich der Unterstützung für GPU-initiierte Kommunikation, einer Split-Kernel-Architektur, die eine Überlappung von Berechnung und Kommunikation ermöglicht, und effizientem Token-Dispatch unter Verwendung minimaler NVSHMEM-Primitive. Während unsere Bibliothek bei Multi-Node-Setups etwa 2-mal langsamer ist als hochspezialisierte Implementierungen, bietet sie eine überlegene Flexibilität über verschiedene Netzwerkumgebungen hinweg (NVLink, CX-7 und EFA).

Da MoE-Modelle weiter skalieren, werden effiziente Kommunikationsstrategien wie unsere für den praktischen Einsatz immer wichtiger. Unsere vollständig Open-Source-Implementierung ist verfügbar unter https://github.com/ppl-ai/pplx-kernels.

Referenzen

DeepEP

Technischer Bericht zu DeepSeek-V3

NVIDIA NVSHMEM

Zitieren

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}