En los últimos años, hemos observado un gran aumento en el tamaño de los modelos de lenguaje de grandes dimensiones (LLMs) que se utilizan para resolver tareas de procesamiento del lenguaje natural (NLP) como la respuesta a preguntas y la resumir texto. Los modelos más grandes, con más parámetros, que están en el orden de cientos de miles de millones en el momento de escribir esto, tienden a producir mejores resultados. Por ejemplo, Llama-3-70B puntúa mejor que su versión más pequeña de 8 mil millones de parámetros en métricas como la comprensión de lectura (SQuAD 85.6 comparado con 76.4). Así, los clientes a menudo experimentan con modelos más grandes y nuevos para construir productos basados en machine learning que aporten valor.
Sin embargo, cuanto más grande es el modelo, más exigente es computacionalmente y más alto es el costo de despliegue. Por ejemplo, en AWS Trainium, Llama-3-70B tiene una latencia mediana por token de 21.4 ms, mientras que Llama-3-8B toma 4.7 ms. De manera similar, Llama-2-70B tiene una latencia mediana por token de 20.6 ms, mientras que Llama-2-7B toma 3.7 ms. Los clientes deben considerar el rendimiento para asegurar que cumplen con las necesidades de sus usuarios. En este blog, exploramos cómo el muestreo especulativo puede ayudar a hacer que la inferencia de modelos de lenguaje grande sea más eficiente en términos de computo y costos en AWS Inferentia y Trainium. Esta técnica mejora el rendimiento de inferencia de LLM y la latencia de salida de token (TPOT).
Los modelos de lenguaje modernos se basan en la arquitectura transformer. Las instrucciones de entrada se procesan primero usando una técnica llamada codificación de contexto, que se ejecuta rápidamente porque es paralelizable. Luego, generamos tokens de forma autorregresiva, es decir, secuencialmente. Para generar N tokens de salida, necesitamos N ejecuciones en serie del decodificador. Cuanto mayor es el modelo, como Llama-3-70B, mayor es el tiempo que tarda en generar el siguiente token.
Desde una perspectiva computacional, la generación de tokens en los LLMs es un proceso limitado por el ancho de banda de la memoria. Cuanto más grande es el modelo, más probable es que esperemos transferencias de memoria, resultando en una infrautilización de las unidades de cálculo y no aprovechando completamente las operaciones de punto flotante disponibles (FLOPS).
El muestreo especulativo es una técnica que mejora la eficiencia computacional para ejecutar inferencias con LLMs, manteniendo la precisión. Funciona utilizando un modelo de borrador más pequeño y rápido para generar múltiples tokens, que luego son verificados por un modelo objetivo más grande y lento. Este paso de verificación procesa múltiples tokens en una sola pasada en lugar de secuencialmente y es más eficiente computacionalmente. Aumentar el número de tokens procesados en paralelo incrementa la intensidad de cálculo porque una mayor cantidad de tokens pueden multiplicarse con el mismo tensor de peso, proporcionando un mejor rendimiento comparado con la ejecución no especulativa, que generalmente está limitada por el ancho de banda de la memoria, derivando en una mejor utilización de los recursos de hardware.
El proceso especulativo involucra una ventana ajustable k, donde el modelo objetivo proporciona un token garantizado correcto, y el modelo borrador especula los próximos k-1 tokens. Si los tokens del modelo borrador son aceptados, el proceso se acelera. Si no, el modelo objetivo toma el control, asegurando la precisión.
Por ejemplo, un caso donde todos los tokens especulados son aceptados resulta en un procesamiento más rápido. El modelo objetivo proporciona un token de salida garantizado, y el modelo de borrador se ejecuta múltiples veces para producir una secuencia de posibles tokens de salida. Estos son verificados por el modelo objetivo y luego aceptados por un método probabilístico.
Por otro lado, un caso donde algunos de los tokens son rechazados implica que obtendremos menos tokens de salida y repetiremos este proceso más veces para completar la respuesta, resultando en un procesamiento más lento en general. Ajustando el tamaño de la ventana k y entendiendo cuándo los modelos borrador y objetivo son propensos a producir resultados similares, podemos maximizar los beneficios del muestreo especulativo.
Demostramos cómo funciona el muestreo especulativo en instancias Amazon EC2 Inf2 alimentadas por Inferentia2 y en instancias EC2 Trn1 alimentadas por Trainium. Usamos un ejemplo donde generamos texto más rápido con Llama-2-70B utilizando Llama-2-7B como modelo de borrador. Aunque el ejemplo se basa en modelos Llama-2, se puede seguir un proceso similar para los modelos Llama-3.
La carga de los modelos Llama-2 con tipo de datos bfloat16 puede hacerse de manera estándar ajustando el parámetro n_positions, que representa la longitud máxima de secuencia permitida para la generación. Para muestreo especulativo, sólo se soporta batch_size de 1. Los modelos combinados necesitan casi 200 GB de memoria del dispositivo para los pesos más memoria adicional en el orden de gigabytes para los cachés de clave-valor (KV).
A medida que más desarrolladores buscan incorporar LLMs en sus aplicaciones, enfrentan la elección entre usar modelos más grandes, costosos y lentos que brindan resultados de mayor calidad, o usar modelos más pequeños, menos costosos y rápidos que podrían reducir la calidad de las respuestas. Ahora, con los chips de inteligencia artificial (AI) de AWS y el muestreo especulativo, los desarrolladores no tienen que tomar esa decisión. Pueden aprovechar las salidas de alta calidad de modelos grandes y la velocidad de los modelos pequeños.
En esta publicación, hemos mostrado cómo podemos acelerar la inferencia de modelos grandes, como Llama-2-70B, utilizando una nueva función llamada muestreo especulativo. Para probarlo usted mismo, puede revisar el ejemplo de muestreo especulativo y ajustar el prompt de entrada y el parámetro k para ver los resultados obtenidos. Para casos de uso más avanzados, puede desarrollar su propia implementación del aceptador de tokens. Para obtener más información sobre cómo ejecutar sus modelos en instancias Inferentia y Trainium, puede leer la documentación de AWS Neuron y visitar el canal de AWS Neuron en repost.aws para discutir sus experimentos con la comunidad de AWS Neuron y compartir ideas.
