Avances en técnicas de aprendizaje automático para el aprendizaje automático

• A medida que los modelos de ML crecen, la eficiencia de ejecución depende en gran medida no solo del hardware sino también de la optimización del compilador, y Google Research y Google DeepMind presentaron un enfoque para volver a mejorar esta área con ML
• Las decisiones del compilador tomadas por debajo de frameworks como TensorFlow, JAX y PyTorch pueden cambiar mucho el tiempo de ejecución y el uso de recursos incluso para el mismo modelo
• TpuGraphs es un conjunto de datos que contiene grafos computacionales de programas de ML para TPU, configuraciones de compilación y tiempos de ejecución, y puede usarse para investigar modelos de costo basados en aprendizaje
• Para reducir la limitación de que es difícil entrenar grafos computacionales grandes de una sola vez, Graph Segment Training divide el grafo en segmentos, reduce el uso de memoria y acorta el tiempo de entrenamiento 3 veces
• En la competencia de Kaggle participaron 792 personas de 616 equipos de 66 países, y se validaron técnicas de mejora de modelos reales de predicción de costo como compresión de grafos, ajuste de valores de padding, adición de características de nodos y attention entre configuraciones

---

Por qué los compiladores de ML determinan el rendimiento de ejecución

• Los modelos modernos de ML realizan tareas como comprensión del lenguaje natural, conversación, generación de imágenes y generación de video, y se escriben y entrenan con frameworks de programación de ML como TensorFlow, JAX y PyTorch
• Los frameworks ofrecen operaciones de álgebra lineal como multiplicación de matrices y convolución, así como capas de redes neuronales como capas de convolución 2D y capas Transformer
• Los usuarios no necesitan manejar directamente los detalles de cómo ejecutar eficientemente el modelo en hardware, ya que el compilador bajo el framework optimiza automáticamente el modelo
• Sin embargo, los compiladores suelen resolver problemas de optimización complejos mediante heurísticas, por lo que no siempre logran el rendimiento óptimo

Grafos computacionales y optimización en dos etapas

• Los compiladores de ML convierten las instrucciones matemáticas escritas por el usuario en instrucciones que pueden ejecutarse en el hardware real
• Los programas de ML pueden representarse como un grafo computacional
  • Los nodos representan operaciones sobre tensores como matrix multiplication
  • Las aristas representan tensores que fluyen de un nodo a otro
• La optimización del compilador se divide en dos grandes tipos
  • Optimización a nivel de grafo: toma decisiones considerando el contexto del grafo completo y transforma todo el grafo
  • Optimización a nivel de kernel: transforma un solo kernel, que es un fused subgraph, de forma independiente de otros kernels

Trade-offs de rendimiento en el layout de memoria

• Los tensores 2D como las matrices pueden almacenarse en memoria como [A B C a b c] o [A a B b C c], lo que corresponde respectivamente a layouts row-major y column-major
• Una de las optimizaciones importantes de los compiladores de ML es asignar un layout de memoria a todos los tensores intermedios del programa
• Un layout específico puede ser el más eficiente para una operación individual, pero si el layout no coincide entre add y convolution, el compilador debe insertar una operación adicional de copy
• A la inversa, aunque el rendimiento de cada operación individual sea un poco menor, una configuración que no requiera conversión de layout puede ser mejor para la ejecución total
• En la suite de benchmarks de XLA, al elegir una configuración de layout óptima en lugar de la configuración predeterminada del compilador, se observó una mejora de velocidad de hasta 32%

Dataset TpuGraphs

• TpuGraphs es un dataset de modelos de costo basados en aprendizaje para programas que se ejecutan en las TPU personalizadas de Google
• El objetivo es entrenar un modelo de costo que reciba el programa de entrada y la configuración del compilador para predecir el tiempo de ejecución del programa
• El dataset cubre dos tipos de configuración del compilador XLA
  • layout: una configuración que generaliza el concepto de row-major y column-major de matrices a tensores de alta dimensión
  • tiling: configuración del tamaño de los tiles
• Cada ejemplo incluye el grafo computacional de una carga de trabajo de ML, la configuración de compilación y el tiempo de ejecución al compilar con esa configuración
• Los grafos se recopilaron de programas de ML de código abierto e incluyen arquitecturas de modelos como ResNet, EfficientNet, Mask R-CNN y Transformer
• El método de descarga y el código inicial están disponibles en TpuGraphs GitHub
• TpuGraphs tiene 25 veces más grafos que el mayor dataset previo de predicción de propiedades de grafos con tamaños de grafo similares, y su tamaño promedio de grafo es 770 veces mayor que el de los datasets previos de predicción de rendimiento de programas de ML

Modelo de costo base y arquitectura GNN

• TpuGraphs también proporciona un modelo de costo base basado en aprendizaje, y como el programa de entrada se representa como un grafo, utiliza una GNN
• Las características de los nodos se componen de dos partes
  • opcode id: la información de nodo más importante, que indica el tipo de operación de tensor
  • Otras características del nodo
• El modelo base convierte el opcode id en un opcode embedding mediante una tabla de lookup de embeddings
• El opcode embedding y el resto de las características del nodo se combinan para usarse como entrada de la GNN
• Los embeddings de nodos generados por la GNN se combinan en un embedding de grafo de tamaño fijo mediante reducciones simples de graph pooling como sum y mean
• El embedding final del grafo se transforma en una única salida escalar a través de una capa feedforward

Entrenamiento de grafos grandes con Graph Segment Training

• Graph Segment Training es una técnica de escalado para entrenamiento de GNN orientada a manejar grafos grandes en dispositivos con capacidad de memoria limitada
• Este método está dirigido a escenarios de graph-level prediction, donde el objetivo de predicción no es un nodo ni una arista sino el grafo completo
• Los grafos computacionales pueden contener cientos de miles de nodos, por lo que usar el grafo completo de una sola vez en Full Graph Training puede ser computacionalmente inviable
• GST divide un grafo grande en segmentos pequeños y actualiza el modelo seleccionando solo un subconjunto aleatorio de segmentos
• Los segmentos restantes generan embeddings sin guardar activations intermedias, lo que reduce el uso de memoria
• Al combinar todos los embeddings de segmentos, se construye el embedding del grafo grande original y se usa para la predicción
• También se introducen una historical embedding table y segment dropout para mitigar la obsolescencia de los historical embeddings
• El método completo reduce el tiempo de entrenamiento end-to-end en 3 veces

Técnicas de mejora validadas en la competencia de Kaggle

• La competencia de Kaggle Fast or Slow? Predict AI Model Runtime se llevó a cabo con base en el dataset TpuGraphs, y participaron 792 personas de 616 equipos de 66 países
• Hubo 10,507 envíos, 153 personas participaron por primera vez en una competencia de Kaggle y, de ellas, 47 quedaron dentro del top 100
• Los equipos participantes probaron varias técnicas
  • Pruning y compresión de grafos: se experimentó con métodos para comprimir grafos grandes en lugar de usar GST, y se utilizó un enfoque que conserva solo subgrafos que incluyen nodos configurables y sus vecinos inmediatos
  • Cambio de valor de padding: como el valor de padding predeterminado 0 entra en conflicto con valores válidos de features, usar -1 mejoró considerablemente la precisión del modelo
  • Adición de características de nodos y cambios en la codificación: características adicionales de nodos como las contracting dimensions de dot general son importantes, y la forma de codificar las características también puede afectar los resultados
  • cross-configuration attention: el equipo ganador diseñó una capa simple que permite al modelo comparar explícitamente configuraciones, y obtuvo resultados mucho mejores que inferir cada configuración por separado
• Los resultados de la competencia y la solución ganadora se presentarán el 16 de diciembre de 2023 en la sesión de competencias del ML for Systems workshop en NeurIPS

Sesión relacionada con NeurIPS Expo

• Para lectores interesados en datos estructurados e investigación en inteligencia artificial, el panel de NeurIPS Expo Graph Learning Meets Artificial Intelligence se realizará el 9 de diciembre de 2023
• Este panel abordará temas como los avances en modelos de costo basados en aprendizaje

[GN⁺]

Opiniones en Hacker News

• Los compiladores de ML están sobrevalorados. Es el mismo compromiso que con los compiladores tradicionales: se obtiene muchísimo más rendimiento que contratando programadores especializados en performance, pero estos últimos normalmente son mucho más rápidos y, en algunos casos, pueden estar varios órdenes de magnitud por delante.
  Se quedan cortos en varios niveles. A nivel de algoritmos, no retroalimentan a la persona con trucos para hacer que la red sea más rápida; apenas dan señales muy básicas. También se pierde la intención. Quien diseña una red de ML especifica la estructura en Python, pero después de varias etapas de transformación hacia abajo puede salir algo completamente distinto. Hace poco vi un compilador que, al hacer un slice update, creaba todos los rangos posibles de índices de un arreglo, los recortaba para obtener los índices a actualizar y luego hacía un scatter; lo reemplacé por una sola llamada a memcpy. Los kernels también son ineficientes. Cada vez que la salida de estos compiladores compite contra un programador de assembly experimentado, el compilador pierde, por lo general por más de 30%. Parece algo fácil de resolver, pero si nadie lo ha resuelto bien en los últimos 50 años, claramente no es tan simple como suena.

  • Si miras el motor de ajedrez Stockfish, en la evaluación del tablero descartaron heurísticas escritas por humanos durante años y las reemplazaron por una pequeña red neuronal, y funciona mejor.
    Los compiladores también tienen muchas heurísticas, como inlining, loop unrolling y vectorización, así que las redes neuronales podrían ayudar y quizá sean más fáciles de mantener que una enorme cantidad de heurísticas escritas por humanos.
  • Dijiste que es el mismo compromiso que con los compiladores tradicionales, y lo curioso es que esos compiladores tradicionales resultaron ser tremendamente útiles.
  • Suena demasiado tajante y cerrado.
  • Exacto. ¿Por qué alguien usaría gcc/clang si puede contratar a alguien para escribir assembly a mano?
  • Ese rendimiento es justamente lo importante. No se puede poner a un experto en performance en cada trabajo de ML.
    Esto sigue siendo mucho mejor que no tener ninguna de estas optimizaciones.

• ¿Alguien puede explicar esto de forma un poco más realista? Me interesa saber cuál es el estado real de los compiladores de ML hoy y qué podemos esperar en el corto plazo.

  • Uno de los enfoques más sencillos es torch.compile. Es la iteración más reciente del compilador de PyTorch; antes estaban TorchScript y FX Tracing.
    Solo escribes model = torch.compile(model). “En estos 163 modelos open source, torch.compile funcionó en el 93% de los casos y aceleró el entrenamiento en GPU NVIDIA A100 un 43%. En precisión Float32 fue 21% más rápido en promedio, y en precisión AMP, 51% más rápido en promedio.”[1] Parece que Google quiere que más gente participe en la I+D de métodos como este.
    [1] https://pytorch.org/get-started/pytorch-2.0/
  • Lo que se espera en el corto plazo es poder usar AMD, CUDA, TPU, CPU, etc., aunque no haya soporte explícito del proveedor para el framework en el que se desarrolló el modelo.
    La realidad es compleja, así que simplificando bastante: se compila el grafo de cómputo a alguna representación intermedia y se implementan backends para ella. Entre los proyectos relacionados están stableHLO, IREE y openXLA. El compilador jit de Jax también puede verse como una forma de este tipo de compilador. Baja las operaciones trazadas a XLA, y XLA a su vez hace varias magias para que funcionen en el backend. Al final, cuanto más se baja en la pila, todo sigue siendo transformación y abstracción.
  • Mira torch.compile.

• En resumen, se trata de mejorar con redes neuronales de grafos (GNN) la predicción del rendimiento en tiempo de ejecución de un grafo de cómputo. Usan un diccionario de embeddings para el opcode de cada nodo junto con otras características del nodo, como shape, bits y window size ([1]).
  Publicaron en [2] un gran dataset de grafos con distintas configuraciones de compilación de XLA y el rendimiento resultante en TPU, y en [3] mejoraron la predicción para grafos más grandes mediante una forma de particionar grafos (es la primera vez que veo METIS graph partition) y varias técnicas de entrenamiento. Esto trata de predecir el rendimiento de un grafo dado, no de mejorar, proponer ni modificar un grafo equivalente nuevo. Como en FunSearch, un modelo con predicciones decentes puede usarse junto con búsqueda 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

• ¿Alguien puede explicar cómo funciona la convolución en ese grafo? Se convoluciona un tensor de shape [2,4,16] con un kernel de shape [4,16,8] y sale un tensor [2,8]; ¿cómo es posible?

  • No sé si ayuda, pero en el tensor de entrada [2,4,16], el 2 puede verse como el tamaño de batch, el 4 como la dimensión de características de entrada y el 16 como la dimensión de canales de entrada.
    En el kernel [4,16,8], el 4 es el tamaño de la ventana del filtro, el 16 coincide con la dimensión de canales de entrada y el 8 es la dimensión de canales de salida. En la salida [2,8], el 2 se mantiene como tamaño de batch y el 8 coincide con la dimensión de canales de salida del kernel. A primera vista las dimensiones no parecen cuadrar, pero la convolución en el grafo aprovecha la estructura de vecindad. El kernel se desplaza sobre el grafo, aplica pesos al nodo actual y a las características de sus vecinos dentro de cierto radio, y reúne esa suma ponderada para crear nuevas características para cada canal de salida. La estructura del grafo, los pesos de las aristas y detalles de implementación como padding y stride también pueden afectar el shape de salida.

• ¿En qué estado está Gemini?

  • Es interesante que GPT-4 siga dominando: https://twitter.com/lmsysorg/status/1735729398672716114
    Solo de los que me vienen a la mente, hay al menos cinco modelos base como Llama, Claude, Gemini, Falcon y Mistral, que se van superando unos a otros, pero GPT sigue estando un nivel arriba, y ya lleva un año así. Se ha vuelto claro que los grandes modelos de lenguaje basados en Transformer son lo bastante simples como para que cualquiera pueda crearlos si puede gastar alrededor de un millón de dólares en tiempo de GPU, pero aun así no logran alcanzar por completo a OpenAI. ¿Cuál será su ingrediente secreto?

• ¿Y qué hay del Transformer en sí? ¿Hay alguna pista de que sea óptimo en algún sentido?

• Siento que enterraron la idea principal en el primer párrafo, pero el resto está genial.

• El ritmo de avance del ML en este momento es asombroso. No creo en la singularidad, pero está cambiando el software y la sociedad de formas que nadie puede predecir.

  • Viendo esto y FunSearch, parece que la singularidad está cerca.
    https://deepmind.google/discover/blog/funsearch-making-new-d...
  • A mí me parece otra fiebre del oro más después de las puntocom, mobile, cloud y VR.
  • En 5 años no creo que la gente programe como ahora.
  • Primero quiero ver que proponga curas para enfermedades difíciles de tratar. La singularidad en sí no significa nada si no beneficia a los humanos, y ese beneficio debería estar sobre todo en mejorar la salud y reducir el sufrimiento.