Crie no Trainium

O Build on Trainium é um programa de investimento de USD 110 milhões focado em pesquisa de IA e educação universitária para apoiar a próxima geração de inovação e desenvolvimento no AWS Trainium. O AWS Trainium é um chip de matriz sistólica de IA projetado exclusivamente para o avanço de ideias e aplicações de IA de última geração. O Build on Trainium financia novas pesquisas de IA no Trainium, investindo nas principais equipes acadêmicas para desenvolver inovações em áreas críticas, incluindo novas arquiteturas de modelos, bibliotecas de ML, otimizações, sistemas distribuídos em larga escala e muito mais. Essa iniciativa plurianual estabelece a base para o futuro da IA, inspirando a comunidade acadêmica a aproveitar, investir e contribuir com a comunidade de código aberto em torno do Trainium. Combinando esses benefícios com o kit de desenvolvimento de software (SDK) da Neuron e o recente lançamento da Neuron Kernel Interface (NKI), os clientes da Trainium agora podem inovar em escala na nuvem.

Criamos um cluster de pesquisa Trainium dedicado com até 40.000 chips Trainium que estará disponível por meio de instâncias Trn1 do Amazon EC2 conectadas em uma única rede de escala de peta-bits sem bloqueio usando o Amazon EC2 UltraClusters. As equipes de pesquisa e os alunos podem acessar esses chips por meio de reservas de blocos de capacidade autogerenciados usando o Blocos de Capacidade do Amazon EC2 para ML.

Estamos conduzindo várias rodadas de chamadas de propostas (CFP) do Amazon Research Awards (ARA) para a ampla comunidade de pesquisa. As propostas selecionadas receberão créditos do AWS Trainium e acesso ao cluster de pesquisa Trainium. O Build on Trainium aceita propostas de pesquisa que utilizarão bibliotecas e estruturas muito usadas de ML de código aberto e contribuirão com o código aberto para aprimorar os recursos para a comunidade de desenvolvedores de ML.

A Neuron Kernel Interface (NKI) é uma nova interface de programação para os chips de IA da AWS, Trainium e Inferentia. A NKI fornece acesso direto a primitivos e instruções de hardware disponíveis no AWS Trainium e no Inferentia, permitindo que os pesquisadores criem e ajustem kernels de computação para obter o desempenho 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. Visite a página de documentação do NKI para saber mais.

No Device Realization Lab do MIT, estamos usando o AWS Trainium para ultrapassar os limites da pesquisa médica de IA. Nossos modelos de segmentação de ultrassom 3D e estimativa da velocidade do som treinam com mais rapidez e eficiência do que nunca, reduzindo o tempo experimental em mais da metade e alcançando precisão de última geração. O AWS Trainium nos permitiu escalar nossa pesquisa de maneiras que eram inviáveis com sistemas de GPU tradicionais. Ao treinar nossas redes neurais 3D totalmente convolucionais no AWS Trainium (trn.32xlarge), alcançamos um desempenho de última geração com throughput 50% maior e menor custo em comparação com as instâncias NVIDIA A100. Usando um cluster Trainium de 32 nós, conduzimos mais de 180 experimentos de ablação de forma eficiente, reduzindo o tempo total de treinamento de meses para semanas e acelerando a inovação da IA médica no MIT. No futuro, planejamos usar o Trainium para treinar modelos de agentes de IA que possam operar e automatizar o fluxo de trabalho de ultrassom digital, economizando tempo significativo do médico e oferecendo melhores cuidados aos pacientes. “

“O grupo de pesquisa CMU Catalyst trabalha na otimização de sistemas de ML. Nosso projeto visa a facilitar a otimização de diferentes sistemas de ML. O Trainium é único a fornecer controle de baixo nível e uma interface de programação acessível por meio da Neuron Kernel Interface (NKI).

com o apoio da AWS por meio do programa Build on Trainium, nosso pesquisador conseguiu explorar otimizações avançadas em um kernel crítico: o FlashAttention. O que mais nos surpreendeu foi a velocidade com que conseguimos iterar: alcançamos melhorias significativas além do estado da arte anterior em apenas uma semana usando o NKI, o Neuron Profiler e a documentação de arquitetura disponíveis publicamente. A combinação de ferramentas poderosas e insights claros sobre o hardware tornou a otimização sofisticada de baixo nível acessível à nossa equipe.

O AWS Trainium e a Neuron Kernel Interface (NKI) capacitam pesquisadores como nós para inovar mais rapidamente, removendo barreiras que normalmente retardam o trabalho de otimização específico para o hardware. “

 

“Por meio do programa Build on Trainium, sua equipe obteve acesso total à nova pilha de compiladores de código aberto NKI da AWS Neuron, incluindo visibilidade direta do ISA e das APIs do Trainium para agendamento preciso e alocação de memória. Esse nível de visibilidade e controle permite que seus alunos analisem com mais facilidade as oportunidades de otimização e descubram implementações de alto desempenho com mais eficiência. “

Christopher Fletcher, professor associado de ciência da computação, University da California, Berkeley
 

“O acesso ao AWS Trainium e ao Inferentia tem sido fundamental para o avanço de nossa pesquisa e capacitação em sistemas de IA eficientes e de grande escala. Usamos essas plataformas para treinamento misto de especialistas e otimizações de inferência, prototipando novas técnicas de runtime e agendamento que melhoram a escalabilidade, a eficiência e a portabilidade em arquiteturas de aceleradores emergentes. Ao utilizar a pilha Neuron Developer, os pesquisadores da UIUC estão desenvolvendo novas técnicas de runtime e agendamento que aumentam a eficiência e a portabilidade dos workloads de IA. A equipe está particularmente impressionada com a abertura da pilha Neuron Developer, a qual torna essas plataformas valiosas para pesquisas de runtime e permitem inovações em dispersão, hierarquias de memória e eficiência de comunicação que vão além das arquiteturas tradicionais de GPU. “

“Ao utilizar o AWS Trainium e o programa Build on Trainium, meus alunos e eu conseguimos acelerar significativamente nossas simulações de circuitos quânticos. O projeto reuniu um sólido grupo de estudantes que construíram colaborativamente um simulador de alto desempenho, possibilitando uma experimentação mais profunda e o aprendizado prático em uma escala que simplesmente não era possível antes. “

“Nossa equipe de pesquisa na UTS está explorando a integração de algoritmos de marca d'água em anéis de árvores desenvolvendo kernels NKI personalizados do Neuron. Ter acesso à pilha Neuron de código aberto por meio do programa Build on Trainium foi transformador. Ele nos dá uma visibilidade sem precedentes da arquitetura Trainium e a capacidade de trabalhar diretamente no nível do hardware. O acesso ao Trainium possibilitou que nossa equipe acelerasse significativamente o workload de marca d'água, reduzindo os ciclos de iteração e permitindo que explorássemos modelos e técnicas mais complexos. Essa profundidade de acesso possibilita que nossos pesquisadores criem novas ideias prototípicas, experimentem otimizações de baixo nível e ultrapassem os limites do que os sistemas de marca d'água podem alcançar nos aceleradores de IA modernos. “