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.