AWS Neuron memperkenalkan Antarmuka Kernel Neuron (NKI), Pelatihan NxD, dan dukungan JAX untuk pelatihan
Hari ini, AWS mengumumkan peluncuran Neuron 2.20, memperkenalkan Neuron Kernel Interface (NKI) (beta), antarmuka pemrograman untuk AWS Trainium dan Inferentia, yang memungkinkan pengembang membangun kernel komputasi yang dioptimalkan untuk fungsionalitas baru, pengoptimalan, dan inovasi sains. Selain itu, rilis ini memperkenalkan NxD Training (beta), pustaka berbasis PyTorch yang memungkinkan pelatihan terdistribusi yang efisien, dengan antarmuka yang mudah digunakan yang kompatibel dengan NeMo. Rilis ini juga memperkenalkan dukungan untuk kerangka kerja JAX (beta).
AWS Neuron adalah SDK untuk instans berbasis AWS Inferentia dan Trainium yang dibuat khusus untuk AI generatif. Neuron terintegrasi dengan kerangka kerja ML populer seperti PyTorch. Neuron termasuk compiler, runtime, alat, dan pustaka untuk mendukung pelatihan kinerja tinggi dan inferensi model AI pada instans Trn1 dan Inf2.
Rilis ini menambahkan fitur dukungan dan peningkatan kinerja untuk pelatihan dan inferensi model. Untuk pelatihan, rilis ini menambahkan model Llama 3.1 8B dan 70B yang mendukung panjang urutan hingga 32K, bersama dengan torch.autocast() untuk dukungan presisi campuran PyTorch asli dan teknik PEFT LoRA. Untuk inferensi, Neuron 2.20 menambahkan dukungan untuk model Llama 3.1 (405b, 70b, 8b) dan Diffusion-Transformers (DiT) seperti Pixart-alpha dan Pixart-sigma. Selain itu, rilis ini menambahkan dukungan inferensi dengan pengambilan sampel top-p pada perangkat dan panjang konteks 128K dengan Flash Decoding. Rilis ini juga menambahkan dukungan untuk sistem operasi Rocky 9.0 dan operator RMSNorm dan RMSNormDx di Neuron Compiler.
Untuk informasi lebih lanjut, lihat Catatan Rilis Neuron.