AWS Neuron

AWS Neuron es el kit de desarrollo de software (SDK) que se utiliza para ejecutar cargas de trabajo de IA generativa y aprendizaje profundo en instancias de Amazon Elastic Compute Cloud (Amazon EC2) impulsadas por AWS Inferentia y AWS Trainium. Incluye un compilador, un tiempo de ejecución, bibliotecas de entrenamiento e inferencia y herramientas para desarrolladores para la supervisión, la creación de perfiles y la depuración. Neuron respalda su ciclo de vida de desarrollo de machine learning (ML) de extremo a extremo, incluida la creación e implementación de modelos de IA y aprendizaje profundo, la optimización para lograr el mayor rendimiento y el menor costo, y la obtención de información más profunda sobre el comportamiento de los modelos.

Neuron se integra de forma nativa en PyTorch, JAX y bibliotecas de ML esenciales, como Optimum Neuron de Hugging Face, PyTorch Lightning y AXLearn. Neuron también es compatible con OpenXLA, incluidos StableHLO y GSPMD, lo que permite a los desarrolladores de PyTorch, XLA y JAX utilizar las optimizaciones del compilador de Neuron para Inferentia y Trainium. Neuron le permite usar instancias basadas en Trainium e Inferentia con servicios como Amazon SageMaker, Amazon EKS, Amazon ECS, AWS ParallelCluster y AWS Batch, así como con servicios de terceros como Ray (Anyscale), Domino Data Lab, Datadog y Weights & Biases.

Neuron incluye optimizaciones listas para usar para el entrenamiento y la inferencia distribuidos con las bibliotecas de PyTorch de código abierto NxD Training y NxD Inference. NxD Training simplifica y optimiza el entrenamiento distribuido a gran escala y admite varias arquitecturas de modelos, estrategias de paralelismo y flujos de trabajo de entrenamiento. NxD Inference proporciona una solución integral para la inferencia optimizada de modelos con características clave como el muestreo en el dispositivo, la fusión de pesos QKV, el procesamiento continuo por lotes, la decodificación especulativa, la agrupación dinámica y la inferencia distribuida. NxD Inference también se integra en soluciones de servicio como vLLM y Hugging Face TGI. Ambos incluyen un centro de modelos para diferentes arquitecturas de modelos.

Neuron tiene varias capacidades de ciencia aplicada que permiten a los científicos e investigadores ampliar los límites de la investigación e innovación en IA de código abierto en Trainium e Inferentia. Neuron Kernel Interface (NKI) proporciona acceso directo a las primitivas de hardware y a las instrucciones disponibles en Trainium e Inferentia, lo que permite a los investigadores crear y ajustar kernels de computación para lograr un rendimiento óptimo. Es un entorno de programación basado en Python que adopta una sintaxis similar a la de Triton y una semántica a nivel de mosaico de uso común. Los investigadores pueden usar NKI para mejorar los modelos de aprendizaje profundo con nuevas funcionalidades, optimizaciones e innovaciones científicas. Los operadores C++ personalizados de Neuron permiten a los desarrolladores ampliar la funcionalidad del SDK mediante la creación de sus propios operadores optimizados para Inferentia y Trainium.

El SDK de AWS Neuron ofrece un conjunto de herramientas completo para proporcionar información detallada sobre la supervisión, la administración y la optimización de los modelos de aprendizaje profundo en las instancias de EC2 impulsadas por AWS Inferentia y Trainium. Proporciona utilidades como neuron-top, neuron-monitor y Neuron Sysfs para supervisar los recursos de hardware, la ejecución del modelo y los detalles del sistema. Para las aplicaciones en contenedores en Kubernetes y EKS, Neuron simplifica la supervisión mediante la integración de Amazon CloudWatch y otras herramientas de observabilidad populares, como Data Dog y Weights & Biases. Además, la herramienta neuron-profile ayuda a identificar y abordar los cuellos de botella de rendimiento tanto en aplicaciones distribuidas como de un solo nodo, y proporciona capacidades de creación de perfiles nativas para los marcos de aprendizaje automático populares.