LHS-Kostenmodell

Kurzfassung:

Auf dieser Seite wird das Kostenmodell beschrieben, das vom Latency Hiding Scheduler verwendet wird. Wenn Sie das Modell abstimmen möchten, können Sie direkt zum Abschnitt zum Abstimmen gehen.

Der Latency Hiding Scheduler (LHS) ist ein Compiler-Pass, der einen HLO-DAG so plant, dass die Echtzeit minimiert wird.

Die Entscheidungen basieren auf dem einheitlichen Kostenmodell, das eine Mischung aus Leistungstabellen und Analysemodellen verwendet. Insbesondere enthält XLA Leistungstabellen für GEMMs und schnelle Interconnect-Collectives und verwendet für andere Fälle ein analytisches Netzwerk- und Fusionskostenmodell. Im Rest des Dokuments wird die Funktionsweise dieser Komponenten auf allgemeiner Ebene beschrieben.

Leistungstabellen – ICI-Kollektive

Die Leistungstabelle besteht aus zwei Hauptkomponenten: einem Collector und einem Interpolator.

Collector

Der Collector ist ein C++-Tool, das für die Generierung der Leistungstabellen für kollektive Vorgänge verantwortlich ist. Sie misst die Leistung einzelner HLO-Operationen (z.B. all-gather, all-reduce) über einen statisch definierten Parameterbereich hinweg.

Funktionsweise

Das Tool führt für einen bestimmten Cluster einen Sweep über eine Reihe von kollektiven Operationen, Übertragungsgrößen und Übertragungsschemas durch. Dabei wird die vorhandene Infrastruktur für die Ausführung von HLOs auf mehreren Hosts und ExecutionProfile-Daten verwendet, um das generierte HLO auszuführen und Leistungsmesswerte zu erfassen.

Parameter für die Datenerhebung

Latenztabellen werden für eine Kombination der folgenden Parameter erfasst:

  • Collective Type (Kollektivtyp):
    • all-reduce
    • all-gather
    • reduce-scatter
  • Übertragungsgröße:
    • Logarithmische Skala von 1.024 B bis 2 GiB (z. B. 1024 B, 2048 B, 4096 B usw.)
  • Übertragungsschema:
    • rail-aligned
    • non-rail-aligned

Dieser Sweep wird für knoteninterne Cluster mit 2, 4 und 8 Geräten ausgeführt.

Ausgabe

Das Ergebnis eines Sammlungslaufs ist eine Latenztabelle im .pbtxt-Format (ca. 116 KB pro Plattform).

Interpolator

Der Interpolator ist die Compilerkomponente, die die generierten Leistungstabellen verwendet, um während der Kompilierung Laufzeitschätzungen zu liefern.

Interne Datenstruktur

Bei der Initialisierung wird die Leistungstabelle vom Interpolator in eine Karte umgewandelt. Diese Map verwendet ein Tupel von (collective_type, transfer_scheme) als Schlüssel.

Der Wert, der jedem Schlüssel zugeordnet ist, ist eine zweidimensionale euklidische Ebene. Auf dieser Ebene wird der Netzwerkdurchsatz (gemessen vom Collector) anhand von zwei Achsen indexiert:

  1. Größe der Übertragung.
  2. Anzahl der beteiligten Geräte.

Nachschlagen und Interpolieren

Wenn der Compiler auf einen kollektiven Vorgang trifft, führt der Interpolator die folgenden Schritte aus:

  1. Sie identifiziert die richtige 2D-Durchsatzebene anhand des (collective_type, transfer_scheme) des Vorgangs als Zuordnungsschlüssel.
  2. Anschließend wird in dieser 2D-Ebene ein gewichteter Durchschnittsabruf (basierend auf der euklidischen Distanz) durchgeführt, wobei der (transfer_size, num_devices) des Vorgangs als Anfragepunkt verwendet wird.
  3. Das Ergebnis dieser Suche ist ein einzelner, eindeutiger Netzwerkdurchsatz-Wert.

Begründung: Durchsatz und Extrapolation

Das System ist so konzipiert, dass es den Netzwerkdurchsatz und nicht die Rohlatenz speichert. Diese Designentscheidung vereinfacht die Extrapolation der Leistung für Übertragungsgrößen, die nicht explizit in der Tabelle enthalten sind, erheblich.

Wenn in den Latenztabellen die Sättigung der Netzwerkbandbreite bei einer kollektiven Größe von S erfasst wird, gilt der Durchsatz T an diesem Punkt als maximal. Für jedes neue Kollektiv der Größe S' > S kann die Laufzeit so geschätzt werden:

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

So kann das Modell die Leistung für Kollektive jeder Größe schätzen, auch für solche, die größer sind als das vom Collector gemessene Maximum von 2 GiB.

  • Unterschätzen Sie den maximalen Durchsatz.
  • Überschätzen Sie daher die Laufzeit für große Übertragungen.

Im Allgemeinen werden Leistungstabellen von den XLA:GPU-Teams verwaltet. Wenn Nutzer jedoch eigene Tabellen bereitstellen, liegt es in ihrer Verantwortung, dafür zu sorgen, dass diese repräsentativ sind und Messungen im bandbreitengesättigten Bereich für die Zielhardware enthalten.

Leistungstabellen – GEMMs

Ähnlich wie beim System für Kollektive werden GEMM-Latenztabellen von zwei Komponenten unterstützt: einem Collector und einem Interpolator.

Collector

Der Collector ist ein C++-Tool, mit dem Leistungstabellen für allgemeine Matrixmultiplikationen (GEMMs) berechnet werden. Damit wird die Leistung von Matrixmultiplikationen auf der HLO-Ebene dot gemessen.

Funktionsweise

Das Tool führt einen Sweep über einen statischen Raum von GEMM-Dimensionen (Batch, zwei nicht kontrahierende und eine kontrahierende Dimension) und Datentypen durch.

  • Standarddatentypen:LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32.
  • Infrastruktur:Der HLO-Op-Profiler wird wiederverwendet.

Parameter für die Erhebung

Latenztabellen werden für eine Kombination der folgenden Dimensionen erfasst:

  • batch:{1, 2, 4}
  • m (nicht kontrahierend): {256, 512, ..., 4096}
  • n (nicht vertraglich): {256, 512, ..., 4096}
  • k (Kontrahieren): {256, 512, ..., 4096}

Ausgabe und Speicherung

Bei einem vollständigen Sweep wird eine .pbtxt-Latenztabelle generiert, die vom Interpolator verwendet werden kann.

Interpolator

Der Interpolator ist die Compilerkomponente, mit der die GEMM-Leistung anhand der generierten Tabellen geschätzt wird.

Begründung: FLOPS-Sättigung

Anhand der erfassten Latenztabellen kann der Interpolator die FLOPS für jeden Eintrag rekonstruieren:

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

Eine wichtige Erkenntnis ist, dass FLOPS ab einem bestimmten Punkt gesättigt sind. Das bedeutet, dass die Hardware ab einer bestimmten Matrixform die maximalen FLOPS erreicht. Durch diese Sättigung kann dieselbe Extrapolationsmethode wie für Kollektive verwendet werden.

Nachschlagen und Interpolieren

Der Interpolator erstellt aus den Tabellendaten einen 4D-euklidischen Raum. Um eine Leistungsschätzung zu ermöglichen, wird in diesem 4D-Raum eine gewichtete Durchschnittsinterpolation durchgeführt. Wenn für einen bestimmten Datentyp keine Tabelle vorhanden ist, wird jede Dimension heuristisch auf die Anzahl der Byte normalisiert.

Analytisches Kostenmodell – DCN

S-Kurven-Modell für kollektive Kosten

Das S-Kurvenmodell ist ein vollständig analytisches Netzwerk-Roofline-Modell.

Übersicht

Das Modell wurde entwickelt, um die Leistung kollektiver Vorgänge auf Grundlage einer Reihe von festen Netzwerkeigenschaften zu schätzen.

Modelleingaben

Für das Modell sind zwei Kategorien von Eingaben erforderlich:

  1. Eigenschaften des Festnetzes (nutzerdefiniert):

    • Kollektiver Start-Overhead
    • NIC-Geschwindigkeit
    • RTT (Umlaufzeit)

    Standardmäßig erkennt XLA automatisch eine Plattform und verwendet Werte für die gängigsten Architekturen. Diese Eigenschaften können vom Nutzer konfiguriert werden. Weitere Informationen finden Sie im Abschnitt zum Optimieren.

  2. Eingaben pro Kollektiv:

    • Kollektiver Typ (z.B. AllGather, ReduceScatter)
    • Übertragungsgröße
    • Anzahl der an der Kommunikation beteiligten Knoten

Integration

Das S-Kurven-Modell ist in XLA:GPU integriert und wird auf Hopper und Blackwell verwendet.

Analytisches Kostenmodell – Zusammenführungen

Bei anderen Kernels verwenden wir das GPU-Leistungskostenmodell, um die richtigen Laufzeiten zu schätzen. Weitere Informationen

Abstimmung

Das S-Kurven-Modell kann durch die richtigen XLA-Flags optimiert werden. Die Standardkonfiguration sollte in den meisten Fällen ausreichend sein. In anderen Fällen ist die Modellsteuerung jedoch verfügbar.

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"