Modello di costo LHS

tldr;

Questa pagina descrive i dettagli interni del modello di costi utilizzato da Latency Hiding Scheduler. Se ti interessa ottimizzare il modello, vai direttamente alla sezione Ottimizzazione.

Lo scheduler per l'occultamento della latenza (LHS) è un passaggio del compilatore che pianifica un DAG HLO in modo da ridurre al minimo il tempo totale di esecuzione.

Le sue decisioni sono guidate dal modello di costi unificato, che utilizza una combinazione di tabelle di rendimento e modelli analitici. In particolare, XLA incorpora tabelle delle prestazioni per GEMM e collettivi di interconnessione rapida e utilizza un modello di costi di fusione e networking analitico per altri casi. Il resto del documento descrive il funzionamento interno di questi a un livello generale.

Tabelle del rendimento - Società di gestione collettiva ICI

La tabella del rendimento è composta da due componenti principali: un raccoglitore e un interpolatore.

Raccoglitore

Il raccoglitore è uno strumento C++ responsabile della generazione delle tabelle delle prestazioni per le operazioni collettive. Misura le prestazioni delle singole operazioni HLO (ad es. all-gather, all-reduce) in uno spazio di parametri definito staticamente.

Come funziona

Lo strumento esegue una scansione su una serie di operazioni collettive, dimensioni di trasferimento e schemi di trasferimento per un determinato cluster. Utilizza l'infrastruttura esistente del runner HLO multihost e i dati ExecutionProfile per eseguire l'HLO generato e raccogliere le metriche di rendimento.

Parametri di raccolta dati

Le tabelle di latenza vengono raccolte per un prodotto incrociato dei seguenti parametri:

  • Tipo di collettivo:
    • all-reduce
    • all-gather
    • reduce-scatter
  • Dimensioni trasferimento:
    • Scala logaritmica da 1024 B fino a 2 GiB (ad es. 1024B, 2048B, 4096B, ...)
  • Transfer Scheme:
    • rail-aligned
    • non-rail-aligned

Questo sweep viene eseguito per i cluster intra-nodo con 2, 4 e 8 dispositivi.

Output

Il risultato di un'esecuzione della raccolta è una tabella di latenza in formato .pbtxt (circa 116 KB per piattaforma).

Interpolator

L'interpolatore è il componente del compilatore che utilizza le tabelle delle prestazioni generate per fornire stime di runtime durante la compilazione.

Struttura dei dati interni

Durante l'inizializzazione, l'interpolatore elabora la tabella del rendimento in una mappa. Questa mappa utilizza una tupla di (collective_type, transfer_scheme) come chiave.

Il valore associato a ogni chiave è un piano euclideo bidimensionale. Questo piano indicizza il throughput di rete (misurato dal Collector) in base a due assi:

  1. Dimensioni trasferimento.
  2. Numero di dispositivi coinvolti.

Ricerca e interpolazione

Quando il compilatore rileva un'operazione collettiva, l'interpolatore esegue i seguenti passaggi:

  1. Identifica il piano di velocità effettiva 2D corretto utilizzando (collective_type, transfer_scheme) dell'operazione come chiave della mappa.
  2. Quindi, utilizza un recupero della media ponderata (in base alla distanza euclidea) all'interno di questo piano 2D, utilizzando (transfer_size, num_devices) dell'operazione come punto di query.
  3. Il risultato di questa ricerca è un singolo valore univoco di velocità effettiva di rete.

Motivazione: throughput ed estrapolazione

Il sistema è progettato per archiviare il throughput di rete anziché la latenza non elaborata. Questa scelta di progettazione semplifica notevolmente l'estrapolazione delle prestazioni per dimensioni di trasferimento non esplicitamente presenti nella tabella.

Se le tabelle di latenza acquisiscono la saturazione della larghezza di banda di rete a una dimensione collettiva S, la velocità effettiva T a quel punto viene considerata il massimo. Per qualsiasi nuovo collettivo di dimensioni S' > S, la durata può essere stimata come:

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

Ciò consente al modello di stimare le prestazioni per i collettivi di qualsiasi dimensione, anche quelli più grandi del massimo di 2 GiB misurato dal raccoglitore.

  • Sottostima il throughput massimo.
  • Di conseguenza, sovrastima il tempo di esecuzione per i trasferimenti di grandi dimensioni.

In generale, i team XLA:GPU gestiscono le tabelle delle prestazioni, ma nei casi in cui l'utente decide di fornire le proprie, è responsabilità dell'utente che genera le tabelle assicurarsi che siano rappresentative e includano misurazioni nella regione satura di larghezza di banda per l'hardware di destinazione.

Tabelle del rendimento - GEM

Analogamente al sistema per i collettivi, le tabelle di latenza GEMM sono supportate da due componenti: un raccoglitore e un interpolatore.

Raccoglitore

Il collector è uno strumento C++ che calcola le tabelle delle prestazioni per le moltiplicazioni di matrici generali (GEMM). Misura il rendimento delle moltiplicazioni di matrici a livello di operazione HLO dot.

Come funziona

Lo strumento esegue una scansione su uno spazio statico di dimensioni GEMM (batch, due dimensioni non contrattuali e una contrattuale) e tipi di dati.

  • Tipi di dati predefiniti: LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32.
  • Infrastruttura:riutilizza il profiler delle operazioni HLO.

Parametri di raccolta

Le tabelle di latenza vengono raccolte per un prodotto incrociato delle seguenti dimensioni:

  • batch: {1, 2, 4}
  • m (non contrattuali): {256, 512, ..., 4096}
  • n (non contrattuale): {256, 512, ..., 4096}
  • k (contrazione): {256, 512, ..., 4096}

Output e spazio di archiviazione

Una scansione completa genera una tabella di latenza .pbtxt, pronta per essere utilizzata dall'interpolatore.

Interpolator

L'interpolatore è il componente del compilatore che utilizza le tabelle generate per stimare le prestazioni di GEMM.

Motivazione: saturazione FLOPS

Le tabelle di latenza raccolte consentono all'interpolatore di ricostruire i FLOPS per ogni voce:

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

Un'informazione chiave è che i FLOPS saturano a un certo punto, ovvero l'hardware raggiunge il picco di FLOPS oltre una certa forma della matrice. Questa saturazione consente l'utilizzo dello stesso metodo di estrapolazione impiegato per i collettivi.

Ricerca e interpolazione

L'interpolatore crea uno spazio euclideo 4D dai dati della tabella. Per fornire una stima del rendimento, esegue un'interpolazione della media ponderata all'interno di questo spazio 4D. Se non è presente una tabella per un determinato tipo di dati, ogni dimensione viene normalizzata in base al numero di byte.

Analytical Cost Model - DCN

Modello di costo collettivo della curva a S

Il modello a S è un modello di networking completamente analitico.

Panoramica

Il modello è progettato per stimare il rendimento delle operazioni collettive in base a un insieme di proprietà di rete fisse.

Input per il modello

Il modello richiede due categorie di input:

  1. Proprietà di rete fisse (definite dall'utente):

    • Overhead di lancio collettivo
    • Velocità NIC
    • RTT (tempo di round trip)

    Per impostazione predefinita, XLA rileva automaticamente una piattaforma e utilizza i valori per le architetture più comuni. Queste proprietà sono configurabili dall'utente. Per ulteriori dettagli, consulta la sezione Ottimizzazione.

  2. Input per collettivo:

    • Tipo di collettivo (ad es. AllGather, ReduceScatter)
    • Dimensioni trasferimento
    • Numero di nodi coinvolti nella comunicazione

Integrazione

Il modello a S è integrato in XLA:GPU e viene utilizzato su Hopper e Blackwell.

Modello di costi analitico - Fusioni

Per gli altri kernel, ci affidiamo al modello di costo delle prestazioni della GPU per stimare i runtime corretti. Per saperne di più, fai clic qui.

Ottimizzazione

Il modello a curva a S può essere ottimizzato emettendo i flag XLA corretti. La configurazione predefinita dovrebbe essere sufficiente nella maggior parte dei casi, ma il controllo del modello è esposto in altri casi.

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"