AWS Neuron introduce Neuron Kernel Interface (NKI), NxD Training e il supporto JAX per l'addestramento
Oggi AWS ha presentato Neuron 2.20, introducendo Neuron Kernel Interface (NKI) (beta), un'interfaccia di programmazione per AWS Trainium e Inferentia, che consente agli sviluppatori di creare kernel di calcolo ottimizzati per nuove funzionalità, ottimizzazioni e innovazioni scientifiche. Inoltre, questa versione introduce NxD Training (beta), una libreria basata su PyTorch che consente un addestramento distribuito efficiente, con un'interfaccia intuitiva compatibile con NeMo. Questa versione introduce anche il supporto per il framework JAX (beta).
AWS Neuron è l'SDK per le istanze basate su AWS Inferentia e Trainium, creato appositamente per l'IA generativa. Neuron si integra con i framework ML più diffusi come PyTorch. Include un compilatore, un runtime, strumenti e librerie per supportare l'addestramento ad alte prestazioni e l'inferenza di modelli IA su istanze Trn1 e Inf2.
Questa versione aggiunge funzionalità di supporto e miglioramenti delle prestazioni per l'inferenza e l'addestramento dei modelli. Per l'addestramento questa versione aggiunge i modelli Llama 3.1 8B e 70B che supportano una lunghezza di sequenza fino a 32K, insieme a torch.autocast() per il supporto nativo di PyTorch con precisione mista e le tecniche PEFT LoRA. Per l'inferenza Neuron 2.20 aggiunge il supporto per i modelli Llama 3.1 (405b, 70b, 8b) e Diffusion-Transformers (DiT) come Pixart-alpha e Pixart-sigma. Inoltre, questa versione aggiunge il supporto per l'inferenza con campionamento top-p sul dispositivo e lunghezza del contesto di 128K con decodifica flash. Questa versione aggiunge anche il supporto per il sistema operativo Rocky 9.0 e gli operatori RMSNorm e RMSNormDx nel compilatore Neuron.
Per ulteriori informazioni consulta le note di rilascio di Neuron.