Cómo Rufus duplicó su velocidad de inferencia y manejó el tráfico de Prime Day con chips AWS Ai y decodificación paralela

Los modelos de idiomas grandes (LLM) han revolucionado la forma en que interactuamos con la tecnología, pero su adopción generalizada ha sido bloqueada por una alta latencia de inferencia, rendimiento limitado y altos costos asociados con la generación de texto. Estas ineficiencias son particularmente pronunciadas durante los eventos de alta demanda como Amazon Prime Day, donde sistemas como Rufus, el Asistente de Compras de Amazon AI, deben manejar una escala masiva mientras se adhieren a los estrictos requisitos de latencia y rendimiento. Rufus es un asistente de compras con IA diseñado para ayudar a los clientes a tomar decisiones de compra informadas. Impulsado por LLMS, Rufus responde a las preguntas del cliente sobre una variedad de necesidades y productos de compras y simplifica la experiencia de compra, como se muestra en la siguiente imagen.

Rufus se basa en muchos componentes para brindar su experiencia al cliente, incluida una base LLM (para generación de respuesta) y un modelo de planificador de consultas (QP) para la clasificación de consultas y la mejora de la recuperación. El modelo analiza las preguntas del cliente para comprender su intención, ya sea un lenguaje natural basado en palabras clave o conversacionales. QP está en la ruta crítica para Rufus porque Rufus no puede iniciar la generación de tokens hasta que QP proporciona su salida completa. Por lo tanto, reducir la latencia de generación de texto de extremo a extremo de QP es un requisito crítico para reducir la primera latencia del fragmento en Rufus, que se refiere al tiempo necesario para generar y enviar la primera respuesta a una solicitud de usuario. Bajar esta latencia mejora la capacidad de respuesta percibida y la experiencia general del usuario. Esta publicación se centra en cómo el modelo QP utilizó un borrador de decodificación especulativa céntrica (SD), también llamada decodificación paralela, con AWS AI Chips para satisfacer las demandas del día principal. Combinando la decodificación paralela con AWS Entrenium y Chips de inferenciaRufus logró dos veces los tiempos de respuesta más rápidos, una reducción del 50% en los costos de inferencia y la escalabilidad perfecta durante el tráfico máximo.

Escalar llms para el día primo

Prime Day es uno de los eventos más exigentes para la infraestructura de Amazon, lo que lleva los sistemas a sus límites. En 2024, Rufus enfrentó un desafío de ingeniería sin precedentes: manejar millones de consultas por minuto y generar miles de millones de tokens en tiempo real, todo mientras se mantiene un SLA de latencia de 300 ms para tareas de QP y minimizando el consumo de energía. Estas demandas requirieron un replanteamiento fundamental de Cómo se implementan LLM a escala conquistando los cuellos de botella de costo y rendimiento. Los desafíos clave del día primario incluyeron:

• Escala masiva: Servir millones de tokens por minuto a los clientes en todo el mundo, con aumentos máximos de tráfico que se esfuerzan incluso los sistemas más robustos.
• SLAS estricto: Entrega de capacidad de respuesta en tiempo real con un límite de latencia difícil de 300 ms, asegurando una experiencia de cliente perfecta.
• Eficiencia de rentabilidad: Minimizar el costo de servir LLM a escala al tiempo que reduce el consumo de energía, un factor crítico para las operaciones sostenibles y económicas.

La generación de texto LLM tradicional es inherentemente ineficiente debido a su naturaleza secuencial. Cada generación de tokens requiere un paso de avance completo a través del modelo, lo que lleva a una alta latencia y subutilización de los recursos computacionales. Si bien se han propuesto técnicas como la decodificación especulativa para abordar estas ineficiencias, su complejidad y gastos generales de la capacitación han limitado su adopción.

AWS AI Chips y decodificación paralela

Para superar estos desafíos, Rufus adoptó la decodificación paralela, una técnica simple pero poderosa para acelerar la generación de LLM. Con la decodificación paralela, la dependencia secuencial se rompe, lo que hace que la generación autorregresiva sea más rápida. Este enfoque introduce cabezas de decodificación adicionales al modelo base que elimina la necesidad de un modelo de borrador separado para tokens especulados. Estas cabezas predicen múltiples tokens en paralelo para futuras posiciones antes de conocer los tokens anteriores, y esto mejora significativamente la eficiencia de la generación.

Para acelerar el rendimiento de la decodificación paralela para la inferencia en línea, Rufus utilizó una combinación de soluciones de AWS: Inferentia2 y Entrenamiento Chips ai, Amazon Elastic Compute Cloud (Amazon EC2) y Balancador de carga de la aplicación. Además, el equipo de Rufus se asoció con Nvidia para alimentar la solución utilizando Servidor de inferencia Triton de NVIDIAproporcionando capacidades para alojar el modelo utilizando chips AWS.

Para obtener la máxima eficiencia de decodificación paralela Núcleos de neuronas de AWStrabajamos en colaboración con el equipo de AWS Neuron para agregar el apoyo arquitectónico de la decodificación paralela en un Inferencia distribuida en neuronx (NXDI) Marco para un solo tamaño de lote.

Rufus extendió la base LLM con múltiples cabezas de decodificación. Estas cabezas son una pequeña capa de red neuronal y están entrenados utilizando las representaciones aprendidas del modelo base para predecir los siguientes tokens en paralelo. Estas cabezas están entrenadas junto con el modelo original, manteniendo el modelo base sin cambios. Porque los tokens no se generan secuencialmente, deben verificarse para asegurarse de que todos los tokens encajen. Para validar los tokens predichos por los cabezales de borrador, Rufus utiliza un mecanismo de atención basado en árboles para verificar e integrar tokens. Cada cabezal de borrador produce varias opciones para cada posición. Estas opciones se organizan en una estructura similar a un árbol para seleccionar la combinación más prometedora. Esto permite que se procesen múltiples tokens candidatos en paralelo, reduciendo la latencia y aumentando la utilización del núcleo de neuronas. La siguiente figura muestra un árbol disperso construido usando nuestro conjunto de calibración, con una profundidad de cuatro, lo que indica la participación de cuatro cabezas en el proceso de cálculo. Cada nodo representa un token de una predicción de Top-K de un cabezal de borrador, y los bordes representan las conexiones entre estos nodos.

Resultados del uso de decodificación paralela

Al integrar la decodificación paralela con los chips AWS AI y el marco NXDI, duplicamos la velocidad de la generación de texto en comparación con la decodificación autorregresiva, lo que lo convierte en una solución ideal para el entorno de alta demanda del día principal. Durante Amazon Prime Day 2024, Rufus demostró el poder de los chips AWS AI con impresionantes métricas de rendimiento:

• Generación dos veces más rápida: AWS AI Chips, optimizado para operaciones de decodificación paralela, habilitada duplicó la velocidad de generación de tokens en comparación con los procesadores tradicionales. Esta capacidad de procesamiento paralelo permitió predecir múltiples tokens futuros simultáneamente, entregando interacciones en tiempo real para millones de clientes.
• Costos de inferencia 50% más bajos: La combinación de chips AWS IA especialmente diseñados y la optimización de decodificación paralela eliminó los cálculos redundantes, reduciendo los costos de inferencia a la mitad mientras mantiene la calidad de la respuesta.
• Implementación simplificada: AWS AI Chips impulsó eficientemente los cabezales de decodificación paralelos del modelo, lo que permite la predicción simultánea del token sin la complejidad de administrar modelos de borradores separados. Esta sinergia arquitectónica simplificó la implementación al tiempo que ofrece una inferencia eficiente a escala.
• Escalabilidad perfecta: La combinación manejó el tráfico máximo sin comprometer la calidad de rendimiento y respuesta.

Estos avances no solo mejoraron la experiencia del cliente, sino que también mostraron el potencial del marco NXDI y la adaptabilidad de los chips AWS AI para optimizar el rendimiento de LLM a gran escala.

Cómo usar la decodificación paralela en entrenamiento e inferentia

La flexibilidad de NXDI combinada con chips de neuronas AWS lo convierte en una solución poderosa para la generación de texto LLM en producción. Ya sea que esté utilizando Trabajo o inferencia para inferencia, NXDI proporciona una interfaz unificada para implementar optimizaciones de decodificación paralela. Esta integración reduce la complejidad operativa y proporciona una ruta directa para las organizaciones que buscan implementar y escalar sus aplicaciones LLM de manera eficiente.

Puede explorar la técnica de decodificación paralela como Medusa Para acelerar sus flujos de trabajo de inferencia en instancias INF2 o TRN1. Para comenzar, necesitará un modelo compatible con Medusa (como la inferencia de texto de texto/Mistral-7B-InNSTRUCT-V0.2-MEDUSA) y una configuración de árbol Medusa. Habilitar Medusa colocando is_medusa=Trueconfigurando su medusa_speculation_length, num_medusa_headsy especificando tu medusa_tree. Al usar la API de Generate () de Huggingface, configure la assistant_model a tu modelo objetivo. Tenga en cuenta que Medusa actualmente admite solo un tamaño por lotes de 1.

 def load_json_file(json_path):
    with open(json_path, "r") as f:
        return json.load(f)
medusa_tree = load_json_file("medusa_mc_sim_7b_63.json")
neuron_config = NeuronConfig(
    is_medusa=True,
    medusa_speculation_length=64,
    num_medusa_heads=4,
    medusa_tree=medusa_tree
)

Conclusión

Prime Day es un testimonio del poder de la innovación para superar los desafíos técnicos. Al usar los chips AWS AI, Rufus no solo cumplió las estrictas demandas del día principal, sino que también estableció un nuevo estándar para la eficiencia de LLM. A medida que los LLM continúan evolucionando, los marcos como NXDI desempeñarán un papel crucial para hacerlos más accesibles, escalables y rentables. Estamos emocionados de ver cómo la comunidad se basará en la Fundación NXDI y los chips AWS AI para desbloquear nuevas posibilidades para las aplicaciones LLM. ¡Pruébelo hoy y experimente la diferencia para usted!

Expresiones de gratitud

Extendemos nuestra gratitud al equipo de AWS Annapurna responsable de los chips AWS AI y el desarrollo del marco. Un agradecimiento especial a los investigadores e ingenieros cuyas contribuciones hicieron posible este logro. Las mejoras en la latencia, el rendimiento y la eficiencia de rentabilidad logradas con la decodificación paralela en comparación con la decodificación autorregresiva han establecido un nuevo punto de referencia para las implementaciones de LLM a escala.

Sobre los autores

Shruti Dubey es ingeniero de software en el equipo de búsqueda principal de Amazon, donde optimiza los sistemas de inferencias de LLM para que la IA sea más rápida y escalable. Le apasiona la IA generativa y le encanta convertir la investigación de vanguardia en el impacto del mundo real. Fuera del trabajo, la encontrarás corriendo, leyendo o tratando de convencer a su perro de que ella es la jefa.

Shivangi Agarwal Es una científica aplicada en el equipo de video principal de Amazon, donde se enfoca en optimizar la inferencia de LLM y desarrollar sistemas de clasificación inteligentes para videos principales utilizando señales de nivel de consulta. Ella está impulsada por una pasión por construir una IA escalable y eficiente que ofrece un impacto en el mundo real. Cuando no esté trabajando, es probable que la encuentre atrapando una buena película, descubrir nuevos lugares o mantenerse al día con su aventurero hijo de 3 años.

Sukhdeep Singh Kharbanda es un gerente de ciencia aplicada en Amazon Core Search. En su papel actual, Sukhdeep está liderando el equipo de inferencia de Amazon para construir soluciones de optimización de inferencia de Genai y sistema de inferencia a escala para una inferencia rápida a bajo costo. El trabajo externo, le gusta jugar con su hijo y cocinar diferentes cocinas.

Rahul goutam Es un gerente de ciencia aplicada en Amazon Core Search, donde lleva a equipos de científicos e ingenieros a construir soluciones de IA escalables que alimentan las experiencias de compra flexibles e intuitivas. Cuando está fuera del reloj, le gusta caminar por un sendero o esquiar en uno.

Yang Zhou es un ingeniero de software que trabaja en la creación y optimización de sistemas de aprendizaje automático. Su enfoque reciente es mejorar el rendimiento y la eficiencia rentable de la inferencia generativa de IA. Más allá del trabajo, le gusta viajar y recientemente ha descubierto una pasión por correr largas distancias.

RJ es ingeniero dentro de Amazon. Construye y optimiza los sistemas para sistemas distribuidos para la capacitación y trabaja en la optimización de los sistemas de adopción para reducir la latencia para la inferencia de ML. El trabajo externo, está explorando el uso de IA generativa para construir recetas de alimentos.

Parque de james es un arquitecto principal de soluciones de especialistas en aprendizaje automático en Amazon Web Services. Trabaja con Amazon para diseñar, construir e implementar soluciones tecnológicas en AWS, y tiene un interés particular en la IA y el aprendizaje automático. En su tiempo libre le gusta buscar nuevas culturas, experiencias y mantenerse al día con las últimas tendencias tecnológicas.