Generación Asistida una nueva dirección hacia la generación de texto de baja latencia

'Generación Asistida nueva dirección para texto de baja latencia.'

Los modelos de lenguaje grandes están de moda en estos días, con muchas empresas invirtiendo recursos significativos para escalarlos y desbloquear nuevas capacidades. Sin embargo, como humanos con una atención cada vez menor, también nos disgusta su lenta respuesta. La latencia es fundamental para una buena experiencia del usuario, y a menudo se utilizan modelos más pequeños a pesar de su menor calidad (por ejemplo, en la finalización de código).

¿Por qué la generación de texto es tan lenta? ¿Qué impide desplegar modelos de lenguaje grandes y de baja latencia sin ir a la quiebra? En esta publicación de blog, repasaremos los cuellos de botella para la generación de texto autoregresiva e introduciremos un nuevo método de decodificación para abordar el problema de la latencia. Verá que al utilizar nuestro nuevo método, la generación asistida, ¡puede reducir la latencia hasta 10 veces en hardware común!

Comprendiendo la latencia en la generación de texto

El núcleo de la generación de texto moderna es fácil de entender. Veamos la pieza central, el modelo de aprendizaje automático (ML). Su entrada contiene una secuencia de texto, que incluye el texto generado hasta ahora y componentes específicos del modelo (por ejemplo, Whisper también tiene una entrada de audio). El modelo toma la entrada y realiza un pase hacia adelante: la entrada se alimenta al modelo y se pasa secuencialmente a lo largo de sus capas hasta que se predicen las probabilidades no normalizadas para el siguiente token (también conocido como logits). Un token puede consistir en palabras completas, subpalabras o incluso caracteres individuales, dependiendo del modelo. El ilustrado GPT-2 es una excelente referencia si desea profundizar en esta parte de la generación de texto.

Un pase hacia adelante del modelo le proporciona los logits para el próximo token, que puede manipular libremente (por ejemplo, establecer la probabilidad de palabras o secuencias indeseables en 0). El siguiente paso en la generación de texto es seleccionar el siguiente token a partir de estos logits. Las estrategias comunes incluyen elegir el token más probable, conocido como decodificación codiciosa, o muestrear de su distribución, también llamado muestreo multinomial. Encadenar pases hacia adelante del modelo con la selección del siguiente token de manera iterativa le proporciona la generación de texto. Esta explicación es solo la punta del iceberg cuando se trata de métodos de decodificación; consulte nuestra publicación de blog sobre generación de texto para una exploración en profundidad.

• Presentando RWKV – Una RNN con las ventajas de un transformador
• Más pequeño es mejor Q8-Chat, una experiencia eficiente de IA generativa en Xeon
• Deduplicación a gran escala detrás de BigCode

A partir de la descripción anterior, el cuello de botella de latencia en la generación de texto es claro: ejecutar un pase hacia adelante del modelo para modelos grandes es lento, y es posible que necesite hacer cientos de ellos en una secuencia. Pero profundicemos más: ¿por qué los pases hacia adelante son lentos? Los pases hacia adelante están típicamente dominados por multiplicaciones de matrices y, después de una visita rápida a la sección correspondiente de Wikipedia, puede decir que el ancho de banda de la memoria es la limitación en esta operación (por ejemplo, desde la RAM de la GPU hasta los núcleos de cálculo de la GPU). En otras palabras, el cuello de botella en el pase hacia adelante proviene de cargar los pesos de las capas del modelo en los núcleos de cálculo de su dispositivo, no de realizar los cálculos en sí.

En este momento, tiene tres vías principales que puede explorar para aprovechar al máximo la generación de texto, todas ellas abordando el rendimiento del pase hacia adelante del modelo. Primero, tiene las optimizaciones del modelo específicas del hardware. Por ejemplo, su dispositivo puede ser compatible con Flash Attention, que acelera la capa de atención mediante una reordenación de las operaciones, o cuantización INT8, que reduce el tamaño de los pesos del modelo.

En segundo lugar, cuando sabe que recibirá solicitudes concurrentes de generación de texto, puede agrupar las entradas y aumentar masivamente el rendimiento con una pequeña penalización en la latencia. Los pesos de las capas del modelo cargados en el dispositivo ahora se utilizan en varias filas de entrada en paralelo, lo que significa que obtendrá más tokens para aproximadamente la misma carga de ancho de banda de memoria. La trampa con la agrupación es que necesita memoria adicional del dispositivo (o descargar la memoria en otro lugar). En el extremo de este espectro, puede ver proyectos como FlexGen, que optimizan el rendimiento a expensas de la latencia.

# Ejemplo que muestra el impacto de la generación agrupada. Dispositivo de medición: RTX3090
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2").to("cuda")
inputs = tokenizer(["Hola mundo"], return_tensors="pt").to("cuda")

def imprimir_tokens_por_segundo(tamaño_lote):
    nuevos_tokens = 100
    tiempo_acumulado = 0

    # calentamiento
    model.generate(
        **inputs, do_sample=True, max_new_tokens=nuevos_tokens, num_return_sequences=tamaño_lote
    )

    for _ in range(10):
        inicio = time.time()
        model.generate(
            **inputs, do_sample=True, max_new_tokens=nuevos_tokens, num_return_sequences=tamaño_lote
        )
        tiempo_acumulado += time.time() - inicio
    print(f"Tokens por segundo: {nuevos_tokens * tamaño_lote * 10 / tiempo_acumulado:.1f}")

imprimir_tokens_por_segundo(1)   # Tokens por segundo: 418.3
imprimir_tokens_por_segundo(64)  # Tokens por segundo: 16266.2 (~39 veces más tokens por segundo)

Finalmente, si tienes varios dispositivos disponibles, puedes distribuir la carga de trabajo utilizando la Paralelización de Tensores y obtener una latencia más baja. Con la Paralelización de Tensores, divides la carga de ancho de banda de la memoria entre varios dispositivos, pero ahora debes considerar los cuellos de botella de comunicación entre dispositivos además del costo monetario de ejecutar varios dispositivos. Los beneficios dependen en gran medida del tamaño del modelo: los modelos que caben fácilmente en un solo dispositivo de consumo tienen beneficios muy limitados. Tomando los resultados de esta publicación del blog de DeepSpeed, puedes ver que puedes distribuir un modelo de 17 mil millones de parámetros en 4 GPUs para reducir la latencia en un 1.5x (Figura 7).

Estas tres mejoras se pueden utilizar en conjunto, lo que resulta en soluciones de alto rendimiento. Sin embargo, después de aplicar optimizaciones específicas del hardware, hay opciones limitadas para reducir la latencia, y las opciones existentes son costosas. ¡Arreglemos eso!

Avance del decodificador de lenguaje, revisado

Has leído anteriormente que cada paso hacia adelante del modelo produce los logits para el siguiente token, pero eso en realidad es una descripción incompleta. Durante la generación de texto, la iteración típica consiste en que el modelo recibe como entrada el último token generado, además de los cálculos internos en caché para todas las entradas anteriores, y devuelve los logits del siguiente token. El caché se utiliza para evitar cálculos redundantes, lo que resulta en avances más rápidos, pero no es obligatorio (y se puede utilizar parcialmente). Cuando se desactiva el caché, la entrada contiene la secuencia completa de tokens generados hasta ahora y la salida contiene los logits correspondientes al siguiente token para todas las posiciones en la secuencia. Los logits en la posición N corresponden a la distribución para el siguiente token si la entrada consistiera en los primeros N tokens, ignorando todos los tokens posteriores en la secuencia. En el caso particular de la decodificación codiciosa, si pasas la secuencia generada como entrada y aplicas el operador argmax a los logits resultantes, obtendrás la secuencia generada de vuelta.

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

inputs = tok(["The"], return_tensors="pt")
generated = model.generate(**inputs, do_sample=False, max_new_tokens=10)
forward_confirmation = model(generated).logits.argmax(-1)

# Excluimos los extremos opuestos de cada secuencia: el avance devuelve
# los logits para el siguiente token, por lo que se desplaza una posición.
print(generated[0, 1:].tolist() == forward_confirmation[0, :-1].tolist())  # True

Esto significa que puedes utilizar un avance del modelo para un propósito diferente: además de alimentar algunos tokens para predecir el siguiente, también puedes pasar una secuencia al modelo y verificar si el modelo generaría esa misma secuencia (o parte de ella).

Consideremos por un momento que tienes acceso a un modelo oráculo mágico sin latencia que genera la misma secuencia que tu modelo, para cualquier entrada dada. Por razones de argumento, no se puede utilizar directamente, está limitado a ser un asistente de tu procedimiento de generación. Utilizando la propiedad descrita anteriormente, podrías usar este modelo asistente para obtener tokens de salida candidatos seguidos de un avance con tu modelo para confirmar que son correctos. En este escenario utópico, la latencia de la generación de texto se reduciría de O(n) a O(1), siendo n el número de tokens generados. Para generaciones largas, estamos hablando de varias órdenes de magnitud.

Dando un paso hacia la realidad, supongamos que el modelo asistente ha perdido sus propiedades oráculo. Ahora es un modelo sin latencia que obtiene algunos de los tokens candidatos incorrectos, según tu modelo. Debido a la naturaleza autoregresiva de la tarea, tan pronto como el asistente obtiene un token incorrecto, todos los candidatos subsiguientes deben ser invalidados. Sin embargo, eso no impide que consultes nuevamente al asistente, después de corregir el token incorrecto con tu modelo, y repitas este proceso de manera iterativa. Incluso si el asistente falla en algunos tokens, la generación de texto tendría una latencia de un orden de magnitud menor que en su forma original.

Obviamente, no existen modelos asistentes sin latencia. Sin embargo, es relativamente fácil encontrar un modelo que aproxime las salidas de generación de texto de otro modelo; las versiones más pequeñas de la misma arquitectura entrenadas de manera similar a menudo cumplen con esta propiedad. Además, cuando la diferencia en los tamaños de los modelos se vuelve significativa, el costo de usar el modelo más pequeño como asistente se convierte en un factor secundario después de tener en cuenta los beneficios de omitir algunos avances. Ahora comprendes el núcleo de la generación asistida.

Decodificación codiciosa con generación asistida

La generación asistida es un acto de equilibrio. Quieres que el asistente genere rápidamente una secuencia candidata mientras sea lo más preciso posible. Si el asistente tiene poca calidad, obtienes el costo de usar el modelo asistente con poco o ningún beneficio. Por otro lado, optimizar la calidad de las secuencias candidatas puede implicar el uso de asistentes lentos, lo que resulta en una desaceleración neta. Aunque no podemos automatizar la selección del modelo asistente para ti, hemos incluido un requisito adicional y una heurística para asegurar que el tiempo gastado con el asistente se mantenga bajo control.

En primer lugar, el requisito: el asistente debe tener el mismo tokenizador exacto que su modelo. Si este requisito no estuviera en su lugar, se agregarían pasos costosos de decodificación y re-codificación de tokens. Además, estos pasos adicionales tendrían que ocurrir en la CPU, lo que a su vez puede requerir transferencias de datos lentas entre dispositivos. El uso rápido del asistente es fundamental para que se muestren los beneficios de la generación asistida.

Por último, la heurística. Hasta este punto, probablemente haya notado las similitudes entre la película Inception y la generación asistida: después de todo, se está ejecutando la generación de texto dentro de la generación de texto. Habrá un pase hacia adelante del modelo asistente por cada token candidato, y sabemos que los pases hacia adelante son costosos. Si bien no se puede saber de antemano la cantidad de tokens que el modelo asistente obtendrá correctamente, se puede realizar un seguimiento de esta información y usarla para limitar la cantidad de tokens candidatos solicitados al asistente: algunas secciones de la salida son más fáciles de anticipar que otras.

Para resumir todo, aquí está nuestra implementación original del bucle de generación asistida (código):

1. Usar decodificación voraz para generar un cierto número de tokens candidatos con el modelo asistente, produciendo candidates. El número de tokens candidatos producidos se inicializa en 5 la primera vez que se llama a la generación asistida.
2. Usando nuestro modelo, hacer un pase hacia adelante con candidates, obteniendo logits.
3. Usar el método de selección de tokens (.argmax() para búsqueda voraz o .multinomial() para muestreo) para obtener los next_tokens de logits.
4. Comparar next_tokens con candidates y obtener la cantidad de tokens coincidentes. Recuerde que esta comparación debe hacerse con causalidad de izquierda a derecha: después de la primera coincidencia incorrecta, todos los candidatos se invalidan.
5. Usar la cantidad de coincidencias para dividir las cosas y descartar variables relacionadas con tokens candidatos no confirmados. En esencia, en next_tokens, mantener los tokens coincidentes más el primer token divergente (que nuestro modelo genera a partir de una subsecuencia de candidatos válidos).
6. Ajustar el número de tokens candidatos que se producirán en la próxima iteración: nuestra heurística original lo aumenta en 2 si TODOS los tokens coinciden y lo disminuye en 1 en caso contrario.

Hemos diseñado la API en 🤗 Transformers de forma que este proceso sea sin complicaciones para usted. ¡Lo único que necesita hacer es pasar el modelo asistente utilizando el nuevo argumento de palabra clave assistant_model y aprovechar las ganancias de latencia! En el momento de la publicación de esta publicación de blog, la generación asistida está limitada a un tamaño de lote de 1.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# ['Alice y Bob están sentados en un bar. Alice está bebiendo una cerveza y Bob está bebiendo a']

¿Vale la pena la complejidad interna adicional? Echemos un vistazo a los números de latencia para el caso de decodificación voraz (los resultados para el muestreo se encuentran en la siguiente sección), considerando un tamaño de lote de 1. Estos resultados se obtuvieron directamente de 🤗 Transformers sin ninguna optimización adicional, por lo que debería poder reproducirlos en su configuración.

Al observar los números recopilados, vemos que la generación asistida puede ofrecer reducciones significativas en la latencia en diversos entornos, pero no es una solución milagrosa; debe evaluarla antes de aplicarla a su caso de uso. Podemos concluir que la generación asistida:

1. 🤏 Requiere acceso a un modelo asistente que sea al menos un orden de magnitud más pequeño que su modelo (cuanto mayor sea la diferencia, mejor);
2. 🚀 Obtiene mejoras de velocidad de hasta 3 veces en presencia de INT8 y hasta 2 veces en otros casos, cuando el modelo cabe en la memoria de la GPU;
3. 🤯 Si está trabajando con modelos que no caben en su GPU y está confiando en la transferencia de memoria, puede ver mejoras de velocidad de hasta 10 veces;
4. 📄 Destaca en tareas basadas en la entrada, como reconocimiento automático del habla o resumen.

Muestra con generación asistida

La decodificación voraz es adecuada para tareas basadas en la entrada (reconocimiento automático del habla, traducción, resumen, …) o búsqueda de conocimiento factual. Las tareas abiertas que requieren altos niveles de creatividad, como la mayoría de los usos de un modelo de lenguaje como chatbot, deben usar la generación de muestras en su lugar. La generación asistida está diseñada naturalmente para la decodificación voraz, ¡pero eso no significa que no puedas usar la generación asistida con muestreo multinomial!

Obtener muestras de una distribución de probabilidad para el siguiente token hará que nuestro asistente voraz falle con más frecuencia, reduciendo sus beneficios de latencia. Sin embargo, podemos controlar qué tan aguda es la distribución de probabilidad para los siguientes tokens, utilizando el coeficiente de temperatura que se encuentra en la mayoría de las aplicaciones basadas en muestreo. En un extremo, con temperaturas cercanas a 0, el muestreo aproximará la decodificación voraz, favoreciendo el token más probable. En el otro extremo, con la temperatura establecida en valores mucho mayores que 1, el muestreo será caótico, seleccionando de una distribución uniforme. Las temperaturas bajas son, por lo tanto, más favorables para su modelo de asistente, retienen la mayoría de los beneficios de latencia de la generación asistida, como podemos ver a continuación.

¿Por qué no lo ves por ti mismo, para que puedas tener una idea de la generación asistida?

Direcciones futuras

La generación asistida muestra que las estrategias modernas de generación de texto están listas para la optimización. Comprender que actualmente es un problema limitado por la memoria, no por la potencia de cálculo, nos permite aplicar heurísticas simples para aprovechar al máximo el ancho de banda de memoria disponible, aliviando el cuello de botella. Creemos que una mayor refinación del uso de modelos de asistentes nos proporcionará reducciones aún mayores en la latencia; por ejemplo, es posible que podamos omitir algunas pasadas adicionales si solicitamos al asistente que genere varias continuaciones candidatas. Naturalmente, la liberación de modelos pequeños de alta calidad para ser utilizados como asistentes será fundamental para comprender y amplificar los beneficios.

Inicialmente lanzado bajo nuestra biblioteca de 🤗 Transformers, para ser utilizado con la función .generate(), esperamos ofrecerlo en todo el universo de Hugging Face. Su implementación también es completamente de código abierto, así que si estás trabajando en generación de texto y no estás utilizando nuestras herramientas, siéntete libre de usarlo como referencia.

Finalmente, la generación asistida plantea una pregunta crucial en la generación de texto. El campo ha estado evolucionando con la restricción de que todos los nuevos tokens son el resultado de una cantidad fija de cálculos, para un modelo dado. Un token por cada pasada hacia adelante homogénea, de manera puramente autoregresiva. Esta publicación de blog refuerza la idea de que no debería ser así: grandes secciones de la salida generada también pueden ser generadas de manera igual por modelos que son una fracción del tamaño. Para eso, necesitaremos nuevas arquitecturas de modelos y métodos de decodificación. ¡Estamos emocionados por ver lo que depara el futuro!

Trabajos relacionados

Después del lanzamiento original de esta publicación de blog, me di cuenta de que otros trabajos han explorado el mismo principio central (usar una pasada hacia adelante para validar continuaciones más largas). En particular, echa un vistazo a los siguientes trabajos:

• Decodificación paralela por bloques, de Google Brain
• Muestreo especulativo, de DeepMind

Cita

@misc {gante2023assisted,
    author       = { {Joao Gante} },
    title        = { Generación asistida: una nueva dirección hacia la generación de texto de baja latencia },
    year         = 2023,
    url          = { https://huggingface.co/blog/assisted-generation },
    doi          = { 10.57967/hf/0638 },
    publisher    = { Hugging Face Blog }
}

Agradecimientos

Me gustaría agradecer a Sylvain Gugger, Nicolas Patry y Lewis Tunstall por compartir muchas sugerencias valiosas para mejorar esta publicación de blog. Por último, felicitaciones a Chunte Lee por diseñar la hermosa portada que puedes ver en nuestra página web.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!