Résumé Technique : Réduire le goulot d'étranglement de la mémoire GPU par la compression sans perte pour le ML

Énoncé du problème

Les charges de travail modernes de l'apprentissage automatique (ML), incluant les réseaux de neurones sur graphes (GNN), les modèles de recommandation de deep learning (DLRM) et les grands modèles de langage (LLM), traitent fréquemment des ensembles de données allant de centaines de gigaoctets à des téraoctets. Ces volumes de données dépassent la capacité de la mémoire GPU disponible, nécessitant des transferts de données à la demande depuis la mémoire CPU, le disque ou le stockage réseau via l'interconnexion PCIe.

La principale contrainte de performance dans ces scénarios est la bande passante PCIe, qui est plus d'un ordre de grandeur inférieure à la bande passante de la mémoire GPU. Par conséquent, l'exécution est souvent limitée par la latence de transfert de données plutôt que par la capacité de calcul du GPU. Bien que des techniques telles que la mise en cache et le chevauchement du calcul avec les transferts offrent un certain soulagement, elles ne peuvent pas atténuer totalement le goulot d'étranglement, en particulier lorsque des échecs de cache (cache misses) se produisent.

Les solutions existantes reposent souvent sur la compression avec perte (par exemple, la quantification) pour réduire le volume de données. Cependant, les méthodes avec perte introduisent une dégradation imprévisible de la précision qui varie selon le modèle et la charge de travail. Dans les déploiements commerciaux, même des pertes de précision mineures (ex. : 0,1 %) sont souvent inacceptables en raison des risques sur les revenus, ce qui rend la compression avec perte difficile à adopter comme solution à usage général.

La compression sans perte offre une alternative compatible avec le déploiement car elle préserve la précision du modèle, mais elle fait face à des défis de mise en œuvre importants dans les pipelines de ML :

1. Compétition de latence : Les latences de transfert CPU-GPU sont compétitives avec les latences de décompression GPU ; les algorithmes de compression traditionnels ajoutent souvent un surcoût qui annule les économies de transfert.
2. Contraintes de mémoire : La décompression nécessite souvent de grands tampons intermédiaires ou des structures de métadonnées complexes qui entrent en compétition avec les entrées du modèle pour la mémoire GPU limitée.
3. Inadéquation du parallélisme : De nombreux algorithmes sans perte standards (ex. : codage de Huffman, LZ4) reposent sur des recherches dans des dictionnaires ou des modèles d'accès mémoire irréguliers qui entravent le parallélisme des warps du GPU et les accès mémoire coalescés.

Méthodologie : Invariant Bit Packing (IBP)

Les auteurs proposent l'Invariant Bit Packing (IBP), un nouvel algorithme de compression sans perte conçu spécifiquement pour les transferts de tenseurs de ML. L'idée centrale est que les tenseurs de ML, même lorsqu'ils sont numériquement divers, présentent souvent des bits invariants à faible entropie à travers des groupes de tenseurs. Par exemple, dans les représentations en virgule flottante, les bits de signe et d'exposant restent souvent constants sur de nombreux tenseurs en raison de la régularisation et de la dynamique d'entraînement.

Algorithme central

L'IBP opère en trois étapes principales :

1. Prétraitement (Génération de métadonnées) :

   • Le système analyse un lot de tenseurs d'entrée pour identifier les positions de bits qui sont invariantes (c'est-à-dire qui ont la même valeur) à travers un pourcentage élevé ($T$) de tenseurs.
   • Il génère un Masque de taille fixe (indiquant quelles positions de bits sont invariantes) et un Bitval (stockant les valeurs invariantes).
   • Ces métadonnées sont stockées en mémoire GPU (généralement de l'ordre de quelques kilo-octets) et sont réutilisées pour l'ensemble du jeu de données, évitant ainsi le surcoût des métadonnées par tenseur.

2. Compression :

   • Les tenseurs d'entrée sont traités par blocs. Pour chaque bloc, l'algorithme vérifie si les bits aux positions invariantes correspondent au Bitval stocké.
   • S'ils correspondent, les bits invariants sont supprimés (éliminés), et les bits restants sont décalés pour combler le vide.
   • Un bit de participation est ajouté devant chaque bloc pour indiquer si celui-ci a été compressé ou laissé non compressé (pour gérer les cas où la compression ne produit pas d'économies d'espace).
   • La compression est effectuée sur le GPU en utilisant la mémoire virtuelle unifiée zero-copy, permettant au GPU d'accéder directement à la mémoire CPU sans configuration explicite de DMA initiée par le CPU pour les petits transferts.

3. Décompression :

   • La décompression est optimisée pour le parallélisme des warps du GPU. Un warp de 32 threads lit de manière coopérative un segment contigu de données compressées depuis la mémoire CPU vers un espace de travail en mémoire partagée.
   • Les threads utilisent des primitives de warp (communication à faible latence au sein d'un warp) pour effectuer un balayage préfixe (prefix scan), calculant l'offset de bit de départ pour le segment assigné à chaque thread.
   • Les threads reconstruisent ensuite leurs données de manière indépendante en réinsérant les bits invariants du Bitval global dans les positions correctes.
   • Le processus utilise des accès mémoire asynchrones pour superposer la récupération des données avec la logique de décompression, masquant ainsi la latence PCIe.

Intégration au système

Les auteurs intègrent l'IBP dans les pipelines de ML de deux manières principales :

• Compression de cache : Compresser les tenseurs stockés dans les caches logiciels du GPU (ex. : dans les frameworks GNN) pour augmenter la capacité effective du cache.
• Compression de transfert PCIe : Compresser les tenseurs en mémoire CPU avant le transfert. Pour gérer les tenseurs compressés de taille variable, le système utilise une compression sur place (allouant des tampons de taille fixe mais ne lisant que la partie compressée) et maintient un masque de bits en mémoire GPU pour suivre quels tenseurs sont compressés.

Contributions clés

1. Analyse de l'invariance des données de ML : L'article identifie et quantifie la présence de bits invariants à travers divers flux de travail et types de données de ML (GNN, DLRM, LLM) et types de données (FP32, FP16, BF16), démontant que les distributions structurées permettent une compression sans perte significative sans perte de précision.
2. Algorithme Invariant Bit Packing (IBP) : Une technique de compression sans perte optimisée pour le GPU qui supprime les bits invariants, stocke les métadonnées globales de manière minimale et utilise la décompression parallèle par warp pour obtenir un débit élevé.
3. Intégration au système : Implémentation de l'IBP dans des frameworks de ML populaires (DGL, Legion, Colossal-AI, FlexGen, InfiniGen) via des API faciles à utiliser (extensions PyTorch et bibliothèques CUDA).
4. Évaluation complète : Benchmarking extensif par rapport aux bibliothèques accélérées par GPU de pointe (nvCOMP, ndzip-gpu) et aux systèmes natifs à travers des jeux de données réels.

Résultats

L'évaluation a été réalisée sur un GPU NVIDIA A100 avec PCIe Gen 4.

• Performance de compression : L'IBP atteint des économies d'espace de 10–12 % pour les jeux de données GNN denses, 8 % pour les poids DLRM, et jusqu'à 27 % pour les poids et les KV-caches de LLM (particulièrement avec les types de données BF16).
• Débit : L'IBP surpasse systématiquement les bibliothèques de compression accélérées par GPU existantes. Pour les jeux de données GNN épars, il atteint un gain de vitesse de 9,7x du débit de transfert par rapport aux transferts non compressés.
• Accélération des charges de travail :
  • Entraînement GNN : Accélération moyenne de 74 % (jusqu'à 2,7x dans certains cas spécifiques) en réduisant le temps de transfert PCIe et en augmentant les taux de succès du cache.
  • Inférence DLRM : Accélération moyenne de 180 % pour les recherches d'embeddings, principalement portée par des transferts initiés par le GPU efficaces et une réduction de la surcharge CPU.
  • Inférence LLM : Accélération moyenne de 24 % pour l'inférence déchargée (offloaded). Pour le déchargement des poids (BF16), la latence a été réduite de 22 % ; pour le déchargage du KV-cache, la latence a été réduite jusqu'à 27 %.
• Surcoût (Overhead) : Le surcoût de décompression est minimal ; pour des économies d'espace inférieures à 50 % (typiques pour les données de ML denses), le débit reste supérieur à 95 % du taux de transfert idéal non compressé. Le temps de prétraitement est modéré et se situe hors du chemin critique.

Signification et affirmations

L'article affirme que l'IBP fournit une solution pratique et sans perte au goulot d'étranglement de la mémoire GPU, sans les risques de précision associés à la quantification ou à la compression avec perte. En exploitant la structure inhérente des données de ML (bits invariants) et en optimisant pour l'architecture GPU (parallélisme de warp, accès coalescés, zero-copy), l'IBP parvient efficacement à masquer le surcoût de décompression derrière la latence de transfert PCIe.

Les auteurs soulignent que l'IBP est un algorithme à usage général applicable à divers types de données (float32, float16, bfloat16, vecteurs épars) et charges de travail (entraînement et inférence). La capacité d'intégrer l'IBP avec un minimum de changements de code (souvent moins de 35 lignes de Python) dans les frameworks existants démontre son potentiel d'adoption dans les déploiements de ML commerciaux pour mettre à l'échelle les modèles au-delà des limites de la mémoire GPU tout en maintenant la performance et la précision.