AWS Neuron introduces Neuron Kernel Interface (NKI), NxD Training, and JAX support for training

The article announces the release of Neuron 2.20, which introduces several new features and improvements for AWS Neuron, the SDK for AWS Inferentia and Trainium based instances designed for generative AI.

Specifically, the article covers:

• Neuron Kernel Interface (NKI) (beta), a programming interface for AWS Trainium and Inferentia, allowing developers to build optimized compute kernels for new functionalities, optimizations, and scientific innovations.
• NxD Training (beta), a PyTorch-based library for efficient distributed training, with a user-friendly interface compatible with NeMo.
• JAX framework support (beta).
• Improved support for training models like Llama 3.1 8B and 70B with up to 32K sequence length, torch.autocast() for native PyTorch mixed precision, and PEFT LoRA techniques.
• Improved support for inference with models like Llama 3.1 (405b, 70b, 8b) and Diffusion-Transformers (DiT) models like Pixart-alpha and Pixart-sigma, top-p sampling on device, and 128K context length with Flash Decoding.
• Support for Rocky 9.0 operating system and new operators like RMSNorm and RMSNormDx in the Neuron Compiler.