L'usage des architectures spécialisées comme AWS Trainium nécessite souvent des adaptations de code. Pour simplifier cette transition, un nouvel outil (Space) sur Hugging Face permet de vérifier si un modèle peut être exécuté directement sur ces puces grâce à TorchNeuron, sans modification structurelle du modèle source.
Des indicateurs de performance et de mémoire
L'interface ne se contente pas de confirmer la compatibilité. Pour chaque modèle répertorié, l'outil fournit des données concrètes sur l'utilisation des ressources :
- La répartition du modèle sur les cœurs Trainium (sharding) ;
- Le temps d'exécution total (incluant la compilation) ;
- La consommation maximale de la mémoire du périphérique (peak device memory) source.
Il est important de noter que ces résultats sont limités : les tests actuels ne concernent que les modèles capables de tenir sur un maximum de quatre instances Trainium2 (soit 16 cœurs) et se limitent aux tâches de génération de texte et de vision-langage source. Un résultat positif (« PASS ») garantit que le modèle s'exécute correctement, mais ne renseigne pas sur le débit réel (throughput) ni sur la capacité de génération longue source.
Consultation et méthodologie
L'outil propose deux modes de consultation. Une sélection de modèles est présentée par défaut, basée sur les statistiques de téléchargement et de tendance du Hub. Pour une recherche spécifique, l'utilisateur peut naviguer dans les classements par tâche et par volume de téléchargements source.
Le processus de test repose sur un script standardisé qui charge le modèle via from_pretrained et utilise torch.compile avec le backend neuron. Pour les modèles qui échouent, l'outil fournit également l'erreur rencontrée, permettant d'identifier les incompatibilités de la stack logicielle source.