Recraft AI es una herramienta de diseño impulsada por modelos de difusión patentados. Su nuevo modelo Red Panda supera a todos los modelos de texto a imagen existentes, incluidos Midjourney, DALL-E 3 y FLUX. Recraft combina una experiencia de usuario única para diseñadores con herramientas de inteligencia artificial de vanguardia. Para dar soporte a millones de usuarios, los modelos de difusión requieren una infraestructura de inferencia sólida que combine hardware potente con software eficiente. En este artículo, exploraremos cómo las herramientas de aceleración de IA de TheStage ayudaron a los ingenieros e investigadores de IA de Recraft a lograr hasta el doble de rendimiento en las GPU de Nvidia a través de una interfaz Python intuitiva.

Introducción

En los últimos años, los modelos de difusión han mostrado resultados extraordinarios para la generación de contenido, incluidas imágenes, música, videos y mallas 3D. Estos modelos utilizan los costos computacionales del tiempo de inferencia para mejorar iterativamente los resultados de generación, actualizando ligeramente el resultado en cada paso de inferencia. Ahora podemos ver una analogía en los modelos de difusión que utilizan el razonamiento a través de inferencias de múltiples pasos para brindar respuestas de alta calidad.

En TheStage AI, estamos construyendo un marco matemático general para modelos arbitrarios de PyTorch para manejar un flujo de aceleración de modelos complicado de forma totalmente automática. Nuestro sistema detecta automáticamente las optimizaciones existentes en su hardware (cuantificación, esparsificación) y selecciona para cada capa el algoritmo adecuado para lograr la mejor calidad con un tamaño de modelo y restricciones de latencia deseados o encontrar la mejor aceleración con restricciones de calidad restringidas. ¡Es un problema matemático difícil que podemos resolver de una manera altamente eficiente! Este artículo explora cómo aplicamos estas herramientas a través de nuestra asociación con Recraft AI.

Al diseñar nuestras herramientas, decidimos respetar los siguientes principios:

• Personalización del hardware. Los productos de IA de alta calidad ya cuentan con su infraestructura preferida
• Preservación de la calidad. Los productos de IA de alta calidad no pueden aceptar la degradación de la calidad.
• Privacidad. Los productos de IA de alta calidad quieren mantener la confidencialidad de sus tecnologías y trabajar con herramientas en su propia infraestructura.
• Redes neuronales de dominio público arbitrarias. Los productos de inteligencia artificial de alta calidad pueden utilizar arquitecturas desarrolladas internamente, mientras que las herramientas de aceleración públicas para modelos de código abierto no pueden manejar arquitecturas de redes neuronales de dominio público complejas para producir resultados correctos.
• Pytorch: el marco de trabajo más popular y conveniente para muchos ingenieros de IA.

Dadas estas condiciones iniciales, nos propusimos crear herramientas con las siguientes características:

• Aceleración controlable. Enmarcamos la optimización de inferencia como un problema de optimización empresarial, que permite a los clientes especificar el tamaño del modelo, la latencia o la calidad que desean para sus datos.
• Compilación sencilla. La compilación de los modelos producidos para una ejecución eficiente del hardware requiere solo una línea de código. También proporcionamos una interfaz sencilla para gestionar las interrupciones de gráficos.
• Arranque en frío rápido. Para lograr el arranque en frío más rápido posible, permitimos guardar los modelos compilados. Por eso no utilizamos compiladores JIT.
• Implementación sencilla. Implementar el modelo optimizado debería ser tan sencillo como implementar el original.

Modelos de difusión de texto a imagen

En cada iteración del proceso de difusión, una red neuronal elimina el ruido de la imagen en el espacio latente de un codificador automático variacional. Luego, la imagen recién obtenida se vuelve a mezclar con ruido, pero con un peso cada vez menor. Durante las iteraciones iniciales, el modelo de difusión esboza la escena principal, aprovechando el peso significativo del ruido para realizar mejoras sustanciales. En iteraciones posteriores, refina los detalles de alta frecuencia. Esta observación nos permite diseñar canales de aceleración específicos al asignar estratégicamente la capacidad de la red en las capas de una iteración a otra, preservando la calidad. Sin embargo, dicha asignación requiere herramientas especializadas que combinen conocimientos matemáticos con ingeniería de sonido: ¡aquí es donde TheStage AI puede ayudar significativamente!

Modelos de difusión Aceleración y compresión

Los modelos de difusión acelerada pueden considerarse como la aceleración de DNN arbitrarias, pero debemos tener en cuenta desafíos específicos. Por ejemplo, la cuantificación estática, que normalmente proporciona una aceleración significativa, presenta un desafío en los modelos de difusión a medida que las distribuciones de activación cambian de una iteración a otra. Para abordar esto, debemos estimar correctamente los valores óptimos para todas las iteraciones o utilizar diferentes configuraciones de cuantificación para cada iteración.

Los modelos de difusión son difíciles de entrenar y lograr un alto rendimiento. Sin embargo, los resultados demostrados por el equipo de Recraft superan a todos los modelos de texto a imagen modernos . Validar la degradación de dichos modelos puede ser difícil, por lo que es crucial utilizar técnicas de aceleración que preserven la semántica de las características originales. Los algoritmos de cuantificación pueden ser una buena opción si pueden manejar el desafío de variar las distribuciones de activación. Echemos un vistazo a nuestras canalizaciones automáticas, que describiremos en las siguientes secciones.

Elaboración de perfiles

Perfilar un modelo dado con unos datos específicos permite:

• Determinar el tamaño de cada parámetro
• Identificar algoritmos de cuantificación, esparsificación y poda aplicables para cada bloque básico.
• Estimar la latencia para bloques individuales con diferentes diseños de memoria
• Compilar toda la información recopilada para ANNA (Acelerador automatizado de redes neuronales)

Compresión y aceleración automáticas

Una vez que el generador de perfiles ha recopilado todos los datos necesarios, podemos iniciar nuestro tablero ANNA y mover el control deslizante para generar diferentes versiones optimizadas del modelo. Nuestros usuarios pueden seleccionar los mejores candidatos en función de la relación calidad-costo de inferencia. Nuestras herramientas manejan estas decisiones subjetivas de calidad de una manera sencilla.

Operaciones de fusión y compilación

Como se mencionó anteriormente, no utilizamos la compilación JIT porque aumenta el tiempo de inicio en frío de un nuevo nodo. Tampoco utilizamos compiladores estándar. En su lugar, compilamos nuestra propia configuración acelerada compleja que puede combinar diferentes algoritmos. Para lograrlo, hemos desarrollado nuestro propio protocolo interno para definir capas aceleradas por DNN de una manera independiente del hardware. Uno de los beneficios clave del marco de aceleración de IA de TheStage es que se puede compilar un único punto de control para una amplia gama de hardware, lo que resuelve la compatibilidad entre plataformas para el software de IA. Esta característica será particularmente importante para la implementación de dispositivos de borde en el desarrollo de aplicaciones.

Los objetivos de los compiladores DNN son:

• Simplificación de grafos. Simplifique el grafo de ejecución mediante la fusión de operaciones matemáticas, reduciendo el tiempo de inferencia.
• Gestión de memoria. Calcule la memoria necesaria para cada operación y administre la programación de asignaciones con una reutilización eficiente de la memoria.
• Implementación óptima. Perfilar la implementación óptima para cada operación básica (una tarea desafiante, ya que la mejor implementación puede requerir diseños de memoria específicos, lo que lleva al análisis de conexiones entre capas).
• Programación de operaciones. Cree una programación de operaciones para el gráfico de ejecución optimizado
• Serialización. Guarde toda esta información para evitar tener que volver a compilar el modelo en ejecuciones posteriores.

Implementación y servicio

Los servidores de inferencia y los canales de escalado automático desempeñan un papel importante en el procesamiento rentable y eficiente de las solicitudes entrantes. También pueden incluir la agrupación de solicitudes específicas y la recopilación de estadísticas para configurar el escalado predictivo para los escaladores automáticos. En nuestros próximos artículos, analizaremos los servidores de inferencia eficientes con más detalle.

Resultados

Al aplicar todas las secuencias de comandos, podemos lograr un rendimiento mejor que el compilador de PyTorch ( torch.compile ) y, por supuesto, significativamente mejor que la ejecución ansiosa de PyTorch con float16. Además, como el compilador de PyTorch utiliza un enfoque de compilación JIT en cada inicialización de modelo, requiere una recompilación para muchos tamaños de entrada, lo que hace que el inicio en frío sea lo suficientemente largo para aplicaciones prácticas donde la latencia es muy importante.

Beneficios empresariales

Estos son los principales beneficios comerciales de la aceleración resultante para el producto de Recraft:

• Reducir los costos de infraestructura al brindar servicio al doble de usuarios con el mismo hardware
• Experiencia de usuario mejorada con tiempos de respuesta de generación de imágenes más rápidos
• Capacidad de atender a más usuarios simultáneos durante picos de carga
• Ventaja competitiva a través de una entrega de servicios más rápida

Las herramientas de optimización de IA de TheStage nos permiten acelerar nuestros modelos de texto a imagen sin degradación de la calidad, creando una mejor experiencia de usuario para nuestros clientes.

Directora ejecutiva de Recraft, Anna Veronika Dorogush

Expresiones de gratitud

Estos resultados brindan una excelente validación de nuestras herramientas e investigaciones sobre productos de carga de trabajo a gran escala. El equipo de IA de TheStage continúa trabajando para ofrecer un rendimiento aún mayor. Para lograrlo, estamos colaborando con socios excepcionales. Estamos profundamente agradecidos a:

• Agradecemos a Anna Veronika, directora ejecutiva de Recraft, por la fructífera colaboración. Estamos encantados de ser, aunque sea una pequeña parte, de su gran recorrido para ofrecer las mejores herramientas de diseño.
• Agradecemos a Pavel Ostyakov, director de inteligencia artificial de Recraft, por su experiencia en DNN, sus valiosos comentarios sobre las herramientas y por establecer objetivos desafiantes para nuestro proyecto de cooperación.
• El equipo de Recraft AI por crear este fantástico producto. Las imágenes de este artículo se generaron con Recraft.
• Al equipo de Nebius por su constante apoyo con una excelente infraestructura de GPU para nuestra investigación.

Contactos / Recursos