AWS Neuron apresenta a Neuron Kernel Interface (NKI), o NxD Training e o suporte do JAX para treinamento

Publicado: 17 de set de 2024

Hoje, a AWS anuncia o lançamento do Neuron 2.20, apresentando a Neuron Kernel Interface (NKI) (beta), uma interface de programação para AWS Trainium e Inferentia, permitindo que os desenvolvedores criem kernels de computação otimizados para novas funcionalidades, otimizações e inovações científicas. Além disso, esta versão apresenta o NxD Training (beta), uma biblioteca baseada em PyTorch que permite treinamento distribuído eficiente, com uma interface amigável compatível com o NeMo. Esta versão também apresenta suporte para a estrutura JAX (beta).

O AWS Neuron é o SDK criado especificamente para IA generativa em instâncias baseadas no AWS Inferentia e no Trainium. O Neuron se integra a estruturas de machine learning populares, como PyTorch. Ele inclui compilador, runtime, ferramentas e bibliotecas para oferecer suporte ao treinamento e inferência de alta performance de modelos de IA nas instâncias Trn1 e Inf2.

Esta versão adiciona recursos de suporte e melhorias de performance para inferência e treinamento de modelos. Para treinamento, esta versão adiciona os modelos Llama 3.1 8B e 70B com suporte de sequência de até 32K, junto com torch.autocast () para suporte nativo de precisão mista do PyTorch e técnicas PEFT LoRA. Para inferência, o Neuron 2.20 adiciona suporte para os modelos Llama 3.1 (405b, 70b, 8b) e Diffusion-Transformers (DiT), como Pixart-alpha e Pixart-sigma. Além disso, esta versão adiciona suporte à inferência com amostragem top-p no dispositivo e tamanho de contexto de 128K com decodificação Flash. Esta versão também adiciona suporte ao sistema operacional Rocky 9.0 e aos operadores RMSNorm e RMSNormDx no Neuron Compiler.

Para obter mais informações, consulte as notas de versão do Neuron.