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 no Trainium2, são projetadas com propósito específico para IA generativa e representam as instâncias do EC2 mais avançadas para treinar e implantar modelos com centenas de bilhões a trilhões de parâmetros. As instâncias Trn2 fornecem uma performance de preço de 30% a 40% superior em comparação com a geração atual das instâncias P5e e P5en do EC2 que são 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 prévia), são perfeitos 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 projeto do UltraServer usa o NeuronLink para conectar 64 chips do Trainium2 distribuídos em quatro instâncias Trn2, em um único nó. Isso disponibiliza novas funcionalidades. Para tarefas de inferência, os servidores UltraServers contribuem para proporcionar tempos de resposta de ponta, 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

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.

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.

Para fornecer alta performance enquanto atende às metas de precisão, os chips do 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.

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.

As instâncias Trn2 são projetadas para serem três vezes mais eficientes em termos de energia do que as instâncias Trn1. As instâncias Trn1 são até 25% mais eficientes em termos de energia do que as instâncias de computação acelerada comparáveis do EC2. Essas instâncias contribuem para o cumprimento de suas metas de sustentabilidade ao treinar modelos muito grandes.

Vídeos

Os bastidores da infraestrutura de IA generativa na Amazon
Acelere o DL e inove mais rápido com o AWS Trainium
Introdução às instâncias Trn1 do Amazon EC2 com a tecnologia AWS Trainium