Avanços em técnicas de machine learning para machine learning
(blog.research.google)- À medida que os modelos de ML crescem, a eficiência de execução passa a depender fortemente não só do hardware, mas também da otimização do compilador, e o Google Research e o Google DeepMind divulgaram uma abordagem para melhorar essa área novamente com ML
- As decisões de compilador tomadas sob frameworks como TensorFlow, JAX e PyTorch podem mudar bastante o tempo de execução e o uso de recursos mesmo para o mesmo modelo
- TpuGraphs é um conjunto de dados que reúne grafos computacionais de programas de ML para TPU, configurações de compilação e tempos de execução, podendo ser usado em pesquisas de modelos de custo baseados em aprendizado
- Para reduzir as limitações de treinar grafos grandes de uma só vez, o Graph Segment Training divide o grafo em segmentos, reduz o uso de memória e encurta o tempo de treinamento em 3 vezes
- Em uma competição do Kaggle, 792 participantes em 616 equipes de 66 países validaram técnicas práticas para melhorar modelos de previsão de custo, como compressão de grafos, ajuste de valores de padding, adição de características de nós e attention entre configurações
Por que compiladores de ML determinam o desempenho de execução
- Modelos modernos de ML executam tarefas como compreensão de linguagem natural, conversação, geração de imagens e geração de vídeo, sendo escritos e treinados com frameworks de programação de ML como TensorFlow, JAX e PyTorch
- Os frameworks fornecem operações de álgebra linear como multiplicação de matrizes e convolução, além de camadas de rede neural como camadas de convolução 2D e camadas Transformer
- O usuário não precisa lidar diretamente com os detalhes de como executar o modelo com eficiência no hardware, porque o compilador abaixo do framework otimiza o modelo automaticamente
- No entanto, como compiladores frequentemente resolvem problemas complexos de otimização com heurísticas, nem sempre conseguem entregar o melhor desempenho possível
Grafo computacional e otimização em duas etapas
- Compiladores de ML convertem instruções matemáticas escritas pelo usuário em instruções que podem ser executadas no hardware real
- Um programa de ML pode ser representado como um grafo computacional
- Os nós representam operações sobre tensores, como matrix multiplication
- As arestas representam tensores que fluem de um nó para outro
- As otimizações de compilador se dividem em dois grandes tipos
- Otimização em nível de grafo: toma decisões considerando o contexto do grafo inteiro e transforma o grafo como um todo
- Otimização em nível de kernel: transforma um único kernel, que é um subgrafo fundido, independentemente dos outros kernels
Trade-offs de desempenho no layout de memória
- Tensores 2D como matrizes podem ser armazenados na memória como
[A B C a b c]ou[A a B b C c], correspondendo respectivamente aos layouts row-major e column-major - Uma das otimizações importantes em compiladores de ML é atribuir layouts de memória a todos os tensores intermediários do programa
- Um layout específico pode ser o mais eficiente para uma operação individual, mas se os layouts entre add e convolution não coincidirem, o compilador precisará inserir uma operação de cópia adicional
- Por outro lado, mesmo que o desempenho de operações individuais seja um pouco menor, uma configuração que não exija conversão de layout pode ser melhor na execução total
- No conjunto de benchmarks do XLA, ao escolher a configuração ótima de layout em vez da configuração padrão do compilador, foi observado um ganho de velocidade de até 32%
Conjunto de dados TpuGraphs
- TpuGraphs é um conjunto de dados de modelo de custo baseado em aprendizado para programas executados nas TPUs customizadas do Google
- O objetivo é treinar um modelo de custo que receba o programa de entrada e a configuração do compilador para prever o tempo de execução do programa
- O conjunto de dados cobre duas configurações do compilador XLA
- layout: uma configuração que generaliza o conceito de row-major e column-major de matrizes para tensores de alta dimensão
- tiling: configuração de tamanho de tile
- Cada exemplo inclui o grafo computacional de uma carga de trabalho de ML, a configuração de compilação e o tempo de execução ao compilar com essa configuração
- Os grafos foram coletados de programas de ML de código aberto e incluem arquiteturas de modelo como ResNet, EfficientNet, Mask R-CNN e Transformer
- O método de download e o código inicial estão disponíveis no GitHub do TpuGraphs
- O TpuGraphs tem 25 vezes mais grafos do que o maior conjunto de dados anterior de previsão de propriedades de grafos com tamanho de grafo semelhante, e seu tamanho médio de grafo é 770 vezes maior do que o de conjuntos de dados anteriores de previsão de desempenho de programas de ML
Modelo de custo de referência e arquitetura GNN
- O TpuGraphs fornece junto um modelo de custo de referência baseado em aprendizado e, como o programa de entrada é representado como grafo, usa GNN
- As características dos nós são compostas por duas partes
- opcode id: a informação de nó mais importante, que indica o tipo de operação de tensor
- outras características de nó
- O modelo de referência converte o opcode id em opcode embedding por meio de uma embedding lookup table
- O opcode embedding é combinado com o restante das características dos nós e usado como entrada para a GNN
- Os node embeddings produzidos pela GNN são combinados em um graph embedding de tamanho fixo por meio de reduções simples de graph pooling, como sum e mean
- O graph embedding final é convertido em uma única saída escalar ao passar por uma feedforward layer
Treinando grafos grandes com Graph Segment Training
- Graph Segment Training é uma técnica de escalonamento de treinamento de GNN para lidar com grafos grandes em dispositivos com capacidade de memória limitada
- Esse método é voltado para cenários de graph-level prediction, em que o alvo da previsão não é um nó ou uma aresta, mas o grafo inteiro
- Grafos computacionais podem incluir centenas de milhares de nós, então o Full Graph Training, que usa o grafo inteiro de uma vez, pode ser computacionalmente inviável
- O GST divide o grafo grande em segmentos menores e atualiza o modelo escolhendo apenas um subconjunto aleatório dos segmentos
- Os demais segmentos geram embeddings sem armazenar activations intermediárias, reduzindo assim o uso de memória
- Todos os segment embeddings são combinados para formar o embedding do grafo grande original, que então é usado na previsão
- Uma historical embedding table e segment dropout são introduzidos em conjunto para aliviar a defasagem dos historical embeddings
- O método completo reduz o tempo de treinamento end-to-end em 3 vezes
Técnicas de melhoria validadas na competição do Kaggle
- A competição do Kaggle Fast or Slow? Predict AI Model Runtime foi realizada com base no conjunto de dados TpuGraphs, e contou com 792 participantes em 616 equipes de 66 países
- Houve 10.507 submissões, e para 153 pessoas essa foi a primeira competição no Kaggle; entre elas, 47 ficaram entre as 100 melhores
- As equipes participantes experimentaram várias técnicas
- Pruning e compressão de grafos: métodos de compressão de grafos grandes foram testados no lugar do GST, incluindo uma abordagem que mantém apenas subgrafos contendo nós configuráveis e seus vizinhos imediatos
- Mudança do valor de padding: como o valor padrão de padding 0 conflita com valores válidos de features, usar -1 pôde melhorar significativamente a precisão do modelo
- Adição de características de nós e mudança de codificação: características extras de nós, como as contracting dimensions de dot general, são importantes, e a forma de codificar as características dos nós também pode afetar os resultados
- Cross-configuration attention: a equipe vencedora projetou uma camada simples que permite ao modelo comparar explicitamente as configurações, obtendo resultados muito melhores do que fazer inferência para cada configuração de forma individual
- Os resultados da competição e a solução vencedora estavam programados para ser apresentados em 16 de dezembro de 2023, na sessão da competição do workshop ML for Systems da NeurIPS
Sessão relacionada à NeurIPS Expo
- Para leitores interessados em dados estruturados e pesquisa em inteligência artificial, o painel da NeurIPS Expo Graph Learning Meets Artificial Intelligence ocorreu em 9 de dezembro de 2023
- Esse painel aborda temas como os avanços em modelos de custo baseados em aprendizado
1 comentários
Comentários no Hacker News
Compiladores de ML estão sendo superestimados. É o mesmo tipo de trade-off dos compiladores tradicionais: você ganha muito mais throughput do que contratando programadores especializados em performance, mas estes últimos normalmente são muito mais rápidos e, em alguns casos, podem ficar à frente por ordens de grandeza
Falham em vários níveis. No nível algorítmico, não conseguem retroalimentar para humanos os truques para tornar a rede mais rápida; dão apenas sinais muito básicos. A intenção também se perde. Um projetista de redes de ML especifica a estrutura em Python, mas depois de várias etapas de lowering a coisa pode virar algo completamente diferente. Recentemente vi um compilador que, ao fazer uma atualização de slice, criava todos os intervalos de índices possíveis de um array, recortava isso para obter os índices a atualizar e então fazia um scatter; substituí por uma única chamada
memcpy. Os kernels também são ineficientes. Sempre que a saída desses compiladores enfrenta um programador de assembly experiente, o compilador perde, geralmente por 30% ou mais. Parece algo fácil de resolver, mas se ninguém conseguiu resolver direito nos últimos 50 anos, claramente não é tão simples quanto pareceCompiladores também têm muitas heurísticas, como inline, desenrolamento de loops e vetorização, então redes neurais podem ajudar, e talvez sejam mais fáceis de manter do que inúmeras heurísticas escritas por humanos
Esse tipo de otimização ainda é muito melhor do que não ter otimização nenhuma
Alguém pode explicar isso de forma mais realista? Tenho curiosidade sobre o estado real atual dos compiladores de ML e o que dá para esperar no curto prazo
Basta escrever
model = torch.compile(model). “Nesses 163 modelos open source, o torch.compile funcionou em 93% dos casos e acelerou o treinamento em 43% em GPUs NVIDIA A100. Na precisão Float32, a aceleração média foi de 21%; na precisão AMP, foi de 51%.”[1] O Google parece querer envolver mais pessoas em P&D desse tipo de método[1] https://pytorch.org/get-started/pytorch-2.0/
A realidade é complexa, então simplificando bastante: a ideia é compilar o grafo computacional para alguma representação intermediária e implementar backends para ela. Projetos relacionados incluem stableHLO, IREE e openXLA. O compilador
jitdo Jax também pode ser visto como uma forma desse tipo de compilador. Ele baixa as operações rastreadas para XLA, e o XLA faz várias mágicas para que isso rode novamente nos backends. No fim, quanto mais você desce, mais é tudo transformação e abstraçãotorch.compileEm resumo, é um trabalho para melhorar a previsão do desempenho em tempo de execução de grafos computacionais usando redes neurais de grafos (GNNs). Ele usa um dicionário de embeddings para o opcode de cada nó, junto com outras características dos nós, como shape, bits e window size ([1])
Um grande conjunto de dados de grafos com diferentes configurações de compilação XLA e o desempenho resultante em TPUs foi publicado em [2], e em [3] eles melhoraram a previsão em grafos maiores do que antes usando uma forma de particionar o grafo (foi a primeira vez que vi METIS graph partition) e várias técnicas de treinamento. Isso trata de prever o desempenho de um grafo dado, não de melhorar, sugerir ou modificar novos grafos equivalentes. Como no FunSearch, um modelo com boa capacidade preditiva pode ser usado junto com busca evolutiva
[1] https://github.com/google-research-datasets/tpu_graphs#featu...
[2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
[3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322
Alguém consegue explicar como a convolução funciona naquele grafo? Foi feita uma convolução de um tensor de shape
[2,4,16]com um kernel de shape[4,16,8], e saiu um tensor[2,8]; como isso é possível?[2,4,16], o2pode ser visto como o tamanho do batch, o4como a dimensão de características de entrada e o16como a dimensão de canais de entradaNo kernel
[4,16,8], o4é o tamanho da janela do filtro, o16corresponde à dimensão dos canais de entrada, e o8é a dimensão dos canais de saída. Na saída[2,8], o2é mantido como tamanho do batch, e o8corresponde à dimensão dos canais de saída do kernel. À primeira vista, as dimensões parecem incompatíveis, mas a convolução em grafos usa a estrutura de vizinhança. O kernel se move sobre o grafo, aplica pesos ao nó atual e às características dos vizinhos dentro de um certo raio, e reúne essa soma ponderada para criar novas características para cada canal de saída. A estrutura do grafo, os pesos das arestas e detalhes de implementação como padding e stride também podem afetar o shape de saídaEm que pé está o Gemini?
Só de cabeça já há pelo menos cinco modelos-base como Llama, Claude, Gemini, Falcon e Mistral, todos disputando posições, mas o GPT ainda está um degrau acima, e isso já faz um ano. Ficou claro que grandes modelos de linguagem baseados em Transformer são simples o bastante para que qualquer um possa criá-los se puder gastar algo como um milhão de dólares em tempo de GPU, mas ainda assim ninguém alcançou totalmente a OpenAI. Qual será o segredo especial deles?
E quanto ao próprio Transformer? Há algum indício de que ele seja ótimo em algum sentido?
Sinto que o ponto principal ficou enterrado no primeiro parágrafo, mas o restante é ótimo
O ritmo de avanço de ML agora é impressionante. Não acredito na singularidade, mas ela está mudando software e sociedade de formas que ninguém consegue prever
https://deepmind.google/discover/blog/funsearch-making-new-d...