Modèle de coûts LHS

tldr;

Cette page décrit les composants internes du modèle de coût utilisé par le planificateur de masquage de la latence. Si vous souhaitez ajuster le modèle, accédez directement à la section sur l'ajustement.

Le planificateur de masquage de la latence (LHS, Latency Hiding Scheduler) est un pass de compilation qui planifie un DAG HLO de manière à minimiser le temps écoulé.

Ses décisions sont guidées par le modèle de coût unifié, qui utilise un mélange de tableaux de performances et de modèles analytiques. En particulier, XLA intègre des tableaux de performances pour les GEMM et les collectifs d'interconnexion rapide, et utilise un modèle de coût analytique pour la mise en réseau et la fusion dans d'autres cas. Le reste du document décrit le fonctionnement interne de ces éléments de manière générale.

Tableaux de performances – Organismes de gestion collective ICI

Le tableau des performances se compose de deux éléments principaux : un collecteur et un interpolateur.

Collecteur

Le collecteur est un outil C++ chargé de générer les tableaux de performances pour les opérations collectives. Elle mesure les performances des opérations HLO individuelles (par exemple, all-gather, all-reduce) dans un espace de paramètres défini de manière statique.

Fonctionnement

L'outil effectue un balayage sur une plage d'opérations collectives, de tailles de transfert et de schémas de transfert pour un cluster donné. Il utilise l'infrastructure d'exécution HLO multihôte existante et les données ExecutionProfile pour exécuter le HLO généré et collecter les métriques de performances.

Paramètres de collecte de données

Les tableaux de latence sont collectés pour un produit croisé des paramètres suivants :

  • Type de collectif :
    • all-reduce
    • all-gather
    • reduce-scatter
  • Taille du transfert :
    • Échelle logarithmique de 1 024 octets à 2 Gio (par exemple, 1024 octets, 2048 octets, 4096 octets, etc.)
  • Schéma de transfert :
    • rail-aligned
    • non-rail-aligned

Cette analyse est exécutée pour les clusters intranœuds avec 2, 4 et 8 appareils.

Sortie

Le résultat d'une exécution de collection est un tableau de latence au format .pbtxt (environ 116 Ko par plate-forme).

Interpolateur

L'interpolateur est le composant du compilateur qui utilise les tables de performances générées pour fournir des estimations d'exécution lors de la compilation.

Structure interne des données

Lors de l'initialisation, l'interpolateur traite le tableau des performances en une carte. Cette carte utilise un tuple de (collective_type, transfer_scheme) comme clé.

La valeur associée à chaque clé est un plan euclidien 2D. Ce plan indexe le débit réseau (mesuré par le collecteur) selon deux axes :

  1. Taille du transfert.
  2. Nombre d'appareils concernés.

Recherche et interpolation

Lorsque le compilateur rencontre une opération collective, l'interpolateur effectue les étapes suivantes :

  1. Il identifie le plan de débit 2D correct en utilisant le (collective_type, transfer_scheme) de l'opération comme clé de carte.
  2. Il utilise ensuite une récupération de moyenne pondérée (basée sur la distance euclidienne) dans ce plan 2D, en utilisant le (transfer_size, num_devices) de l'opération comme point de requête.
  3. Le résultat de cette recherche est une valeur unique de débit réseau.

Raison : Débit et extrapolation

Le système est conçu pour stocker le débit réseau plutôt que la latence brute. Ce choix de conception simplifie considérablement l'extrapolation des performances pour les tailles de transfert qui ne sont pas explicitement présentes dans le tableau.

Si les tableaux de latence capturent la saturation de la bande passante réseau à une taille collective S, le débit T à ce moment-là est considéré comme le maximum. Pour tout nouveau collectif de taille S' > S, la durée d'exécution peut être estimée comme suit :

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

Cela permet au modèle d'estimer les performances pour les collectifs de n'importe quelle taille, même ceux qui dépassent la limite maximale de 2 Gio mesurée par le collecteur.

  • Sous-estimer le débit maximal.
  • Par conséquent, surestimez la durée d'exécution pour les transferts volumineux.

En général, les équipes XLA:GPU tiennent à jour des tableaux de performances. Toutefois, si les utilisateurs décident de fournir les leurs, il leur incombe de s'assurer qu'ils sont représentatifs et qu'ils incluent des mesures dans la région saturée en bande passante pour le matériel cible.

Tableaux de performances : GEMM

Comme pour les collectifs, les tables de latence GEMM sont compatibles avec deux composants : un collecteur et un interpolateur.

Collecteur

Le collecteur est un outil C++ qui calcule les tableaux de performances pour les multiplications matricielles générales (GEMM). Elle mesure les performances des multiplications de matrices au niveau de l'opération HLO dot.

Fonctionnement

L'outil effectue un balayage sur un espace statique de dimensions GEMM (batch, deux dimensions non contractuelles et une dimension contractuelle) et de types de données.

  • Types de données par défaut : LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32.
  • Infrastructure : réutilise le profileur d'opérations HLO.

Paramètres de la collection

Les tableaux de latence sont collectés pour un produit croisé des dimensions suivantes :

  • batch : {1, 2, 4}
  • m (non contractuel) : {256, 512, ..., 4096}
  • n (sans contraction) : {256, 512, ..., 4096}
  • k (contraction) : {256, 512, ..., 4096}

Sortie et stockage

Un balayage complet génère une table de latence .pbtxt, prête à être utilisée par l'interpolateur.

Interpolateur

L'interpolateur est le composant du compilateur qui utilise les tables générées pour estimer les performances GEMM.

Explication : saturation des FLOPS

Les tableaux de latence collectés permettent à l'interpolateur de reconstruire les FLOPS pour chaque entrée :

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

Une information clé est que les FLOPS saturent à un certain point, c'est-à-dire que le matériel atteint le pic de FLOPS au-delà d'une certaine forme de matrice. Cette saturation permet d'utiliser la même méthode d'extrapolation que celle employée pour les collectifs.

Recherche et interpolation

L'interpolateur crée un espace euclidien 4D à partir des données du tableau. Pour fournir une estimation des performances, il effectue une interpolation de moyenne pondérée dans cet espace à quatre dimensions. S'il n'existe pas de tableau pour un certain type de données, chaque dimension est normalisée en fonction du nombre d'octets, à titre d'heuristique.

Modèle de coût analytique – DCN

Modèle de coût collectif de courbe en S

Le modèle de courbe en S est un modèle de roofline de mise en réseau entièrement analytique.

Présentation

Le modèle est conçu pour estimer les performances des opérations collectives en fonction d'un ensemble de propriétés réseau fixes.

Entrées du modèle

Le modèle nécessite deux catégories d'entrées :

  1. Propriétés réseau fixes (définies par l'utilisateur) :

    • Latence collective de lancement
    • Vitesse de la carte d'interface réseau
    • DAR (délai aller-retour)

    Par défaut, XLA détecte automatiquement une plate-forme et utilise des valeurs pour les architectures les plus courantes. Ces propriétés sont configurables par l'utilisateur. Pour en savoir plus, consultez la section Réglage.

  2. Entrées par collectif :

    • Type collectif (par exemple, AllGather, ReduceScatter)
    • Taille du transfert
    • Nombre de nœuds impliqués dans la communication

Intégration

Le modèle de courbe en S est intégré à XLA:GPU et est utilisé sur Hopper et Blackwell.

Modèle de coût analytique : fusions

Pour les autres kernels, nous nous appuyons sur le modèle de coût des performances du GPU pour estimer les bons temps d'exécution. Pour en savoir plus, cliquez ici.

Réglage

Le modèle de courbe en S peut être ajusté en émettant les bons indicateurs XLA. La configuration par défaut devrait suffire dans la majorité des cas, mais le contrôle du modèle est exposé dans d'autres cas.

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"