Aceleradores AI AWS Trainium: 55 bilhões de transistores, 3 GHz, 512 GB HBM e 840 teraflops em FP32

As GPUs têm sido usadas há muito tempo para acelerar cálculos e, nos últimos anos, cresceram demais com suporte para formatos de dados específicos típicos de algoritmos de aprendizado de máquina, tendo praticamente perdido seus próprios blocos gráficos. Mas, em um futuro próximo, eles podem ser superados de várias maneiras por processadores de IA especializados, incluindo um novo desenvolvimento da AWS, o chip Trainium.

No AWS Re: Invent, a empresa destacou seu progresso em aprendizado de máquina com suas instâncias P3dn (Nvidia V100) e P4 (Nvidia A100). A primeira variante estreou em 2018, quando o modelo BERT-Large era um exemplo de complexidade e, com 256 GB de memória e rede 100 GbE, apresentou resultados impressionantes. No entanto, a cada ano, a complexidade dos modelos de aprendizado de máquina aumenta em quase uma ordem de magnitude, e o crescimento dos recursos dos aceleradores de IA fica claramente aquém dessas taxas.

A complexidade dos modelos de aprendizado de máquina crescerá cada vez mais rápido

Quando a variante P4d foi introduzida no ano passado, ela quadruplicou seu poder de processamento e quadruplicou sua memória, enquanto o famoso modelo GPT-3 ultrapassou o BERT-Large em 500 vezes em complexidade. E agora 175 bilhões de parâmetros deste último não são nada comparados a 10 trilhões em novos modelos. Também temos que aumentar a quantidade de memória local (Trainium tem 512 GB de HBM com uma largura de banda total de 13,1 TB / s) e usar mais ativamente o aprendizado distribuído.

Para a última abordagem, o subsistema de rede tornou-se um gargalo e, ao desenvolver a pilha do Elastic Fabric Adapter (EFA), a empresa levou isso em consideração, dando às novas instâncias Trn1 uma conexão a 800 Gb / s (o dobro do P4d) e com latência ultrabaixa e uma versão mais otimizada do Trn1n, que tem o dobro da largura de banda e chega a 1,6 Tbit / s. Para comunicação entre os próprios chips dentro da instância, a interconexão NeuroLink é usada a uma velocidade de 768 GB / s.

Progresso do subsistema de rede e memória em instâncias AWS AI

Mas não se trata apenas de poder treinar o GPT-3 em menos de duas semanas: a quantidade de recursos usados ​​também é importante. No caso do P3d, isso exigiria 600 instâncias em execução simultaneamente, e até mesmo a mudança para a arquitetura Ampere reduziria esse número para 200. Mas o treinamento baseado em chips Trainium requer apenas 130 instâncias Trn1. Graças às otimizações, os custos de comunicação para novas instâncias são de apenas 7% contra 14% para Ampere e incríveis 49% para Volta.

Menos instâncias, maior eficiência com igual tempo de treinamento – isso é o que o Trainium dará

Trainium depende de uma matriz sistólica (o Google usou a mesma abordagem para suas TPUs), ou seja, consiste em muitas unidades computacionais intimamente relacionadas que processam independentemente os dados recebidos dos vizinhos e transmitem o resultado ao próximo vizinho. Esta abordagem, em particular, elimina vários registros e acessos à memória, o que é típico para GPUs “clássicas”, mas priva esses aceleradores de flexibilidade.

No Trainium, de acordo com a AWS, a flexibilidade é preservada – o acelerador possui 16 manipuladores totalmente programáveis ​​(em C / C ++). Ele também tem outras otimizações. Por exemplo, a aceleração de hardware de arredondamento estocástico, que em modelos ultra-grandes se torna muito “cara” devido à sobrecarga, embora possa melhorar a eficiência do treinamento com precisão mista. Tudo isso permite que você obtenha até 3,4 Pflops em cálculos de baixa precisão e até 840 Teraflops em cálculos FP32.

A AWS tentou fazer a transição para o Trainium o mais indolor possível para os desenvolvedores, pois o SDK do AWS Neuron oferece suporte a estruturas populares de aprendizado de máquina. No entanto, a empresa não vai forçar os clientes a instâncias Trn1 e continuará a fornecer outros aceleradores para escolher, uma vez que a transição, por exemplo, do ecossistema CUDA pode ser difícil. No entanto, em questões de aprendizado de máquina para suas próprias necessidades, a Amazon agora é completamente independente – ela tem uma CPU moderna Graviton3 e o acelerador de inferência Inferentia.

avalanche

Postagens recentes

Intel pensa em vender terrenos nos EUA para economizar dinheiro

A Intel Corporation está atualmente reestruturando seus ativos e tentando encontrar novos proprietários para alguns…

1 hora atrás

O robô humanóide Figure AI aprendeu a executar tarefas 4 vezes mais rápido e 7 vezes mais preciso

No início deste ano, a Figura AI apresentou seu robô humanóide de segunda geração, a…

2 horas atrás

Um concorrente da SpaceX aparecerá na Europa para entregar carga à ISS

A empresa europeia The Exploration Company poderá no futuro tornar-se concorrente da americana SpaceX no…

7 horas atrás

A Rostec importou um dos importantes elementos da eletrônica, rádios e navegadores

A holding Ruselectronics, que faz parte da estatal Rostec, criou uma linha de osciladores de…

8 horas atrás

O desenvolvimento de novos recursos da Apple para iOS estagnou

O jornalista da Bloomberg, Mark Gurman, relatou que a Apple adiou o lançamento de um…

9 horas atrás