Compatibilité des modèles PyTorch avec le nouveau matériel

Un modèle PyTorch conçu à l'origine pour les GPU Nvidia peut désormais fonctionner sur les puces d'Amazon ou de Google avec des modifications minimales du code ou de la configuration, selon les deux entreprises. Cette compatibilité est rendue possible grâce à de nouveaux outils comme TorchNeuron et TorchTPU, qui permettent aux modèles de s'exécuter sur différentes plateformes matérielles moyennant peu de réécriture.

Le portage du modèle Nova-3 par Deepgram

Deepgram, premier client à détailler publiquement ce processus, a réussi à faire fonctionner son modèle de transcription Nova-3 sur les puces Trainium2 d'Amazon. Cela a nécessité un travail de refonte important, notamment la réécriture des cœurs de calcul avec des ingénieurs d'AWS, le développement d'une nouvelle couche de convolution à partir de zéro, et la création d'un moteur de service en Rust.

Réécriture des noyaux personnalisés pour le nouveau matériel

Amazon comme Google exigent que les noyaux personnalisés écrits en CUDA ou en Triton soient réécrits dans le langage propre à leurs puces respectives. Ces noyaux sont essentiels pour optimiser les performances sur chaque plateforme, et aucune des deux entreprises n'a encore fourni de documentation ni d'outils permettant de traduire automatiquement ces noyaux.

TorchNeuron et ses prérequis

TorchNeuron, actuellement en version bêta publique, nécessite PyTorch 2.14.0 et doit prendre en charge la version 2.15 d'ici le 28 octobre. Il remplace d'anciens ponts comme PyTorch/XLA et utilise l'interface PrivateUse1, réservée dans PyTorch aux matériels tiers.

TorchTPU et son intégration avec Transformers

TorchTPU de Google, codéveloppé avec l'équipe PyTorch de Meta, remplace PyTorch/XLA pour la prise en charge des TPU. Il utilise le même pipeline de compilation PyTorch/XLA, incluant OpenXLA et StableHLO, mais contourne Inductor. Google recommande de migrer les charges de travail existantes vers TorchTPU.

Prise en charge des modèles PyTorch sur différentes plateformes

Hugging Face indique que 89 % des 100 modèles de génération de texte les plus utilisés sont nativement pris en charge par Transformers, contre 46 % avec l'ancienne bibliothèque Optimum Neuron. Cela suggère que les outils plus récents comme TorchNeuron et TorchTPU offrent une meilleure compatibilité avec les modèles actuels.

Limites et développement en cours

TorchNeuron et TorchTPU présentent tous deux des limites connues. TorchNeuron ne prend en charge que certaines versions de Python ainsi que les instances Trn2 et Trn3 d'Amazon, tandis que TorchTPU prend en charge plusieurs générations de TPU mais ne précise pas clairement ses exigences de version. Certaines opérations nécessitent encore une recompilation, et certaines fonctionnalités restent à un stade expérimental.

Affirmations de performance et tests de référence

Amazon affirme que son système peut faire fonctionner 100 000 modèles Hugging Face avec une grande précision, sans toutefois publier de taux de réussite. Google évoque des améliorations de performance pour l'entraînement de modèles comme Qwen 3 8B, mais les détails restent limités. Deepgram fait état d'une amélioration de 60 % à 70 % de l'efficacité du traitement audio, mais les tests de référence exacts ne sont pas entièrement transparents.

Le projet d'optimisation x86 d'AMD et Intel

AMD et Intel ont annoncé leur intention de regrouper leurs optimisations PyTorch pour les processeurs x86 au sein de la bibliothèque oneDNN, suivant les recommandations du groupe consultatif sur l'écosystème x86. Cette démarche vise à offrir une voie unifiée pour l'accélération x86 dans PyTorch, avec des mises à jour qui devraient inclure de nouveaux jeux d'instructions pour la multiplication matricielle.