Tesla quer levar silício de aprendizado de máquina para o Dojo

Para saciar a sede de modelos de IA e aprendizado de máquina cada vez maiores, a Tesla revelou uma riqueza de detalhes no Hot Chips 34 em sua arquitetura de supercomputação totalmente personalizada chamada Dojo.

O sistema é essencialmente um supercomputador combinável massivo, embora diferente do que vemos no 500 topo, ele é construído a partir de uma arquitetura totalmente personalizada que abrange computação, rede e silício de entrada/saída (E/S) para arquitetura de conjunto de instruções (ISA), fornecimento de energia, empacotamento e resfriamento. Tudo isso foi feito com o propósito expresso de executar algoritmos de treinamento de aprendizado de máquina específicos e personalizados em escala.

“O processamento de dados do mundo real só é viável por meio de técnicas de aprendizado de máquina, seja processamento de linguagem natural, condução em ruas que são feitas para a visão humana para a interface robótica com o ambiente cotidiano”, disse Ganesh Venkataramanan, diretor sênior de engenharia de hardware da Tesla. durante seu discurso de abertura.

No entanto, ele argumentou que os métodos tradicionais para dimensionar cargas de trabalho distribuídas não conseguiram acelerar na taxa necessária para acompanhar as demandas de aprendizado de máquina. Com efeito, a Lei de Moore não é suficiente e nem os sistemas disponíveis para treinamento de IA/ML em escala, ou seja, alguma combinação de CPU/GPU ou, em circunstâncias mais raras, usando aceleradores de IA especializados. 

“Tradicionalmente construímos chips, os colocamos em pacotes, os pacotes vão em PCBs, que vão para os sistemas. Os sistemas vão para racks”, disse Venkataramanan. O problema é que cada vez que os dados se movem do chip para o pacote e para fora do pacote, há uma penalidade de latência e largura de banda.

Um sanduíche de datacenter

Então, para contornar as limitações, Venkataramanan e sua equipe começaram do zero.

“Desde a minha entrevista com Elon, ele me perguntou o que você pode fazer de diferente de CPUs e GPUs para IA. Sinto que toda a equipe ainda está respondendo a essa pergunta.”

Isso levou ao desenvolvimento do bloco de treinamento Dojo, um cluster de computação independente que ocupa meio pé cúbico capaz de 556 TFLOPS de desempenho FP32 em um pacote de 15 kW refrigerado a líquido.

Cada bloco está equipado com 11 GB de SRAM e é conectado em uma malha de 9 TB/s usando um protocolo de transporte personalizado em toda a pilha.

“Este bloco de treinamento representa uma quantidade incomparável de integração do computador à memória, ao fornecimento de energia e à comunicação, sem a necessidade de interruptores adicionais”, disse Venkataramanan.

No centro do bloco de treinamento está o D1 da Tesla, uma matriz de 50 bilhões de transistores, baseada no processo de 7 nm da TSMC. A Tesla diz que cada D1 é capaz de 22 TFLOPS de desempenho FP32 a um TDP de 400W. No entanto, Tesla observa que o chip é capaz de executar uma ampla variedade de cálculos de ponto flutuante, incluindo alguns personalizados.

“Se você comparar transistores por milímetro quadrado, isso provavelmente é o limite de qualquer coisa que esteja por aí”, disse Venkataramanan.

A Tesla então pegou 25 D1s, separou-os para matrizes boas conhecidas e, em seguida, empacotou-os usando a tecnologia system-on-wafer da TSMC para “alcançar uma enorme quantidade de integração de computação com latência muito baixa e largura de banda muito alta”, disse ele.

No entanto, o design do sistema em wafer e a arquitetura empilhada verticalmente apresentaram desafios no que diz respeito ao fornecimento de energia.

De acordo com Venkataramanan, a maioria dos aceleradores hoje coloca a energia diretamente adjacente ao silício. E, embora comprovada, essa abordagem significa que uma grande área do acelerador deve ser dedicada a esses componentes, o que a tornou impraticável para o Dojo, explicou. Em vez disso, a Tesla projetou seus chips para fornecer energia diretamente pela parte inferior do molde. 

Juntando tudo

“Poderíamos construir um datacenter inteiro ou um edifício inteiro com esse bloco de treinamento, mas o bloco de treinamento é apenas a parte de computação. Também precisamos alimentá-lo”, disse Venkataramanan.

Para isso, a Tesla também desenvolveu o Dojo Interface Processor (DIP), que funciona como uma ponte entre a CPU do host e os processadores de treinamento. O DIP também serve como uma fonte de memória compartilhada de alta largura de banda (HBM) e como uma NIC de alta velocidade de 400 Gbit/s.

Cada DIP possui 32 GB de HBM e até cinco dessas placas podem ser conectadas a um bloco de treinamento a 900 GB/s para um agregado de 4.5 TB/s ao host para um total de 160 GB de HBM por bloco.

Os pares de configuração V1 da Tesla desses blocos - ou 150 matrizes D1 - em matriz suportavam quatro CPUs host, cada uma equipada com cinco placas DIP para obter um desempenho exaflop reivindicado de BF16 ou CFP8.

Juntos, Venkataramanan diz que a arquitetura – detalhada em profundidade aqui by A próxima plataforma – permite que a Tesla supere as limitações associadas aos aceleradores tradicionais de empresas como Nvidia e AMD.

“Como os aceleradores tradicionais funcionam, normalmente você tenta encaixar um modelo inteiro em cada acelerador. Replicá-lo e, em seguida, transmitir os dados através de cada um deles”, disse ele. “O que acontece se tivermos modelos cada vez maiores? Esses aceleradores podem falhar porque ficam sem memória.”

Este não é um problema novo, ele observou. O switch NV da Nvidia, por exemplo, permite que a memória seja agrupada em grandes bancos de GPUs. No entanto, Venkataramanan argumenta que isso não apenas adiciona complexidade, mas introduz latência e comprometimentos na largura de banda.

“Pensamos nisso desde o início. Nossos blocos de computação e cada um dos moldes foram feitos para encaixar modelos grandes”, disse Venkataramanan.

Software

Essa arquitetura de computação especializada exige uma pilha de software especializada. No entanto, Venkataramanan e sua equipe reconheceram que a programabilidade faria ou quebraria o Dojo.

“Facilidade de programação para contrapartes de software é primordial quando projetamos esses sistemas”, disse ele. “Os pesquisadores não vão esperar que seu pessoal de software escreva um kernel manuscrito para se adaptar a um novo algoritmo que queremos executar.”

Para fazer isso, Tesla abandonou a ideia de usar kernels e projetou a arquitetura do Dojo em torno de compiladores.

“O que fizemos foi usar o PiTorch. Criamos uma camada intermediária, que nos ajuda a paralelizar para dimensionar o hardware abaixo dela. Por baixo de tudo está o código compilado”, disse ele. “Esta é a única maneira de criar pilhas de software adaptáveis ​​a todas essas cargas de trabalho futuras.”

Apesar da ênfase na flexibilidade do software, Venkataramanan observa que a plataforma, que está atualmente em execução em seus laboratórios, está limitada ao uso da Tesla por enquanto.

“Estamos focados em nossos clientes internos em primeiro lugar”, disse ele. “Elon tornou público que, com o tempo, vamos disponibilizar isso para os pesquisadores, mas não temos um prazo para isso. ®