O que é o AWS Neuron?
O AWS Neuron é o kit de desenvolvimento de software (SDK) usado para executar workloads de aprendizado profundo e de IA generativa em instâncias do Amazon Elastic Compute Cloud (Amazon EC2) com tecnologia do AWS Inferentia e do AWS Trainium. Essa solução inclui um compilador, um runtime, bibliotecas de treinamento e de inferência, e ferramentas de desenvolvedor para monitoramento, criação de perfil e depuração. O Neuron oferece suporte ao ciclo de vida completo de desenvolvimento de machine learning (ML), incluindo a criação e a implementação de modelos de aprendizado profundo e IA, otimização para alcançar a melhor performance e o menor custo, e obtenção de análises mais aprofundadas sobre o comportamento dos modelos.
Integração nativa com estruturas e bibliotecas conhecidas de machine learning
O Neuron se integra nativamente com PyTorch e JAX, bem como com bibliotecas essenciais de machine learning, como Hugging Face Optimum Neuron, PyTorch Lightning e AXLearn. Além disso, o Neuron oferece suporte ao OpenXLA, incluindo StableHLO e GSPMD. Isso possibilita que os desenvolvedores de PyTorch, XLA e JAX usem as otimizações do compilador do Neuron para o Inferentia e o Trainium. O Neuron possibilita o uso de instâncias baseadas no Trainium e no Inferentia com serviços como o Amazon SageMaker, o Amazon EKS, o Amazon ECS, o AWS ParallelCluster e o AWS Batch, bem como com serviços de entidades externas como Ray (Anyscale), Domino Data Lab, Datadog e Weights & Biases.
Bibliotecas para treinamento e inferência distribuídos
O Neuron inclui otimizações prontas para uso para treinamento e inferência distribuídos com as bibliotecas de código aberto NxD Training e NxD Inference do PyTorch. O NxD Training simplifica e otimiza o treinamento distribuído em grande escala, fornecendo suporte a diversas arquiteturas de modelo, estratégias de paralelismo e fluxos de trabalho de treinamento. O NxD Inference disponibiliza uma solução completa para inferência otimizada de modelos, com recursos importantes como amostragem no dispositivo, fusão de pesos QKV, processamento contínuo em lote, decodificação especulativa, geração de buckets dinâmica e inferência distribuída. Além disso, o NxD Inference se integra a soluções de serviço como o vLLM e o TGI da Hugging Face. Ambos incluem um hub de modelos para diferentes arquiteturas de modelo.
Funcionalidades avançadas de ciência aplicada
O Neuron contém diversas funcionalidades de ciência aplicada para capacitar cientistas e pesquisadores a expandir os limites da pesquisa e de inovação em IA de código aberto utilizando o Trainium e o Inferentia. A Neuron Kernel Interface (NKI) disponibiliza acesso direto às primitivas de hardware e às instruções presentes no Trainium e no Inferentia, possibilitando o desenvolvimento e o ajuste de kernels de computação para a obtenção de uma performance ideal. É um ambiente de programação baseado em Python que adota uma sintaxe semelhante à do Triton e uma semântica em nível de blocos comumente usadas. Os pesquisadores podem usar o NKI para aprimorar os modelos de aprendizado profundo com novas funcionalidades, otimizações e inovações científicas. Os operadores personalizados em C++ do Neuron permitem que os desenvolvedores expandam a funcionalidade do SDK ao criar seus próprios operadores otimizados para o Inferentia e o Trainium.
Ferramentas avançadas para desenvolvedores
O AWS Neuron SDK oferece um conjunto abrangente de ferramentas para fornecer insights detalhados sobre o monitoramento, o gerenciamento e a otimização de modelos de aprendizado profundo em instâncias do EC2 com tecnologia do AWS Inferentia e do Trainium. Essa solução disponibiliza utilitários, como neuron-top, neuron-monitor e Neuron Sysfs, para monitorar os recursos de hardware, a execução de modelos e os detalhes do sistema. Para as aplicações conteinerizadas no Kubernetes e no EKS, o Neuron simplifica o monitoramento por meio da integração com o Amazon CloudWatch e com outras ferramentas conhecidas de observabilidade, como Data Dog e Weights & Biases. Além disso, a ferramenta neuron-profile auxilia na identificação e na resolução de gargalos de performance em aplicações de nó único e distribuídas, além de fornecer funcionalidades nativas de perfilamento para estruturas de machine learning (ML) conhecidas.
