🤖 machine learning

AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers

AoiZora es un planificador de topología mediado por compilador que optimiza la inferencia de difusión de video de baja latencia en sub-slices de TPU al integrar la conciencia de la interconexión física en el proceso de búsqueda de auto-paralelismo, logrando una reducción de latencia de hasta 1.42x para Wan 2.1 en comparación con las soluciones existentes.

Autores originales: Kaijian Wang, Yuanyuan Xu, Fanjiang Ye, Ye Cao, Jingwei Zuo, T. S. Eugene Ng, Yarong Mu, Yuke Wang

Publicado 2026-06-17

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Kaijian Wang, Yuanyuan Xu, Fanjiang Ye, Ye Cao, Jingwei Zuo, T. S. Eugene Ng, Yarong Mu, Yuke Wang

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando organizar una carrera de relevos masiva y de alto riesgo para crear un video. Los "corredores" son chips informáticos (específicamente los chips TPU de Google) y el "testigo" es la información necesaria para generar cada fotograma del video.

El proceso de generación de video es como un bucle largo y repetitivo donde los corredores tienen que limpiar una imagen con ruido una y otra vez (llamado "denoising" o eliminación de ruido) hasta que se vea clara. Debido a que el video es enorme y el proceso es lento, no puedes hacerlo en un solo chip; necesitas a todo un equipo de chips trabajando juntos.

El Problema: El error del "Mapa de Asientos"
Actualmente, cuando los ingenieros configuran este equipo, se enfocan en el plan lógico: "El Corredor A se encargará del fondo, el Corredor B se encargará del cielo y el Corredor C se encargará de las personas". Esto es como asignar tareas en una pizarra.

Sin embargo, a menudo ignoran la realidad física: ¿Dónde están sentados exactamente estos corredores en el estadio?

En un estadio real (la red física de TPU), algunos corredores están sentados justo al lado de otros y pueden pasar el testigo instantáneamente. Otros están en el lado opuesto del campo, y pasar el testigo toma tiempo y causa atascos de tráfico en los caminos compartidos entre ellos.

El artículo encontró que incluso si les das a los corredores las mismas tareas (el mismo plan lógico), simplemente cambiar dónde se sientan en el estadio puede hacer que la carrera sea un 16.6% más lenta o más rápida. El principal culpable no es la distancia que corren, sino el hecho de que algunos corredores se ven obligados a usar el mismo pasillo estrecho (un enlace de red compartido) al mismo tiempo, lo que provoca un atasco de tráfico.

La Solución: AoiZora (El Entrenador Inteligente)
Los autores construyeron una nueva herramienta llamada AoiZora (llamada así por las palabras japonesas para "Cielo Azul"). Piensa en AoiZora como un entrenador súper inteligente que no solo escribe la jugada, sino que también diseña el mapa de asientos basado en el diseño real del estadio.

AoiZora resuelve un problema complicado: revisar todas las posibles disposiciones de asientos es demasiado lento y costoso. Si intentaras probar cada una de las formas de sentar a los corredores, pasarías más tiempo planeando que corriendo la carrera.

Cómo funciona AoiZora (La estrategia de dos etapas)
AoiZora utiliza un ingenioso enfoque de "filtrar y refinar", como un proceso de contratación:

Etapa 1: El cribado rápido de currículums (Pre-compilación)
Primero, AoiZora observa los "currículums" (los planes de tareas lógicos) utilizando un método barato y rápido. No construye todo el estadio todavía; solo verifica si las tareas tienen sentido sobre el papel. Descarta rápidamente los malos planes que claramente fallarían (como intentar dividir una tarea que no se puede dividir). Es rápido y económico, como una entrevista de 30 segundos.
Etapa 2: La simulación completa (Post-compilación)
Para los pocos candidatos de élite que sobrevivieron a la primera ronda, AoiZora realiza el trabajo duro. Construye una simulación detallada del estadio real (los chips TPU físicos) y realiza una prueba completa. Ve exactamente dónde ocurrirán los atascos de tráfico en los enlaces de red específicos. Luego, elige la única mejor disposición de asientos que evite esos atascos.

El Resultado
Al usar este método de dos pasos, AoiZora encontró la combinación perfecta de "qué hacer" y "dónde hacerlo".

La Analogía: Imagina que estás organizando una cena.
- La forma antigua: Decides quién trae cada plato (plan lógico) y les dices que se sienten en cualquier lugar. A veces, la persona que trae la sopa pesada tiene que pasar junto a la persona que lleva el pastel frágil y chocan entre sí, derramándolo todo.
- La forma de AoiZora: Primero descartas rápidamente a los invitados que no pueden traer un plato. Luego, para los invitados restantes, miras el diseño real de la mesa y sientas al portador de la sopa justo al lado de la puerta de la cocina y al portador del pastel al otro lado, asegurándote de que nadie choque con nadie.

La Conclusión
En los chips TPU de Google, AoiZora hizo que la generación de video fuera 1.42 veces más rápida que los métodos anteriores. No cambió el videojuego en sí ni las reglas de la carrera; simplemente descubrió la forma más inteligente de sentar a los corredores para que pudieran pasar el testigo sin tropezarse unos con otros. Esto permite una generación de video más rápida y de menor latencia sin necesidad de reescribir el software subyacente.

Resumen Técnico: AoiZora

Planteamiento del Problema
La inferencia de difusión de video, particularmente para modelos como Wan 2.1, se ha convertido en una carga de trabajo generativa crítica. Sin embargo, el proceso de inferencia es computacionalmente intensivo, requiriendo cientos de miles de tokens espacio-temporales y pasos de eliminación de ruido repetidos sobre grandes tensores de activación. Esto requiere inferencia distribuida a través de múltiples aceleradores. Los sistemas de auto-paralelismo actuales (por ejemplo, aquellos basados en JAX/XLA) buscan eficazmente mallas de dispositivos lógicos óptimas (determinando cómo se fragmentan los tensores), pero ignoran en gran medida la colocación física de estos rangos lógicos en el hardware.

En los sub-slices de TPU (subconjuntos rectangulares contiguos de un pod de TPU), la interconexión entre chips (ICI) es una malla de vecinos más cercanos sin enlaces de retorno (wrap-around) en todas las dimensiones. En consecuencia, la misma estrategia de fragmentación lógica puede resultar en latencias finales significativamente diferentes dependiendo de cómo los rangos lógicos se mapeen a los chips físicos. El artículo observa que este "efecto de colocación" por sí solo puede causar una variación de hasta un 16.6% en la latencia de un paso de eliminación de ruido en un slice de v5e-8, debido principalmente a la contención de enlaces compartidos en lugar del conteo de saltos (hop count). Los sistemas existentes vinculan los rangos lógicos a los dispositivos físicos de forma temprana en el flujo de trabajo (a menudo mediante un orden por defecto) antes de que el compilador exponga el grafo concreto de comunicación colectiva, dejando la el rendimiento dependiente de la topología sin optimizar.

Metodología: Un Planificador Mediante Compilador de Dos Etapas
AoiZora aborda esta brecha situando la planificación consciente de la topología en la frontera entre el análisis del compilador y la ejecución. Aprovecha la "asimetría de costo/fidelidad" del pipeline de compilación de JAX/XLA: las representaciones intermedias (IR) tempranas son económicas pero carecen de detalles concretos de comunicación, mientras que las optimizaciones de alto nivel (HLO) totalmente compiladas son costosas pero exponen el grafo de colectivos y los grupos de dispositivos exactos.

El planificador opera en dos etapas distintas:

Etapa 1: Poda de Fragmentación Lógica Pre-Compilación
- Objetivo: Eliminar eficientemente los candidatos de fragmentación lógica ineficientes.
- Mecanismo: AoiZora enumera las asignaciones de ejes de carga de trabajo legales (por ejemplo, paralelismo de tensor, de secuencia, de datos) y las reduce a las IR de StableHLO/Shardy.
- Análisis: Construye un DAG de predicción que incluye comunicaciones explícitas y nodos de comunicación implícita sintética (por ejemplo, para re-fragmentación, reducciones fragmentadas y límites de kernels opacos) para estimar el volumen de comunicación lógica.
- Calificación: Utiliza un límite de rendimiento ajeno a la colocación ( $Q_1$ ) para calificar los candidatos. Esta etapa filtra las fragmentaciones con malas huellas de memoria o comunicación lógica excesiva, reteniendo solo un pequeño conjunto "Top-K" de supervivientes.
Etapa 2: Clasificación Topológica Post-Compilación
- Objetivo: Seleccionar la colocación física óptima para las fragmentaciones lógicas supervivientes.
- Mecanismo: Los candidatos lógicos Top-K son compilados completamente por XLA para generar HLO concreto.
- Análisis: El sistema enumera distintos órdenes de ejes físicos (mapeando ejes lógicos a coordenadas de sub-slice de TPU). Analiza el HLO compilado para extraer el grafo de colectivos exacto, incluyendo grupos de réplica, tamaños de carga útil y dependencias.
- Calificación: Un modelo de costo consciente de la topología ( $Q_2$ ) clasifica cada par de fragmentación-colocación. Este modelo reconstruye rutas físicas para los colectivos, estima las cargas de los enlaces y penaliza las colocaciones que causan contención en los enlaces ICI compartidos.
- Selección: Se selecciona el par con el costo estimado más bajo.

Contribuciones Clave

Caracterización de la Sensibilidad de la Colocación: Los autores demuestran que, en los sub-slices de TPU, la colocación física es una variable de optimización de primer orden. Muestran que fragmentaciones lógicas idénticas pueden producir diferencias de latencia de hasta un 16.6% debido a la contención de enlaces compartidos.
Búsqueda Consciente de la Fidelidad del Compilador: El artículo identifica y explota la asimetría entre las IR de pre-compilación baratas (aptas para la poda de alto recall) y el HLO compilado costoso (necesario para el ranking de topología preciso).
Diseño del Sistema AoiZora: Un planificador mediado por el compilador que separa la poda lógica del ranking físico. Crucialmente, se integra con el stack estándar de JAX/XLA sin modificar el código del modelo, las reglas de reducción, las implementaciones de kernels colectivos o el enrutamiento de red. Produce una estrategia de fragmentación y un orden de eje físico que el compilador estándar luego ejecuta.
Evaluación: Experimentos extensos con Wan 2.1 a través de sub-slices de TPU v5e (4, 8 y 16 chips) y diversas resoluciones de video (480p/720p) y conteos de fotogramas.

Resultados Experimentales
Evaluado en Wan 2.1, AoiZora supera consistentemente a dos líneas base:

Despliegue Existente: Una estrategia de paralelismo configurada por expertos y ajustada a mano (usando FSDP, paralelismo de secuencia y de tensor) con colocación física fija.
Línea Base de Búsqueda Lógica: Un sistema inspirado en Alpa que optimiza la fragmentación lógica utilizando un modelo de costo abstracto pero que no realiza el ranking de colocación física post-compilación.

Ganancias de Rendimiento: AoiZora logra aceleraciones de hasta 1.42× sobre la solución existente y 1.37× sobre la línea base de búsqueda lógica. Las aceleraciones de la media geométrica son de 1.09× y 1.11×, respectivamente.
Fidelidad: La etapa de pre-compilación retiene exitosamente la mejor fragmentación lógica medida en todos los entornos probados (100% de recall en el top-1), mientras que la etapa de post-compilación logra una alta correlación de rango ( $\rho_s$ hasta 0.985) entre la latencia predicha y la medida.
Sobrecarga (Overhead): El costo de planificación offline es práctico, variando desde ~47 segundos (v5e-4) hasta ~376 segundos (v5e-16), lo cual se amortiza a lo largo de solicitudes de inferencia repetidas de la misma forma.

Significancia y Reivindicaciones
El artículo afirma que AoiZora cierra la brecha entre el paralelismo lógico y la topología física de la TPU sin requerir cambios en el stack de servicio subyacente. Al diferir las decisiones de colocación física hasta que el compilador ha materializado el grafo de colectivos concreto, AoiZora evita la "omisión" de los sistemas de auto-paralelismo actuales que tratan las mallas lógicas como abstracciones completas.

Los autores posicionan a AoiZora como una solución práctica para las restricciones específicas de los sub-slices de TPU, donde la pérdida de los enlaces de retorno hace que la colocación consciente de la topología sea crítica. Aseguran que su enfoque permite a los desarrolladores de modelos permanecer agnósticos a los detalles de la topología mientras logran una latencia de inferencia casi óptima para las cargas de trabajo de difusión de video. El trabajo sugiere que para arquitecturas DiT de video densas con formas de tensores estáticas, un paso de planificación offline de una sola vez puede mejorar significativamente el rendimiento del servicio al resolver los problemas de contención inherentes a las colocaciones físicas fijas.

Más como este