GPU ocupadas: el método de muestreo y canalización acelera el aprendizaje profundo en gráficos grandes

Los gráficos, una red potencialmente extensa de nodos conectados por bordes, se pueden usar para expresar e interrogar relaciones entre datos, como conexiones sociales, transacciones financieras, tráfico, redes de energía e interacciones moleculares. A medida que los investigadores recopilen más datos y construyan estas imágenes gráficas, necesitarán métodos más rápidos y eficientes, así como más poder computacional, para llevar a cabo un aprendizaje profundo en ellos, en forma de redes neuronales gráficas (GNN).  

Ahora, un nuevo método, llamado SALIENT (SAmpling, sLIcing, and data movemeNT), desarrollado por investigadores del MIT e IBM Research, mejora el rendimiento del entrenamiento y la inferencia al abordar tres cuellos de botella clave en la computación. Esto reduce drásticamente el tiempo de ejecución de GNN en grandes conjuntos de datos que, por ejemplo, contienen una escala de 100 millones de nodos y 1 millones de bordes. Además, el equipo descubrió que la técnica escala bien cuando se agrega potencia computacional de una a 16 unidades de procesamiento gráfico (GPU). El trabajo fue presentado en las V Jornadas de Machine Learning y Sistemas.

“Comenzamos a observar los desafíos que experimentaban los sistemas actuales al escalar técnicas de aprendizaje automático de última generación para gráficos a conjuntos de datos realmente grandes. Resultó que había mucho trabajo por hacer, porque muchos de los sistemas existentes estaban logrando un buen rendimiento principalmente en conjuntos de datos más pequeños que caben en la memoria de la GPU”, dice Tim Kaler, autor principal y posdoctorado en Ciencias de la Computación del MIT. y Laboratorio de Inteligencia Artificial (CSAIL).

Por vastos conjuntos de datos, los expertos se refieren a escalas como toda la red de Bitcoin, donde ciertos patrones y relaciones de datos podrían explicar tendencias o juegos sucios. "Hay casi mil millones de transacciones de Bitcoin en la cadena de bloques, y si queremos identificar actividades ilícitas dentro de una red conjunta de este tipo, entonces nos enfrentamos a un gráfico de tal escala", dice el coautor Jie Chen, científico investigador principal y gerente. de IBM Research y el MIT-IBM Watson AI Lab. “Queremos construir un sistema que sea capaz de manejar ese tipo de gráfico y permita que el procesamiento sea lo más eficiente posible, porque todos los días queremos seguir el ritmo de los nuevos datos que se generan”.

Los coautores de Kaler y Chen incluyen a Nickolas Stathas MEng '21 de Jump Trading, quien desarrolló SALIENT como parte de su trabajo de posgrado; la ex becaria del MIT-IBM Watson AI Lab y estudiante de posgrado del MIT Anne Ouyang; el postdoctorado MIT CSAIL Alexandros-Stavros Iliopoulos; Tao B. Schardl, científico investigador de MIT CSAIL; y Charles E. Leiserson, profesor de ingeniería eléctrica Edwin Sibley Webster en el MIT e investigador del MIT-IBM Watson AI Lab.     

Para este problema, el equipo adoptó un enfoque orientado a los sistemas al desarrollar su método: SALIENT, dice Kaler. Para hacer esto, los investigadores implementaron lo que consideraron optimizaciones básicas importantes de los componentes que encajan en los marcos de aprendizaje automático existentes, como PyTorch Geometric y la biblioteca de gráficos profundos (DGL), que son interfaces para construir un modelo de aprendizaje automático. Stathas dice que el proceso es como cambiar motores para construir un auto más rápido. Su método fue diseñado para adaptarse a las arquitecturas GNN existentes, de modo que los expertos de dominio pudieran aplicar fácilmente este trabajo a sus campos específicos para acelerar el entrenamiento del modelo y obtener información durante la inferencia más rápido. El equipo determinó que el truco consistía en mantener ocupado todo el hardware (CPU, enlaces de datos y GPU) en todo momento: mientras la CPU toma muestras del gráfico y prepara minilotes de datos que luego se transferirán a través del enlace de datos. , la GPU más crítica está trabajando para entrenar el modelo de aprendizaje automático o realizar inferencias. 

Los investigadores comenzaron analizando el rendimiento de una biblioteca de aprendizaje automático de uso común para GNN (PyTorch Geometric), que mostró una utilización sorprendentemente baja de los recursos de GPU disponibles. Aplicando optimizaciones simples, los investigadores mejoraron la utilización de la GPU del 10 al 30 por ciento, lo que resultó en una mejora del rendimiento de 1.4 a dos veces en relación con los códigos de referencia públicos. Este código de referencia rápido podría ejecutar una pasada completa sobre un gran conjunto de datos de entrenamiento a través del algoritmo (una época) en 50.4 segundos.                          

Buscando mejoras adicionales en el rendimiento, los investigadores se dispusieron a examinar los cuellos de botella que ocurren al comienzo de la canalización de datos: los algoritmos para el muestreo de gráficos y la preparación de mini lotes. A diferencia de otras redes neuronales, las GNN realizan una operación de agregación de vecindario, que calcula información sobre un nodo utilizando información presente en otros nodos cercanos en el gráfico; por ejemplo, en un gráfico de red social, información de amigos de amigos de un usuario. A medida que aumenta la cantidad de capas en la GNN, la cantidad de nodos a los que la red tiene que llegar para obtener información puede explotar, excediendo los límites de una computadora. Los algoritmos de muestreo de vecindario ayudan al seleccionar un subconjunto aleatorio más pequeño de nodos para recopilar; sin embargo, los investigadores descubrieron que las implementaciones actuales de esto eran demasiado lentas para mantenerse al día con la velocidad de procesamiento de las GPU modernas. En respuesta, identificaron una combinación de estructuras de datos, optimizaciones algorítmicas, etc., que mejoraron la velocidad de muestreo y, en última instancia, mejoraron la operación de muestreo solo unas tres veces, lo que llevó el tiempo de ejecución por época de 50.4 a 34.6 segundos. También encontraron que el muestreo, a un ritmo apropiado, se puede realizar durante la inferencia, mejorando la eficiencia y el rendimiento energético general, un punto que se había pasado por alto en la literatura, señala el equipo.      

En sistemas anteriores, este paso de muestreo era un enfoque de múltiples procesos, creando datos adicionales y movimiento de datos innecesario entre los procesos. Los investigadores hicieron que su método SALIENT fuera más ágil al crear un proceso único con subprocesos livianos que mantuvieron los datos en la CPU en la memoria compartida. Además, SALIENT aprovecha un caché de procesadores modernos, dice Stathas, paralelizando el corte de funciones, que extrae información relevante de los nodos de interés y sus vecinos y bordes circundantes, dentro de la memoria compartida del caché central de la CPU. Esto nuevamente redujo el tiempo de ejecución general por época de 34.6 a 27.8 segundos.

El último cuello de botella que abordaron los investigadores fue canalizar las transferencias de datos en minilotes entre la CPU y la GPU mediante un paso de obtención previa, que prepararía los datos justo antes de que se necesiten. El equipo calculó que esto maximizaría el uso del ancho de banda en el enlace de datos y llevaría el método a una utilización perfecta; sin embargo, solo vieron alrededor del 90 por ciento. Identificaron y corrigieron un error de rendimiento en una popular biblioteca de PyTorch que provocaba comunicaciones de ida y vuelta innecesarias entre la CPU y la GPU. Con este error corregido, el equipo logró un tiempo de ejecución de 16.5 segundos por época con SALIENT.

“Creo que nuestro trabajo mostró que el diablo está en los detalles”, dice Kaler. “Cuando presta mucha atención a los detalles que afectan el rendimiento al entrenar una red neuronal gráfica, puede resolver una gran cantidad de problemas de rendimiento. Con nuestras soluciones, terminamos completamente obstruidos por el cómputo de la GPU, que es el objetivo ideal de dicho sistema”.

La velocidad de SALIENT se evaluó en tres conjuntos de datos estándar ogbn-arxiv, ogbn-products y ogbn-papers100M, así como en configuraciones de varias máquinas, con diferentes niveles de fanout (cantidad de datos que la CPU prepararía para la GPU) y a través de varias arquitecturas, incluida la última tecnología de punta, GraphSAGE-RI. En cada configuración, SALIENT superó a PyTorch Geometric, sobre todo en el gran conjunto de datos ogbn-papers100M, que contiene 100 millones de nodos y más de mil millones de bordes. Aquí, fue tres veces más rápido, ejecutándose en una GPU, que la línea de base optimizada que se creó originalmente para este trabajo; con 16 GPU, SALIENT fue ocho veces más rápido. 

Mientras que otros sistemas tenían hardware y configuraciones experimentales ligeramente diferentes, por lo que no siempre fue una comparación directa, SALIENT aún los superó. Entre los sistemas que lograron una precisión similar, las cifras de rendimiento representativas incluyen 99 segundos con una GPU y 32 CPU, y 13 segundos con 1,536 CPU. Por el contrario, el tiempo de ejecución de SALIENT con una GPU y 20 CPU fue de 16.5 segundos y fue de solo dos segundos con 16 GPU y 320 CPU. “Si observa los números finales que informan trabajos anteriores, nuestro tiempo de ejecución de 16 GPU (dos segundos) es un orden de magnitud más rápido que otros números que se informaron anteriormente en este conjunto de datos”, dice Kaler. Los investigadores atribuyeron sus mejoras de rendimiento, en parte, a su enfoque de optimizar su código para una sola máquina antes de pasar a la configuración distribuida. Stathas dice que la lección aquí es que, por su dinero, "tiene más sentido usar el hardware que tiene de manera eficiente y extrema, antes de comenzar a escalar a varias computadoras", lo que puede proporcionar ahorros significativos en costos y emisiones de carbono. que puede venir con el entrenamiento modelo.

Esta nueva capacidad ahora permitirá a los investigadores abordar y profundizar en gráficos cada vez más grandes. Por ejemplo, la red Bitcoin que se mencionó anteriormente contenía 100,000 1,000 nodos; el sistema SALIENT puede manejar hábilmente un gráfico XNUMX veces (o tres órdenes de magnitud) más grande.

“En el futuro, buscaríamos no solo ejecutar este sistema de entrenamiento de redes neuronales gráficas en los algoritmos existentes que implementamos para clasificar o predecir las propiedades de cada nodo, sino que también queremos realizar tareas más profundas, como identificar patrones comunes en un gráfico (patrones de subgráficos), [que] pueden ser realmente interesantes para indicar delitos financieros”, dice Chen. “También queremos identificar nodos en un gráfico que sean similares en el sentido de que posiblemente correspondan al mismo malhechor en un delito financiero. Estas tareas requerirían desarrollar algoritmos adicionales y posiblemente también arquitecturas de redes neuronales”.

Esta investigación fue apoyada por el MIT-IBM Watson AI Lab y en parte por el Laboratorio de Investigación de la Fuerza Aérea de EE. UU. y el Acelerador de Inteligencia Artificial de la Fuerza Aérea de EE. UU.