AWS Trainium
Obtenha alta performance para treinamento de aprendizado profundo e IA generativa e, ao mesmo tempo, reduza os custos
Por que o Trainium?
Os chips do AWS Trainium correspondem a uma família de chips de IA desenvolvidos pela AWS com o propósito específico de treinamento e inferência de IA para fornecer alta performance e redução de custos.
O chip do AWS Trainium de primeira geração é responsável pelas instâncias Trn1 do Amazon Elastic Compute Cloud (Amazon EC2), que apresentam custos de treinamento até 50% menores do que as instâncias comparáveis do Amazon EC2. Diversos clientes, como Databricks, Ricoh, NinjaTech AI e Arcee AI, estão experimentando benefícios de performance e de redução de custos com as instâncias Trn1.
O chip do AWS Trainium2 proporciona uma performance até quatro vezes superior quando comparado com o chip do Trainium de primeira geração. As instâncias Trn2 do Amazon EC2 baseadas em Trainium2 foram desenvolvidas com propósito específico para IA generativa e oferecem um custo-benefício 30-40% melhor em comparação com a geração atual de instâncias P5e e P5en do EC2 baseadas em GPU. As instâncias Trn2 são equipadas com 16 chips do Trainium2 interconectados pelo NeuronLink, nosso interconector proprietário “chip-to-chip”. É possível usar as instâncias Trn2 para treinar e implantar modelos mais complexos, incluindo grandes modelos de linguagem (LLMs), modelos multimodais e transformadores de difusão, para o desenvolvimento de uma ampla variedade de aplicações de IA generativa da próxima geração. As instâncias Trn2 UltraServers, uma oferta totalmente nova do EC2 (disponível em versão preliminar), são perfeitas para os modelos mais complexos que necessitam de mais memória e largura de banda de memória do que as instâncias autônomas do EC2 conseguem disponibilizar. O design do UltraServer utiliza o NeuronLink para conectar em um único nó 64 chips do Trainium2 distribuídos em quatro instâncias Trn2, desbloqueando novas capacidades. Para tarefas de inferência, os servidores UltraServers contribuem para proporcionar menores tempos de resposta, garantindo as melhores experiências em tempo real. Para tarefas de treinamento, os UltraServers aumentam a velocidade e a eficiência do treinamento de modelos com uma comunicação coletiva mais rápida para o paralelismo de modelos, em comparação com as instâncias autônomas.
É possível começar a treinar e a implantar modelos nas instâncias Trn2 e Trn1 com suporte nativo para estruturas conhecidas de machine learning (ML), como PyTorch e JAX.
Benefícios
IA generativa econômica e com alta performance
As instâncias e os servidores Trn2 UltraServers fornecem uma performance revolucionária no Amazon EC2 para treinamento e inferência de IA generativa. Cada servidor Trn2 UltraServer contém 64 chips do Trainium2 que estão interconectados pelo NeuronLink, nosso interconector proprietário “chip-to-chip”, e oferece até 83,2 petaflops de computação FP8, 6 TB de HBM3 com 185 terabytes por segundo (TBps) de largura de banda de memória, e 12,8 terabits por segundo (Tbps) de rede do Elastic Fabric Adapter (EFA). Cada instância Trn2 contém 16 chips do Trainium2 que estão conectados pelo NeuronLink e oferece até 20,8 petaflops de computação FP8, 1,5 TB de HBM3 com 46 TBps de largura de banda de memória e 3,2 Tbps de rede do EFA. A instância Trn1 contém até 16 chips do Trainium e oferece até 3 petaflops de computação FP8, 512 GB de HBM com 9,8 TBps de largura de banda de memória e até 1,6 Tbps de rede do EFA.
Suporte nativo para frameworks e bibliotecas de ML
O AWS Neuron SDK ajuda você a extrair toda a performance das instâncias Trn2 e Trn1, permitindo que você se concentre no desenvolvimento e na implantação de modelos, acelerando o tempo de lançamento no mercado. O AWS Neuron fornece integração nativa com JAX, PyTorch e bibliotecas fundamentais como Hugging Face, PyTorch Lightning e NeMo. O AWS Neuron é compatível com mais de 100 mil modelos no hub de modelos da Hugging Face, incluindo modelos conhecidos, como a família Llama da Meta e o Stable Diffusion XL. Essa solução otimiza os modelos automaticamente para treinamento e inferência distribuídos, ao mesmo tempo em que oferece insights aprofundados para criação de perfis e depuração. O AWS Neuron se integra com serviços como o Amazon SageMaker, o Amazon Elastic Kubernetes Service (Amazon EKS), o Amazon Elastic Container Service (Amazon ECS), o AWS ParallelCluster e o AWS Batch, bem como com serviços de entidades externas, como Ray (Anyscale), Domino Data Lab e Datadog.
Otimizações de IA de última geração
Para fornecer alta performance enquanto atende às metas de precisão, os chips Trainium são otimizados para os tipos de dados FP32, TF32, BF16, FP16 e o novo tipo de dado configurável FP8 (cFP8). Para sustentar o ritmo acelerado de inovação na IA generativa, o Trainium2 oferece otimizações de hardware com dispersão de quatro vezes (16:4), escalabilidade mínima, arredondamento estocástico e mecanismos coletivos dedicados.
Projetado para pesquisa em IA
A Neuron Kernel Interface (NKI) possibilita o acesso direto à arquitetura do conjunto de instruções (ISA) ao usar um ambiente baseado em Python com uma interface semelhante ao Triton, permitindo que você realize inovações em novas arquiteturas de modelos e núcleos de computação altamente otimizados que superam as técnicas existentes.