Nos últimos anos, a complexidade dos modelos de IA dobrou em média a cada dois meses e, até agora, essa tendência continua. Há apenas três anos, o Google treinou um modelo “modesto” de BERT com 340 milhões de parâmetros em 9 dias Pflop. Em 2020, demorou cerca de 25-30 mil dias Pflop para treinar o modelo Micrsofot MSFT-1T com 1 trilhão de parâmetros. Os processadores e GPUs de uso geral têm cada vez mais dificuldade em lidar com essas tarefas, por isso várias empresas estão envolvidas no desenvolvimento de aceleradores especializados: Google, Groq, Graphcore, SambaNova, Enflame, etc.
A Cerebras se destacará por ter escolhido uma forma especial de escalar o poder de computação. Em vez de imprimir dezenas de chips em um grande wafer de silício, cortando-os do wafer e depois conectando-os, a empresa desenvolveu em 2019 um chip gigante Wafer-Scale Engine 1 (WSE-1) que ocupa quase todo o wafer . 400 mil núcleos feitos com uma tecnologia de processo de 16 nm consomem 15 kW, mas em uma série de tarefas eles acabam sendo centenas de vezes mais rápidos do que um supercomputador de 450 kW baseado em aceleradores NVIDIA.
Neste ano, a empresa lançou a segunda geração desses chips – WSE-2, em que, graças à transição para um processo técnico de 7 nm, foi possível aumentar o número de núcleos tensores para 850 mil, e trazer o L2 volume de cache de 40 GB, o que é cerca de 1000 vezes mais do que qualquer GPU. Naturalmente, essa abordagem de produção diminui o rendimento de wafers adequados e aumenta drasticamente o custo dos produtos, mas a Cerebras, em cooperação com a TSMC, conseguiu reduzir parcialmente a gravidade desse problema devido à redundância inerente ao projeto WSE.
Devido à identidade de todos os núcleos, mesmo que alguns deles falhem, o produto como um todo permanece operacional. No entanto, o custo de um wafer de 7 nm de 300 mm é de vários milhares de dólares, enquanto o custo do chip WSE é estimado em US $ 2 milhões, uma ordem de magnitude maior do que as GPUs mais produtivas. Um dos motivos para essa diferença é a grande quantidade de memória on-board rápida e a velocidade dos dados entre os núcleos.
No entanto, nem todo modelo agora é capaz de caber em um único chip WSE, então, de acordo com o CEO da Cerebras, Andrew Feldman, o foco da empresa agora é construir sistemas eficientes compostos de vários chips WSE. A taxa de crescimento da complexidade dos modelos supera a possibilidade de aumentar o poder de computação adicionando novos núcleos e memória ao prato, pois isso leva a um aumento excessivo no custo de um sistema já caro.
Os engenheiros da empresa veem a desagregação como a única maneira de fornecer o nível necessário de desempenho e escalabilidade. Esta abordagem implica a separação de memória e unidades computacionais para poder escaloná-las independentemente umas das outras – os parâmetros do modelo são colocados em um armazenamento separado e o próprio modelo pode ser espalhado por vários nós computacionais CS, unidos em um cluster .
No Hot Chips 33, a empresa introduziu um armazenamento especial chamado MemoryX, combinando DRAM e memória flash com uma capacidade total de 2,4 PB, que pode armazenar até 120 trilhões de parâmetros. Isso, segundo a empresa, permite construir modelos próximos à escala do cérebro humano, que possui cerca de 80 bilhões de neurônios e 100 trilhões. conexões entre eles. A propósito, Kioxia também está desenvolvendo memória flash do tamanho de uma placa inteira de 300 mm.
Para garantir a escalabilidade nos níveis de cluster WSE e CS, a Cerebras desenvolveu a tecnologia Weight Streaming. Com ele, a camada de ativação para modelos supergrandes (que em breve se tornará a norma) pode ser armazenada no WSE, e o fluxo de parâmetros vem de fora. A desagregação dos parâmetros de computação e armazenamento elimina os gargalos de latência e largura de banda de memória enfrentados por grandes clusters de processadores.
Isso abre amplas possibilidades para o dimensionamento independente do tamanho e velocidade do cluster, permitindo que você armazene trilhões de pesos WSE-2 no MemoryX e use de 1 a 192 CS-2 sem alterar o software. Em sistemas tradicionais, à medida que mais nós computacionais são adicionados ao cluster, cada um deles contribui cada vez menos para a solução do problema. A Cerebras desenvolveu a interconexão SwarmX, que permite conectar até 163 milhões de núcleos enquanto mantém ganhos lineares de desempenho.
Além disso, a empresa atentou para a esparsidade, ou seja, a exclusão de uma parte dos pesos que são insignificantes para o resultado final. Estudos têm mostrado que otimizar adequadamente o modelo pode alcançar um aumento de 10 vezes no desempenho, mantendo a precisão dos cálculos. O CS-2 apresenta Selectable Sparsity, uma tecnologia de dispersão dinâmica que permite aos usuários selecionar o nível desejado de redução do modelo para reduzir o tempo computacional.
«Grandes redes como GPT-3 já remodelaram a indústria de processamento de linguagem natural, tornando possível o que antes era impossível de imaginar. A indústria mudou para modelos com 1 trilhão de parâmetros e estamos expandindo esse limite em duas ordens de magnitude, criando redes neurais com 120 trilhões de parâmetros, comparáveis em escala ao cérebro ”, disse Feldman.