Optimización de la memoria para la inferencia y el ajuste de modelos de lenguaje grandes

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

Comprender los requisitos de memoria

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

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

Por ejemplo, cargar el modelo GPT-175 con parámetro 3B requeriría aproximadamente 350 GB de VRAM con precisión bfloat16. A día de hoy, las GPU más grandes disponibles comercialmente, como NVIDIA A100 y H100, ofrecen solo 80 GB de VRAM, lo que requiere técnicas de paralelismo tensorial y de paralelismo de modelos.

Durante la inferencia, la huella de la memoria está dominada por los parámetros del modelo y los tensores de activación temporal 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.

Cuantificar la memoria de inferencia

Cuantifiquemos los requisitos de memoria para la inferencia utilizando el modelo OctoCode, que tiene alrededor de 15 mil millones de parámetros en formato bfloat16 (aproximadamente 31 GB). Usaremos Biblioteca de transformadores para cargar el modelo y generar texto:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]

def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024

bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Salida:

El uso máximo de memoria de la GPU es de alrededor de 29 GB, lo que se alinea 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 cuantificación

Si bien bfloat16 es la precisión común utilizada para entrenar LLM, los investigadores han descubierto que cuantificar los pesos del modelo a tipos de datos de menor precisión, 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 generación de texto.

Veamos el ahorro de memoria de la cuantificación de 8 y 4 bits del modelo OctoCode:

</div>
# 8-bit quantization
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:

# 4-bit quantization
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 cuantificación de 8 bits, el requisito de memoria se reduce de 31 GB a 15 GB, mientras que la cuantificación de 4 bits lo reduce aún más a solo 9.5 GB. Esto permite ejecutar el modelo OctoCode de parámetros 15B en GPU de consumo como la RTX 3090 (24 GB de VRAM).

Sin embargo, tenga en cuenta que una cuantificación más agresiva, como la de 4 bits, a veces puede provocar una degradación de la precisión en comparación con la de 8 bits o bfloat16. Existe un equilibrio entre el ahorro de memoria y la precisión que los usuarios deben evaluar según su caso de uso.

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

Estimación de memoria para ajuste fino

Si bien la cuantificación se utiliza principalmente para una inferencia eficiente, técnicas como el paralelismo tensorial y el paralelismo de modelos son cruciales para gestionar los requisitos de memoria durante el entrenamiento o sintonia FINA de grandes modelos lingüísticos.

El consumo máximo de memoria durante el ajuste fino suele ser entre 3 y 4 veces mayor que el de la inferencia debido a los requisitos de memoria adicionales para:

• Gradientes
• Estados del optimizador
• Activaciones del pase hacia adelante almacenadas para propagación hacia atrás

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

Por ejemplo, ajustar el modelo LLaMA del parámetro 7B requeriría aproximadamente 7 * 8 = 56 GB de VRAM por GPU con precisión bfloat16. Esto excede la capacidad de memoria de las GPU actuales, lo que requiere técnicas de ajuste distribuidas.

Técnicas de ajuste distribuido

Se han propuesto varios métodos de ajuste 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 múltiples GPU mientras divide y distribuye los lotes de datos de entrenamiento. Esto reduce el tiempo de entrenamiento de forma lineal con la cantidad de GPU, pero no reduce el requisito máximo de memoria en cada GPU.
2. Cero etapa 3: una forma avanzada de paralelismo de datos que divide los parámetros del modelo, los gradientes y los estados del optimizador en las GPU. 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 de entrenamiento.
3. Paralelismo tensorial: En lugar de replicar el modelo, el paralelismo tensorial divide los parámetros del modelo en filas o columnas y los distribuye entre las GPU. Cada GPU opera con un conjunto particionado de parámetros, gradientes y estados del optimizador, lo que genera ahorros sustanciales de memoria.
4. Paralelismo de tuberías: Esta técnica divide las capas del modelo en diferentes GPU/trabajadores, y cada dispositivo ejecuta un subconjunto de las capas. Las activaciones se transmiten entre trabajadores, lo que reduce el pico de memoria pero aumenta la sobrecarga de comunicación.

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

La solución LLMem

Los investigadores propusieron recientemente LLMem, una solución que estima con precisión el consumo de memoria de la GPU al aplicar métodos de ajuste distribuidos a LLM en varias GPU.

Estimación del uso de memoria de GPU para ajustar el LLM previamente capacitado

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

Los resultados experimentales muestran que LLMem puede estimar el uso máximo de memoria de la GPU para ajustar los LLM en una sola GPU con tasas de error de hasta el 1.6 %, superando la tasa de error promedio de DNNMem de última generación de un 42.6%. Al aplicar métodos de ajuste distribuido a LLM con más de mil millones de parámetros en múltiples GPU, LLMem logra una impresionante tasa de error promedio de un 3.0%.