AWS Neuron présente l'interface Neuron Kernel (NKI), la formation NxD et la prise en charge de JAX pour la formation

Publié le: 17 sept. 2024

AWS annonce aujourd'hui la sortie de Neuron 2.20, qui présente Neuron Kernel Interface (NKI) (bêta), une interface de programmation pour AWS Trainium et Inferentia, qui permet aux développeurs de créer des noyaux de calcul optimisés pour de nouvelles fonctionnalités, des optimisations et des innovations scientifiques. En outre, cette version introduit NxD Training (bêta), une bibliothèque basée sur PyTorch permettant une formation distribuée efficace, avec une interface conviviale compatible avec NeMo. Cette version introduit également la prise en charge du framework JAX (bêta).

AWS Neuron est le kit SDK pour les instances basées sur AWS Inferentia et Trainium, spécialement conçu pour l'IA générative. Neuron s'intègre aux cadres de machine learning populaires tels que PyTorch. Il inclut un compilateur, un environnement d'exécution, des outils et des bibliothèques pour prendre en charge la formation et l'inférence haute performance de modèles d'IA sur les instances Trn1 et Inf2.

Cette version ajoute des fonctionnalités de prise en charge et des améliorations de performances pour l’entraînement et l'inférence des modèles. Pour la formation, cette version ajoute les modèles Llama 3.1 8B et 70B prenant en charge une longueur de séquence allant jusqu'à 32 K, ainsi que torch.autocast () pour le support de précision mixte natif de PyTorch et les techniques PEFT LoRA. À des fins d'inférence, Neuron 2.20 prend en charge les modèles Llama 3.1 (405b, 70b, 8b) et Diffusion-Transformers (DiT) tels que Pixart-alpha et Pixart-Sigma. En outre, cette version ajoute la prise en charge de l'inférence avec un échantillonnage de pointe sur l'appareil et une longueur de contexte de 128 Ko avec le décodage Flash. Cette version ajoute également la prise en charge du système d'exploitation Rocky 9.0 et des opérateurs RMSNorm et RMSNormDx dans le Neuron Compiler.

Pour en savoir plus, consultez les notes de version de Neuron.