GPUs ocupadas: o método de amostragem e pipelining acelera o aprendizado profundo em gráficos grandes

Grafos, uma rede potencialmente extensa de nós conectados por arestas, podem ser usados ​​para expressar e interrogar relacionamentos entre dados, como conexões sociais, transações financeiras, tráfego, redes de energia e interações moleculares. À medida que os pesquisadores coletam mais dados e constroem essas imagens gráficas, os pesquisadores precisarão de métodos mais rápidos e eficientes, bem como de mais poder computacional, para conduzir o aprendizado profundo sobre eles, na forma de redes neurais de grafo (GNN).  

Agora, um novo método, chamado SALIENT (SAmpling, sLIcing e data movemeNT), desenvolvido por pesquisadores do MIT e da IBM Research, melhora o treinamento e o desempenho de inferência abordando três principais gargalos na computação. Isso reduz drasticamente o tempo de execução de GNNs em grandes conjuntos de dados, que, por exemplo, contêm na escala de 100 milhões de nós e 1 bilhão de arestas. Além disso, a equipe descobriu que a técnica escala bem quando o poder computacional é adicionado de uma a 16 unidades de processamento gráfico (GPUs). O trabalho foi apresentado na Quinta Conferência sobre Aprendizado de Máquina e Sistemas.

“Começamos a observar os desafios que os sistemas atuais enfrentam ao dimensionar técnicas de aprendizado de máquina de última geração para gráficos em conjuntos de dados realmente grandes. Descobriu-se que havia muito trabalho a ser feito, porque muitos dos sistemas existentes estavam obtendo bom desempenho principalmente em conjuntos de dados menores que cabem na memória da GPU”, diz Tim Kaler, autor principal e pós-doutorando no MIT Computer Science. e Laboratório de Inteligência Artificial (CSAIL).

Por vastos conjuntos de dados, os especialistas querem dizer escalas como toda a rede Bitcoin, onde certos padrões e relacionamentos de dados podem indicar tendências ou jogo sujo. “Existem quase um bilhão de transações de Bitcoin na blockchain e, se quisermos identificar atividades ilícitas dentro de uma rede tão conjunta, estamos diante de um gráfico dessa escala”, diz o coautor Jie Chen, cientista sênior de pesquisa e gerente. da IBM Research e do MIT-IBM Watson AI Lab. “Queremos construir um sistema que seja capaz de lidar com esse tipo de gráfico e que permita que o processamento seja o mais eficiente possível, porque a cada dia queremos acompanhar o ritmo dos novos dados que são gerados.”

Os co-autores de Kaler e Chen incluem Nickolas Stathas MEng '21 da Jump Trading, que desenvolveu SALIENT como parte de seu trabalho de pós-graduação; a ex-estagiária do MIT-IBM Watson AI Lab e estudante de pós-graduação do MIT Anne Ouyang; Pós-doutorado MIT CSAIL Alexandros-Stavros Iliopoulos; Cientista Pesquisador do MIT CSAIL Tao B. Schardl; e Charles E. Leiserson, professor Edwin Sibley Webster de Engenharia Elétrica no MIT e pesquisador do MIT-IBM Watson AI Lab.     

Para esse problema, a equipe adotou uma abordagem orientada a sistemas no desenvolvimento de seu método: SALIENT, diz Kaler. Para fazer isso, os pesquisadores implementaram o que consideraram importantes otimizações básicas de componentes que se encaixam nas estruturas de aprendizado de máquina existentes, como o PyTorch Geometric e a biblioteca de gráficos profundos (DGL), que são interfaces para a construção de um modelo de aprendizado de máquina. Stathas diz que o processo é como trocar os motores para construir um carro mais rápido. Seu método foi projetado para se adequar às arquiteturas GNN existentes, para que os especialistas de domínio pudessem aplicar facilmente esse trabalho a seus campos específicos para agilizar o treinamento do modelo e obter insights durante a inferência mais rapidamente. O truque, determinado pela equipe, era manter todo o hardware (CPUs, links de dados e GPUs) ocupado o tempo todo: enquanto a CPU faz amostras do gráfico e prepara mini-lotes de dados que serão transferidos pelo link de dados , a GPU mais crítica está trabalhando para treinar o modelo de aprendizado de máquina ou conduzir a inferência. 

Os pesquisadores começaram analisando o desempenho de uma biblioteca de aprendizado de máquina comumente usada para GNNs (PyTorch Geometric), que mostrou uma utilização surpreendentemente baixa dos recursos de GPU disponíveis. Aplicando otimizações simples, os pesquisadores melhoraram a utilização da GPU de 10% para 30%, resultando em uma melhoria de desempenho de 1.4 a duas vezes em relação aos códigos de referência públicos. Esse código de linha de base rápido pode executar uma passagem completa em um grande conjunto de dados de treinamento por meio do algoritmo (uma época) em 50.4 segundos.                          

Buscando mais melhorias de desempenho, os pesquisadores se propuseram a examinar os gargalos que ocorrem no início do pipeline de dados: os algoritmos para amostragem de gráficos e preparação de mini-batch. Ao contrário de outras redes neurais, as GNNs realizam uma operação de agregação de vizinhança, que calcula informações sobre um nó usando informações presentes em outros nós próximos no grafo — por exemplo, em um grafo de rede social, informações de amigos de amigos de um usuário. À medida que o número de camadas no GNN aumenta, o número de nós que a rede precisa alcançar para obter informações pode explodir, excedendo os limites de um computador. Os algoritmos de amostragem de vizinhança ajudam selecionando um subconjunto aleatório menor de nós a serem coletados; no entanto, os pesquisadores descobriram que as implementações atuais disso eram muito lentas para acompanhar a velocidade de processamento das GPUs modernas. Em resposta, eles identificaram uma combinação de estruturas de dados, otimizações algorítmicas e assim por diante que melhoraram a velocidade de amostragem, melhorando a operação de amostragem sozinha em cerca de três vezes, levando o tempo de execução por período de 50.4 para 34.6 segundos. Eles também descobriram que a amostragem, em uma taxa apropriada, pode ser feita durante a inferência, melhorando a eficiência e o desempenho geral da energia, um ponto que foi negligenciado na literatura, observa a equipe.      

Em sistemas anteriores, essa etapa de amostragem era uma abordagem multiprocesso, criando dados extras e movimentação desnecessária de dados entre os processos. Os pesquisadores tornaram seu método SALIENT mais ágil, criando um único processo com threads leves que mantinham os dados na CPU na memória compartilhada. Além disso, o SALIENT aproveita um cache de processadores modernos, diz Stathas, paralelizando o fatiamento de recursos, que extrai informações relevantes de nós de interesse e seus vizinhos e arestas circundantes, dentro da memória compartilhada do cache central da CPU. Isso novamente reduziu o tempo de execução geral por época de 34.6 para 27.8 segundos.

O último gargalo que os pesquisadores abordaram foi canalizar as transferências de dados em minilotes entre a CPU e a GPU usando uma etapa de pré-busca, que prepararia os dados antes de serem necessários. A equipe calculou que isso maximizaria o uso da largura de banda no link de dados e levaria o método a uma utilização perfeita; no entanto, eles viram apenas cerca de 90 por cento. Eles identificaram e corrigiram um bug de desempenho em uma biblioteca popular do PyTorch que causava comunicações de ida e volta desnecessárias entre a CPU e a GPU. Com esse bug corrigido, a equipe alcançou um tempo de execução de 16.5 segundos por época com o SALIENT.

“Acho que nosso trabalho mostrou que o diabo mora nos detalhes”, diz Kaler. “Quando você presta muita atenção aos detalhes que afetam o desempenho ao treinar uma rede neural de gráfico, pode resolver um grande número de problemas de desempenho. Com nossas soluções, acabamos ficando completamente engasgados na computação da GPU, que é o objetivo ideal de um sistema desse tipo.”

A velocidade do SALIENT foi avaliada em três conjuntos de dados padrão ogbn-arxiv, ogbn-products e ogbn-papers100M, bem como em configurações de várias máquinas, com diferentes níveis de fanout (quantidade de dados que a CPU prepararia para a GPU) e em várias arquiteturas, incluindo a mais recente, GraphSAGE-RI. Em cada configuração, o SALIENT superou o PyTorch Geometric, principalmente no grande conjunto de dados ogbn-papers100M, contendo 100 milhões de nós e mais de um bilhão de arestas. Aqui, foi três vezes mais rápido, rodando em uma GPU, do que a linha de base otimizada que foi originalmente criada para Este trabalho; com 16 GPUs, o SALIENT foi oito vezes mais rápido. 

Enquanto outros sistemas tinham hardware ligeiramente diferente e configurações experimentais, nem sempre era uma comparação direta, o SALIENT ainda os superava. Entre os sistemas que alcançaram precisão semelhante, os números de desempenho representativos incluem 99 segundos usando uma GPU e 32 CPUs e 13 segundos usando 1,536 CPUs. Em contraste, o tempo de execução do SALIENT usando uma GPU e 20 CPUs foi de 16.5 segundos e foi de apenas dois segundos com 16 GPUs e 320 CPUs. “Se você olhar para os números finais que o trabalho anterior relata, nosso tempo de execução de 16 GPUs (dois segundos) é uma ordem de grandeza mais rápida do que outros números relatados anteriormente neste conjunto de dados”, diz Kaler. Os pesquisadores atribuíram suas melhorias de desempenho, em parte, à abordagem de otimizar seu código para uma única máquina antes de passar para a configuração distribuída. Stathas diz que a lição aqui é que, pelo seu dinheiro, “faz mais sentido usar o hardware que você tem de forma eficiente e ao extremo, antes de começar a escalar para vários computadores”, o que pode proporcionar economias significativas em custos e emissões de carbono que pode vir com o treinamento do modelo.

Essa nova capacidade agora permitirá que os pesquisadores abordem e se aprofundem em gráficos cada vez maiores. Por exemplo, a rede Bitcoin mencionada anteriormente continha 100,000 nós; o sistema SALIENT pode lidar com um gráfico 1,000 vezes (ou três ordens de magnitude) maior.

“No futuro, estaríamos olhando não apenas para executar este sistema de treinamento de rede neural de gráfico nos algoritmos existentes que implementamos para classificar ou prever as propriedades de cada nó, mas também queremos fazer tarefas mais aprofundadas, como identificar padrões comuns em um gráfico (padrões de subgráficos), [que] podem ser realmente interessantes para indicar crimes financeiros”, diz Chen. “Também queremos identificar nós em um gráfico que sejam semelhantes no sentido de que possivelmente correspondam ao mesmo malfeitor em um crime financeiro. Essas tarefas exigiriam o desenvolvimento de algoritmos adicionais e, possivelmente, também arquiteturas de redes neurais”.

Esta pesquisa foi apoiada pelo MIT-IBM Watson AI Lab e em parte pelo Laboratório de Pesquisa da Força Aérea dos EUA e pelo Acelerador de Inteligência Artificial da Força Aérea dos EUA.