Optimización de la memoria para inferencia y ajuste fino de modelos de lenguaje grande

Los modelos de lenguaje grande (LLM) como GPT-4, Bloom y LLaMA han logrado capacidades notables al escalar hasta miles de millones de parámetros. Sin embargo, desplegar estos modelos masivos para inferencia o ajuste fino es un desafío debido a sus enormes requisitos de memoria. En este blog técnico, exploraremos técnicas para estimar y optimizar el consumo de memoria durante la inferencia y el ajuste fino de LLM en varios entornos de hardware.

Comprensión de los requisitos de memoria

La memoria necesaria para cargar un LLM se determina principalmente por la cantidad de parámetros y la precisión numérica utilizada para almacenar los parámetros. Una regla general es:

• Cargar un modelo con X mil millones de parámetros requiere aproximadamente 4X GB de VRAM en precisión de 32 bits float
• Cargar un modelo con X mil millones de parámetros requiere aproximadamente 2X GB de VRAM en precisión de 16 bits bfloat16/float16

Por ejemplo, cargar el modelo GPT-3 de 175 mil millones de parámetros requeriría aproximadamente 350 GB de VRAM en precisión bfloat16. Actualmente, las GPU más grandes disponibles comercialmente, como la NVIDIA A100 y H100, ofrecen solo 80 GB de VRAM, lo que hace necesarias técnicas de paralelismo de tensor y paralelismo de modelo.

Durante la inferencia, la huella de memoria está dominada por los parámetros del modelo y los tensores de activación temporales producidos. Una estimación de alto nivel para el uso máximo de memoria durante la inferencia es la suma de la memoria necesaria para cargar los parámetros del modelo y la memoria para las activaciones.

Cuantificación de la memoria de inferencia

Veamos los requisitos de memoria para la inferencia utilizando el modelo OctoCode, que tiene alrededor de 15 mil millones de parámetros en formato bfloat16 (~ 31 GB). Utilizaremos la biblioteca Transformers para cargar el modelo y generar texto:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Pregunta: Por favor, escriba una función de Python para convertir bytes a gigabytes.\n\nRespuesta:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Salida:

El uso máximo de memoria de la GPU es de aproximadamente 29 GB, lo que coincide con nuestra estimación de 31 GB para cargar los parámetros del modelo en formato bfloat16.

Optimización de la memoria de inferencia con cuantización

Si bien bfloat16 es la precisión común utilizada para entrenar LLM, los investigadores han encontrado que cuantizar los pesos del modelo a tipos de datos de precisión más baja como enteros de 8 bits (int8) o enteros de 4 bits puede reducir significativamente el uso de memoria con una pérdida mínima de precisión para tareas de inferencia como la generación de texto.

Veamos los ahorros de memoria de la cuantización de 8 bits y 4 bits del modelo OctoCode:

</div>
# Cuantización de 8 bits
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Salida:

# Cuantización de 4 bits
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Salida:

Con la cuantización de 8 bits, el requisito de memoria disminuye de 31 GB a 15 GB, mientras que la cuantización de 4 bits lo reduce aún más a solo 9,5 GB. Esto permite ejecutar el modelo OctoCode de 15 mil millones de parámetros en GPUs de consumo como la RTX 3090 (24 GB de VRAM).

Sin embargo, tenga en cuenta que la cuantización más agresiva como la de 4 bits puede llevar a una degradación de la precisión en comparación con la precisión de 8 bits o bfloat16. Hay un equilibrio entre ahorro de memoria y precisión que los usuarios deben evaluar para su caso de uso.

La cuantización es una técnica poderosa que puede permitir el despliegue de LLM en entornos con recursos limitados como instancias en la nube, dispositivos de borde o incluso teléfonos móviles al reducir drásticamente la huella de memoria.

Estimación de la memoria para el ajuste fino

Mientras que la cuantización se utiliza principalmente para la inferencia eficiente, técnicas como el paralelismo de tensor y el paralelismo de modelo son cruciales para gestionar los requisitos de memoria durante el entrenamiento o ajuste fino de modelos de lenguaje grande.

El consumo máximo de memoria durante el ajuste fino es generalmente 3-4 veces mayor que el de la inferencia debido a los requisitos de memoria adicionales para:

• Gradientes
• Estados del optimizador
• Activaciones de la pasada hacia adelante almacenadas para la retropropagación

Una estimación conservadora es que el ajuste fino de un LLM con X mil millones de parámetros requiere alrededor de 4 * (2X) = 8X GB de VRAM en precisión bfloat16.

Por ejemplo, el ajuste fino del modelo LLaMA de 7 mil millones de parámetros requeriría aproximadamente 7 * 8 = 56 GB de VRAM por GPU en precisión bfloat16. Esto supera la capacidad de memoria de las GPUs actuales, lo que hace necesarias técnicas de ajuste fino distribuido.

Técnicas de ajuste fino distribuido

Se han propuesto varios métodos de ajuste fino distribuido para superar las limitaciones de memoria de la GPU para modelos grandes:

1. Paralelismo de datos: El enfoque clásico de paralelismo de datos replica el modelo completo en varias GPUs mientras divide y distribuye los lotes de datos de entrenamiento. Esto reduce el tiempo de entrenamiento linealmente con la cantidad de GPUs, pero no reduce el requisito de memoria máximo en cada GPU.
2. ZeRO Stage 3: Una forma avanzada de paralelismo de datos que divide los parámetros del modelo, los gradientes y los estados del optimizador en las GPUs. Reduce la memoria en comparación con el paralelismo de datos clásico al mantener solo los datos particionados necesarios en cada GPU durante las diferentes fases del entrenamiento.
3. Paralelismo de tensor: En lugar de replicar el modelo, el paralelismo de tensor divide los parámetros del modelo en filas o columnas y los distribuye en las GPUs. Cada GPU opera en un conjunto particionado de parámetros, gradientes y estados del optimizador, lo que lleva a un ahorro de memoria sustancial.
4. Paralelismo de tubería: Esta técnica divide las capas del modelo en diferentes GPUs/trabajadores, con cada dispositivo ejecutando un subconjunto de las capas. Las activaciones se pasan entre los trabajadores, lo que reduce la memoria máxima pero aumenta la sobrecarga de comunicación.

Estimar el uso de memoria para estos métodos distribuidos es no trivial, ya que la distribución de parámetros, gradientes, activaciones y estados del optimizador varía entre técnicas. Además, diferentes componentes como el cuerpo del transformador y la cabeza de modelado de lenguaje pueden exhibir diferentes comportamientos de asignación de memoria.

La solución LLMem

Los investigadores han propuesto recientemente LLMem, una solución que estima con precisión el consumo de memoria de la GPU cuando se aplican métodos de ajuste fino distribuido a LLM en varias GPUs.

Estimación del uso de memoria de la GPU para el ajuste fino de LLM preentrenados

LLMem considera factores como la recombinación de parámetros antes del cálculo (ZeRO Stage 3), la recopilación de salidas en la pasada hacia atrás (paralelismo de tensor) y las diferentes estrategias de asignación de memoria para el cuerpo del transformador y la cabeza de modelado de lenguaje.

Los resultados experimentales muestran que LLMem puede estimar el uso máximo de memoria de la GPU para el ajuste fino de LLM en una sola GPU con tasas de error de hasta 1,6%, superando la tasa de error promedio de 42,6% de DNNMem. Cuando se aplican métodos de ajuste fino distribuido a LLM con más de mil millones de parámetros en varias GPUs, LLMem logra una tasa de error promedio de 3,0%.