Renforcer la compétitivité des GPU AMD pour l’inférence LLM (2023)

• MLC-LLM permet de compiler et déployer des LLM sur des GPU AMD avec ROCm, ce qui permet au Radeon RX 7900 XTX d’atteindre 80 % de la vitesse du RTX 4090 et 94 % de celle du RTX 3090 Ti en inférence sur Llama 2 7B/13B
• Les trois GPU comparés disposent tous de 24 Go de mémoire et d’une bande passante similaire, si bien que, pour l’inférence LLM sensible à la latence, le goulot d’étranglement mémoire pèse davantage que les performances de calcul FP16
• L’écart d’AMD venait davantage du manque de prise en charge logicielle et d’optimisation que du matériel lui-même, et les investissements dans ROCm ainsi que la compilation machine learning contribuent à réduire le coût de prise en charge de plusieurs backends
• ROCm est appliqué au Radeon 7900 XTX, et Vulkan à l’APU AMD du SteamDeck, montrant qu’il est possible d’exécuter Llama-7B quantifié en 4 bits sur des appareils AMD grand public
• Les résultats reposent sur des benchmarks en batch unique, et les optimisations d’attention peuvent modifier les chiffres côté AMD comme NVIDIA ; ils doivent donc être interprétés avec une marge d’erreur de 10 %

---

Inférence LLM sur GPU AMD avec ROCm

• MLC-LLM permet de compiler et déployer des LLM sur des GPU AMD en utilisant ROCm
• Les performances du Radeon RX 7900 XTX sur Llama 2 7B/13B atteignent :
  • 80 % de la vitesse du NVIDIA GeForce RTX 4090
  • 94 % de la vitesse du NVIDIA GeForce RTX 3090 Ti
• En plus de ROCm, MLC prend aussi en charge Vulkan, ce qui élargit le champ de déploiement des LLM à des appareils comme le SteamDeck équipé d’un APU AMD

Comparaison matérielle et véritables goulots d’étranglement

• L’AMD RX 7900 XTX occupe, sur le papier, une position comparable aux RTX 4090 et RTX 3090 Ti
  • Les trois GPU disposent tous de 24 Go de mémoire, ce qui permet de charger des modèles de même taille
  • Leur bande passante mémoire est également similaire
  • Le RTX 4090 offre des performances FP16 2 fois supérieures à celles du RX 7900 XTX, et le RTX 3090 Ti 1,3 fois supérieures
• L’inférence LLM sensible à la latence est majoritairement limitée par la mémoire, si bien qu’ici l’écart de performances FP16 n’est pas le goulot d’étranglement principal
• Le RX 7900 XTX est 40 % moins cher que le RTX 4090
• Le RTX 3090 Ti appartient à la génération précédente, ce qui rend la comparaison de prix plus délicate ; il sert surtout de point de référence

Origine de l’écart d’AMD et approche de MLC

• Le retard d’AMD tenait davantage au manque de prise en charge logicielle et d’optimisations pour ces modèles qu’au matériel
• Deux évolutions contribuent à réduire cet écart :
  • AMD investit dans la stack ROCm pour rattraper son retard
  • La compilation machine learning réduit le coût du support logiciel générique pour plusieurs backends
• La compilation machine learning automatise la compilation et l’optimisation des charges de travail ML, au lieu d’écrire manuellement des kernels distincts pour ROCm, CUDA, etc.
• MLC-LLM est une solution de déploiement de LLM fondée sur la compilation machine learning, construite sur Apache TVM Unity
  • Elle propose un flux de développement centré sur Python
  • Elle inclut la transformation des graphes de calcul, l’optimisation des layouts et des schedules de kernels GPU, ainsi qu’une API de déploiement native
  • Elle prend en charge CUDA, Metal, ROCm, Vulkan et OpenCL
  • Elle couvre aussi bien les GPU serveur que le mobile comme l’iPhone ou Android

Méthodes de prise en charge des GPU et APU AMD

• Les voies de prise en charge des GPU AMD incluent ROCm, OpenCL, Vulkan et WebGPU
  • ROCm est la stack qu’AMD pousse activement ces derniers temps, avec des composants comparables à ceux de CUDA
  • Vulkan est un standard graphique moderne largement pris en charge sur les appareils GPU
  • WebGPU est un standard web récent permettant d’exécuter des calculs dans le navigateur
• Peu de logiciels de machine learning construisent des solutions hors CUDA, et le coût d’ingénierie pour répliquer une stack à chaque nouveau matériel ou modèle de programmation GPU est élevé
• MLC prend en charge plusieurs voies grâce à la génération automatique de code, sans réécrire les kernels GPU pour chaque backend
• Les performances réelles dépendent de la qualité du runtime GPU bas niveau et de sa disponibilité sur chaque plateforme

Mise en œuvre des optimisations ROCm

• ROCm a été retenu pour le Radeon 7900 XTX, et Vulkan pour l’APU du SteamDeck
• La stack ROCm a fonctionné immédiatement, et grâce au pipeline de développement Python de TVM Unity, une version optimisée a pu être ajoutée en quelques heures
• La prise en charge ROCm réutilise des composants existants :
  • L’ensemble du pipeline MLC pour des cibles existantes comme CUDA et Metal
    • planification mémoire
    • fusion d’opérateurs
  • L’espace générique d’optimisation des kernels GPU écrit en TVM TensorIR
  • Le flux de génération de code ROCm de TVM, qui produit des kernels ROCm bas niveau via LLVM
• Le code généré peut être exporté sous forme de bibliothèque partagée ou statique, puis appelé depuis une CLI, Python ou une API REST

Conditions et interprétation des benchmarks Llama 2

• Les benchmarks mesurent Llama 2 7B et 13B en quantification 4 bits
• Les performances de décodage sont mesurées en fournissant un seul token de prompt puis en générant 512 tokens
• Tous les résultats sont basés sur une inférence en batch unique
• Avec la version ROCm 5.6, les performances en inférence batch unique atteignent 80 % de la vitesse du NVIDIA 4090
• La baseline CUDA est considérée comme l’une des meilleures disponibles à l’époque pour cette tâche
• Il reste une marge d’amélioration, notamment avec de meilleures optimisations d’attention, et si elles sont intégrées à MLC, elles pourraient améliorer les résultats AMD comme NVIDIA
• Si ces optimisations n’étaient implémentées que côté NVIDIA, l’écart pourrait passer de 20 % à 30 %, d’où la recommandation d’interpréter ces chiffres avec une marge d’erreur de 10 %

Conditions d’exécution directe et exemple

• Des wheels précompilées et des instructions d’exécution sont fournies pour reproduire les benchmarks
• Les conditions d’exécution sont un GPU AMD sous Linux avec ROCm 5.6 ou ultérieur
• L’installation du paquet MLC précompilé avec ROCm activé suit la documentation d’essai de MLC-LLM
• L’exemple Python charge le modèle Llama-2-7b-chat-hf-q4f16_1 avec mlc_chat.ChatModule et mesure les performances via benchmark_generate("Hi", generate_length=512)
• MLC-LLM propose aussi une CLI interactive, mais sur ROCm la CLI doit être compilée depuis les sources en suivant la documentation de build de la CLI

Vulkan et mémoire unifiée sur SteamDeck

• Le SteamDeck est utilisé comme exemple d’une gamme plus large d’appareils AMD équipés d’un APU
• En ROCm, la VRAM GPU disponible est limitée à 4 Go dans le BIOS
• Le pilote Mesa Vulkan permet, via la mémoire unifiée, d’utiliser des buffers au-delà de cette limite jusqu’à 16 Go
• Cette capacité mémoire suffit pour exécuter Llama-7B quantifié en 4 bits
• Cela montre qu’une prise en charge des LLM est possible sur divers appareils AMD grand public

Travaux futurs et liens du projet

• À l’ère de l’IA générative, la disponibilité du matériel devient un enjeu important
• La compilation machine learning peut atténuer ce problème en permettant un déploiement générique haute performance sur plusieurs backends matériels
• Les recherches actuelles se concentrent sur les GPU grand public
• D’après l’expérience acquise, les optimisations MLC pour GPU grand public se généralisent parfois aussi aux GPU cloud, par exemple du RTX 4090 vers les A100 et A10g
• Les prochains axes de travail incluent :
  • le batching et la prise en charge multi-GPU
  • l’intégration à l’écosystème PyTorch
  • la prise en charge de davantage de schémas de quantification et d’architectures de modèles
  • l’extension de l’optimisation automatique à davantage de backends matériels
• NVIDIA conserve une position de leader grâce à une innovation continue, et le paysage peut évoluer avec l’arrivée de nouveaux matériels comme le H100 et les avancées logicielles
• Les informations sur MLC-LLM sont disponibles sur la page du projet, et le code source sur le dépôt GitHub