AWS Trainium

Obtenga un alto rendimiento para la formación en aprendizaje profundo e IA generativa y, al mismo tiempo, reduzca los costos

¿Por qué Trainium?

Los chips de AWS Trainium son una familia de chips de IA diseñados específicamente por AWS para el entrenamiento y la inferencia de la IA con el fin de ofrecer un alto rendimiento y reducir los costos.

El chip de AWS Trainium de primera generación impulsa las instancias Trn1 de Amazon Elastic Compute Cloud (Amazon EC2), que tienen costos de capacitación hasta un 50 % más bajos que las instancias de Amazon EC2 similares. Muchos clientes, incluidos Databricks, Ricoh, NinjaTech AI y Arcee AI, están descubriendo los beneficios de rendimiento y costo de las instancias Trn1.

El chip AWS Trainium2 ofrece hasta cuatro veces más rendimiento que el Trainium de primera generación. Las instancias Trn2 de Amazon EC2 basadas en Trainium2 están diseñadas específicamente para la IA generativa y son las instancias de EC2 más potentes para entrenar e implementar modelos con entre cientos de miles de millones y más de un billón de parámetros. Las instancias Trn2 ofrecen entre un 30 y un 40 % de mejor rendimiento de precio que la generación actual de instancias P5e y P5en de EC2 basadas en GPU. Las instancias Trn2 cuentan con 16 chips Trainium2 interconectados con NeuronLink, nuestra interconexión chip a chip patentada. Puede usar las instancias de Trn2 para entrenar e implementar los modelos más exigentes, incluidos los modelos de lenguaje grande (LLM), los modelos multimodales y los transformadores de difusión, para crear un amplio conjunto de aplicaciones de IA generativa de última generación. Los UltraServers Trn2 son una oferta de EC2 totalmente nueva (disponible en versión preliminar) y son ideales para los modelos de mayor tamaño que requieren más memoria y ancho de banda de memoria que los que pueden proporcionar las instancias de EC2 independientes. El diseño de UltraServer utiliza NeuronLink para conectar 64 chips Trainium2 a cuatro instancias Trn2 en un nodo, lo que permite nuevas capacidades. A modo de inferencia, los UltraServers ayudan a ofrecer un tiempo de respuesta líder en la industria para crear las mejores experiencias en tiempo real. En cuanto al entrenamiento, los UltraServers aumentan la velocidad y la eficiencia del entrenamiento de modelos con una comunicación colectiva más rápida para lograr el paralelismo de los modelos en comparación con las instancias independientes.

Puede empezar a entrenar e implementar modelos en instancias Trn2 y Trn1 con soporte nativo para marcos de machine learning (ML) populares, como PyTorch y JAX.

Beneficios

Los UltraServer y las instancias Trn2 ofrecen un rendimiento sin precedentes en Amazon EC2 para el entrenamiento y la inferencia de la IA generativa. Cada UltraServer Trn2 tiene 64 chips Trainium2 interconectados con NeuronLink, nuestra interconexión exclusiva de chip a chip, y ofrece hasta 83,2 petaflops de procesamiento FP8, 6 TB de HBM3 con 185 terabytes por segundo (TBps) de ancho de banda de memoria y 12,8 terabits por segundo (Tbps) de red Elastic Fabric Adapter (EFA). Cada instancia Trn2 tiene 16 chips Trainium2 conectados con NeuronLink y ofrece hasta 20,8 petaflops de procesamiento FP8, 1,5 TB de HBM3 con 46 TBps de ancho de banda de memoria y 3,2 Tbps de red EFA. La instancia Trn1 cuenta con hasta 16 chips Trainium y ofrece hasta 3 petaflops de procesamiento FP8, 512 GB de HBM con 9,8 Tbps de ancho de banda de memoria y hasta 1,6 Tbps de redes EFA.

El SDK de AWS Neuron lo ayuda a extraer todo el rendimiento de las instancias Trn2 y Trn1 para que pueda centrarse en crear e implementar modelos, y acelerar el tiempo de comercialización. AWS Neuron se integra de forma nativa con JAX, PyTorch y bibliotecas esenciales como Hugging Face, PyTorch Lightning y NeMo. AWS Neuron admite más de 100 000 modelos en el centro de modelos Hugging Face, incluidos los populares como la familia de modelos Meta Llama y Stable Diffusion XL. Optimiza los modelos listos para usar para el entrenamiento y la inferencia distribuidos, al mismo tiempo que proporciona información profunda para la creación de perfiles y la depuración. AWS Neuron se integra con servicios como Amazon SageMaker, Amazon Elastic Kubernetes Service (Amazon EKS), Amazon Elastic Container Service (Amazon ECS), AWS ParallelCluster y AWS Batch, así como con servicios de terceros como Ray (Anyscale), Domino Data Lab y Datadog.

Para ofrecer un alto rendimiento y, al mismo tiempo, cumplir con los objetivos de precisión, los chips de Trainium están optimizados para FP32, TF32, BF16, FP16, y el nuevo tipo de datos FP8 configurable (cFP8). Para respaldar el rápido ritmo de innovación en la IA generativa, Trainium2 cuenta con optimizaciones de hardware para una dispersión 4 veces mayor (16:4), el microescalado, el redondeo estocástico y los motores colectivos dedicados.

Neuron Kernel Interface (NKI) permite el acceso directo a la arquitectura del conjunto de instrucciones (ISA) mediante un entorno basado en Python con una interfaz similar a Triton, lo que le permite innovar en arquitecturas de nuevos modelos y kernels de computación altamente optimizados que superan a las técnicas existentes.

Las instancias Trn2 están diseñadas para ser tres veces más eficientes desde el punto de vista energético que las instancias Trn1. Las instancias Trn1 son hasta un 25 % más eficientes desde el punto de vista energético que las instancias de EC2 de computación acelerada comparables. Estas instancias ayudan a alcanzar sus objetivos de sostenibilidad cuando entrena modelos ultragrandes.

Videos

Detrás de cámaras en la infraestructura de IA generativa en Amazon
Aceleración del DL e innovación más rápida con AWS Trainium
Introducción a las instancias Trn1 de Amazon EC2 impulsadas por AWS Trainium