Optimización de redes neuronales: cómo desenredar el grafo computacional ✛
Contexto
Las redes neuronales profundas para sistemas de recomendación son monstruos de complejidad computacional. Cuanto más preciso es el modelo, más enredado es su paisaje interno: miles de operaciones se entrelazan en grafos con topologías caprichosas. Cada vuelta de más, cada división y unión repetida de tensores es como un sendero sin despejar que frena al viajero. A nivel de hardware, esto se traduce en una avalancha de lanzamientos de núcleo (kernel launch) en la GPU, que consumen tiempo y memoria. Los ingenieros intentan poner orden manualmente, pero los modelos cambian cada semana y las arquitecturas son cada vez más rebuscadas. Surgió la necesidad de un «enderezador» automático de senderos de grafos, capaz de encontrar y eliminar fragmentos de código ineficientes sin intervención humana.
Métodos
El núcleo de Optimus es una colección de más de treinta patrones de optimización, cada uno de los cuales describe un «patrón basura» típico en el grafo: por ejemplo, una cadena sin sentido de «cortar el tensor en trozos y luego volver a unirlos». La búsqueda de estos patrones se realiza mediante un algoritmo voraz que parte de «anclas» — operaciones clave como split u otras comunes en arquitecturas de recomendación. Desde el ancla, el algoritmo recorre los vecinos más cercanos, comparándolos con el patrón de referencia, y al coincidir activa la sustitución: en lugar de la construcción engorrosa se inserta un equivalente compacto. Se distinguen dos tipos de reparación: horizontal, cuando varias operaciones idénticas independientes (por ejemplo, una docena de capas lineales) se fusionan en una sola, y vertical, cuando se eliminan transiciones redundantes, como los famosos split-then-cat. Todo el procedimiento está integrado en el compilador PyTorch 2.x y funciona en cualquiera de sus representaciones intermedias (Torch-IR, Aten-IR), lo que lo hace universal para modelos de diversa índole.
Resultados
En cinco modelos industriales de recomendación de Meta de diferentes tipos arquitectónicos (tipo DHEN, tipo Wukong, tipo Interformer, entre otros), Optimus aumentó de forma consistente el rendimiento, con hasta un 63 % más de ejemplos de entrenamiento por segundo. La principal contribución vino de las optimizaciones verticales: la eliminación de los ciclos parásitos split-cat suprimió cientos de núcleos Triton superfluos. La ganancia de las fusiones horizontales dependió del número de operaciones idénticas: cuantas más, mayor el beneficio. El consumo máximo de memoria en modelos grandes con cuantización de activaciones cayó entre un 2 % y un 6 %, y el tiempo de compilación en frío se desplomó: en una de las pruebas, pasó de más de 4000 segundos a aproximadamente 3600, es decir, un ahorro superior a 400 segundos. En un ejemplo ilustrativo abierto de TorchBench, el número de núcleos (kernel) en entrenamiento se redujo de 12 772 a 24, y en inferencia, de 6320 a 6, lo que supuso una aceleración del entrenamiento de 11 veces y de la inferencia de 364 veces. La comparación con alternativas (TASO, TENSAT) mostró que estas no pudieron mejorar el grafo para este modelo, mientras que Optimus lo consiguió gracias a su flexible sistema de reglas.
Implicaciones
Optimus no es un artilugio de usar y tirar, sino un armazón modular integrado en el código abierto de PyTorch. Cualquier investigador o ingeniero puede añadir un nuevo patrón de optimización registrándolo a través de la interfaz estándar y aplicarlo de inmediato a su grafo de redes neuronales. Este enfoque convierte la antes tediosa manipulación manual del código en una limpieza sistemática a nivel del compilador. Para la industria, esto significa que las costosas GPU se utilizarán de manera más eficiente sin reescribir los modelos y, al mismo tiempo, se reducirá la huella de carbono del entrenamiento de sistemas de recomendación gigantescos.
Desarrollo futuro
Actualmente, las reglas principales funcionan sobre Torch-IR y Aten-IR, pero los autores planean extender el método a Inductor-IR, una representación de más bajo nivel donde se puede exprimir aún más rendimiento. Otra dirección es la selección automática del orden y del conjunto de patrones para un modelo concreto: hoy en día la secuencia de optimizaciones se elige manualmente, y lo deseable sería delegar esta tarea en un algoritmo que analice la estructura del grafo.
Impacto
La tecnología acelerará directamente el entrenamiento y el despliegue de sistemas de recomendación, un eslabón clave de las plataformas de internet modernas. También se puede adaptar a cualquier otra red neuronal profunda, especialmente aquellas con grafos computacionales ramificados: procesamiento del lenguaje natural, visión por computadora, redes neuronales de grafos.
Próximos pasos
Trasladar los avances al nivel de Inductor-IR e implementar un planificador automático de la secuencia de patrones. También está previsto ampliar la biblioteca de patrones para cubrir nuevos motivos arquitectónicos.
Problemas abiertos clave
El problema de la ejecución eficiente de redes neuronales en el hardware actual choca con la brecha entre la complejidad de las arquitecturas y el ancho de banda limitado de la memoria de las GPU. Optimus supone un paso hacia la eliminación automática de esta brecha a nivel de compilador, sin exigir al desarrollador que rehaga el código. Esto se inscribe en un contexto más amplio de búsqueda de métodos escalables de optimización numérica de cómputos.
La mirada de la máquina del conocimiento
Esta sección está dirigida al autor del trabajo y a sus colegas. Los trabajos vecinos los encontró la búsqueda por significado en nuestro archivo, no por coincidencia de palabras. No evaluamos: mostramos el campo alrededor.
Qué hace el trabajo
En este trabajo se presenta Optimus, un marco modular para la optimización automática de grafos de cómputo en PyTorch 2.x, orientado a modelos de recomendación. Utiliza una colección de más de treinta plantillas para buscar y reemplazar patrones ineficientes (como split-then-cat) mediante un algoritmo voraz basado en anclas. En cinco modelos industriales de Meta se logró una aceleración del entrenamiento de hasta el 63%, una reducción del pico de consumo de memoria del 2–6% y una disminución significativa del tiempo de compilación (de más de 4000 a aproximadamente 3600 segundos en una prueba).
Dónde está su fuerza
La fortaleza del trabajo radica en la integración práctica con el compilador de PyTorch y en su sistema flexible de reglas, que permite encontrar y eliminar automáticamente patrones ineficientes en grafos industriales reales, donde alternativas como TASO y TENSAT no logran hacerlo.
Hacia dónde puede avanzar
- Podría ser interesante revisar la experiencia de SpecFANN en la aceleración de cálculos mediante arquitecturas especializadas: allí se utiliza un emulador rápido para generar espectros, lo que se relaciona con la idea de reemplazar fragmentos pesados por equivalentes ligeros. Quizás los principios de reemplazo basados en sustitutos entrenados puedan adaptarse para ampliar el conjunto de plantillas de Optimus.se apoya en: SpecFANN: un boceto relámpago del espectro estelar arXiv:2607.17348v1
- Cerca se encuentra un trabajo sobre una red neuronal que detecta la tensión en piezas de cualquier forma: allí se aplican métodos de análisis de grafos para identificar nodos críticos. Esto podría sugerir cómo mejorar la heurística de selección de anclas en Optimus para encontrar optimizaciones más globales.se apoya en: La red neuronal que detecta la tensión en cualquier pieza, sin importar su forma arXiv:2606.08287
Vale la pena probar
- Probar el uso de un sustituto entrenado para predecir la ganancia de aplicar una plantilla, con el fin de priorizar los reemplazos.
- Probar el análisis de la 'tensión' en el grafo (en analogía con el trabajo sobre piezas) para seleccionar anclas que tengan más probabilidades de conducir a optimizaciones en cascada.
- Verificar si es posible generar automáticamente nuevas plantillas generalizando patrones frecuentes en grafos reales.
Trabajos cercanos en nuestro archivo
- La red neuronal que detecta la tensión en cualquier pieza, sin importar su formacercanía 0.62
- SpecFANN: un boceto relámpago del espectro estelarcercanía 0.62
- Cómo un gemelo digital salva el cerebro en minutoscercanía 0.62
- Una red neuronal aprende a fluir a través de un coladorcercanía 0.60
- Etiquetas para grafos: cuando el enfoque cuántico cede ante la simplicidadcercanía 0.60
- Las neuronas revelan sus secretos en minutos, no en horascercanía 0.59
El trabajo más cercano del archivo está a 0.62 por significado: el nivel habitual de nuestro corpus.
Resumen original
tal como lo escribieron los autores, en inglés arXiv:2607.02945
In large-scale industrial applications, deep learning models that power recommendation and ranking have complex and diverse model architectures. These models are continuously developed and refined by large teams of machine learning engineers, rendering manual optimization infeasible. Consequently, graph-based optimization techniques have become an industry standard for boosting performance, with PyTorch FX transformations leading the charge. These transformations typically rely on a set of human-engineered module-level rewrite rules which are not scalable to diverse model architectures. To address this limitation, we introduce Optimus, a general-purpose model transformation framework built in the PyTorch 2.x (PT2) machine learning compiler. With a concise set of predefined patterns, Optimus applies an efficient greedy search algorithm for pattern matching and replacement, while preserving model semantic. It is designed and implemented as a highly customizable and extensible framework integrated into the PT2 stack. Our evaluation shows that the framework can achieve up to 63% speedup, 6% peak memory reduction, and over 400 second compile time decrease for our industry-scale recommendation models compared to baselines. Optimus is open-sourced together with PyTorch 2.x as a customizable model transformation layer.
Reducir el tiempo de compilación en 400 segundos es como si, en lugar de preparar café, revisar el correo y esperar a que se abra el entorno de desarrollo, te pusieras manos a la obra de inmediato. Y la caída del número de lanzamientos de núcleos en la inferencia del modelo ilustrativo de 6320 a 6: esa compactación del grafo podría compararse con sustituir una multitud de mensajeros por un solo camión.
Cifras clave
- aumento del rendimiento (QPS): hasta 63%
- reducción de la memoria máxima: hasta 6%
- reducción del tiempo de compilación: más de 400 segundos
- núcleos en entrenamiento TorchBench antes/después: 12 772 -> 24
- núcleos en inferencia TorchBench antes/después: 6 320 -> 6