Carga de modelos grandes en memoria

Al cargar un modelo preentrenado en PyTorch, el flujo de trabajo habitual es el siguiente:

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

En términos simples, esos pasos son:

1. Cree el modelo con pesos inicializados aleatoriamente
2. Cargue los pesos del modelo (en un diccionario, generalmente llamado state dict) desde el disco
3. Cargue esos pesos dentro del modelo

Aunque esto funciona muy bien para modelos de tamaño regular, este flujo de trabajo tiene algunas limitaciones claras cuando se trata de un modelo enorme: en el paso 1, cargamos una versión completa del modelo en la RAM y dedicamos tiempo a inicializar aleatoriamente los pesos (que se descartarán en el paso 3). En el paso 2, cargamos otra versión completa del modelo en la RAM, con los pesos preentrenados. Si estás cargando un modelo con 6 mil millones de parámetros, esto significa que necesitarás 24 GB de RAM para cada copia del modelo, por lo que serían 48 GB en total (la mitad para cargar el modelo en FP16).

Esta API es bastante nueva y aún se encuentra en su etapa experimental. Aunque nos esforzamos por proporcionar una API estable, es posible que algunas partes pequeñas de la API pública cambien en el futuro.

Cómo funciona el proceso: una visión general

Cómo funciona el proceso: trabajar con código

Instanciar un modelo vacío

La primera herramienta que Accelerate introduce para ayudar con modelos grandes es un gestor de contexto init_empty_weights() que te permite inicializar un modelo sin usar RAM, de modo que el paso 1 pueda realizarse con modelos de cualquier tamaño. Así funciona:

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

Por ejemplo:

with init_empty_weights():
    model = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

inicializa un modelo vacío con un poco más de 100B de parámetros. Detrás de escena, esto se basa en el dispositivo meta introducido en PyTorch 1.9. Durante la inicialización bajo el gestor de contexto, cada vez que se crea un parámetro, se mueve instantáneamente a ese dispositivo.

No se puede mover directamente a CPU u otro dispositivo un modelo inicializado de esta manera, ya que no contiene ningún dato. Además, es muy probable que una pasada hacia adelante con ese modelo vacío falle, ya que no todas las operaciones están soportadas en el dispositivo meta.

Puntos de control fragmentados

Es posible que tu modelo sea tan grande que ni siquiera una copia quepa en la RAM. Eso no significa que no se pueda cargar: si tienes una o varias GPUs, tienes más memoria disponible para almacenar tu modelo. En este caso, es mejor que tu checkpoint esté dividido en varios archivos más pequeños que llamamos fragmentos de checkpoint.

Accelerate gestionará los puntos de control fragmentados siempre que sigas el siguiente formato: tu punto de control debe estar en una carpeta, con varios archivos que contienen los diccionarios de estado parciales, y debe haber un índice en formato JSON que contenga un diccionario que mapee los nombres de los parámetros al archivo que contiene sus pesos. Puedes fragmentar fácilmente tu modelo con save_model(). Por ejemplo, podríamos tener una carpeta que contenga:

first_state_dict.bin
index.json
second_state_dict.bin

donde index.json es el siguiente archivo:

{
  "linear1.weight": "first_state_dict.bin",
  "linear1.bias": "first_state_dict.bin",
  "linear2.weight": "second_state_dict.bin",
  "linear2.bias": "second_state_dict.bin"
}

y first_state_dict.bin que contiene los pesos para "linear1.weight" y "linear1.bias", second_state_dict.bin los de "linear2.weight" y "linear2.bias"

Carga de pesos

La segunda herramienta que introduce Accelerate es la función load_checkpoint_and_dispatch(), que te permitirá cargar un checkpoint dentro de tu modelo vacío. Esto admite checkpoints completos (un único archivo que contiene todo el state dict) así como checkpoints fragmentados. También distribuirá automáticamente esos pesos entre los dispositivos que tengas disponibles (GPUs, RAM de CPU), por lo que, si cargas un checkpoint fragmentado, el uso máximo de RAM será el tamaño del fragmento más grande.

Si deseas usar la inferencia de modelos grandes con modelos de Transformers, consulta esta documentación.

A continuación se muestra cómo podemos usar esto para cargar el modelo GPT2-1.5B.

Descarguemos la versión fragmentada de este modelo.

pip install huggingface_hub

from huggingface_hub import snapshot_download
checkpoint = "marcsun13/gpt2-xl-linear-sharded"
weights_location = snapshot_download(repo_id=checkpoint)

Para inicializar el modelo, utilizaremos la biblioteca minGPT.

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    model = GPT(model_config)

Luego, carga el punto de control que acabas de descargar con:

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map="auto", no_split_module_classes=['Block']
)

Al pasar device_map="auto", le indicamos a Accelerate que determine automáticamente dónde colocar cada capa del modelo en función de los recursos disponibles:

• primero, usamos el espacio máximo disponible en la(s) GPU
• si aún necesitamos espacio, almacenamos los pesos restantes en la CPU
• si no hay suficiente RAM, almacenamos los pesos restantes en el disco duro como tensores mapeados en memoria

no_split_module_classes

Este parámetro indicará que algunos de los módulos con el nombre "Block" no deben dividirse entre diferentes dispositivos. Debe establecer aquí todos los bloques que incluyan una conexión residual de algún tipo.

El device_map

Puedes ver el device_map que seleccionó Accelerate accediendo al atributo hf_device_map de tu modelo:

model.hf_device_map

{'transformer.wte': 0,
 'transformer.wpe': 0,
 'transformer.drop': 0,
 'transformer.h.0': 0,
 ...
 'transformer.h.21': 0, 
 'transformer.h.22': 1, 
 'transformer.h.23': 1, 
 'transformer.h.24': 1,
 ...
 'transformer.h.47': 1, 
 'transformer.ln_f': 1, 
 'lm_head': 1}

Es totalmente posible crear su propio mapa de dispositivos para las capas y especificar el dispositivo de GPU a utilizar (un número), "cpu" o "disk", y pasar este mapa de la siguiente manera:

device_map = {
    "transformer.wte": "cpu",
    "transformer.wpe": 0,
    "transformer.drop": "cpu",
    "transformer.h.0": "disk"
}

model = load_checkpoint_and_dispatch(
    model, checkpoint=weights_location, device_map=device_map
)

Ejecutar el modelo

Ahora que hemos hecho esto, nuestro modelo se distribuye en varios dispositivos y, posiblemente, en el disco duro. Pero aún se puede usar como un modelo de PyTorch regular:

from mingpt.bpe import BPETokenizer
tokenizer = BPETokenizer()
inputs = tokenizer("Hello, my name is").to(0)

outputs = model.generate(x1, max_new_tokens=10, do_sample=False)[0]
tokenizer.decode(outputs.cpu().squeeze())

Detrás de escena, Accelerate añadió hooks al modelo, de modo que:

• en cada capa, las entradas se colocan en el dispositivo correcto (de modo que funcione incluso si tu modelo está distribuido en varias GPU)
• para los pesos descargados en la CPU, se colocan en una GPU justo antes de la pasada hacia adelante y se limpian justo después
• para los pesos descargados en el disco duro, se cargan en la RAM y luego se colocan en una GPU justo antes de la pasada hacia adelante y se limpian justo después

De esta manera, tu modelo puede ejecutarse para inferencia incluso si no cabe en una de las GPUs o en la RAM de la CPU.

Esto solo admite la inferencia de tu modelo, no el entrenamiento. La mayor parte del cómputo se realiza dentro de los administradores de contexto torch.no_grad() para evitar gastar memoria de GPU con activaciones intermedias.

Diseño de un mapa de dispositivos

Puedes dejar que Accelerate calcule el mapa de dispositivos estableciendo device_map en una de las opciones admitidas ("auto", "balanced", "balanced_low_0", "sequential") o crear uno tú mismo si deseas tener más control sobre dónde debe ir cada capa.

Puede derivar todos los tamaños del modelo (y por lo tanto calcular un device_map) en un modelo que está en el dispositivo meta.

Todas las opciones producirán el mismo resultado cuando no se dispone de suficiente memoria de GPU para alojar el modelo completo (es decir, para colocar en la GPU todo lo que quepa y, a continuación, descargar los pesos a la CPU o incluso al disco si no hay suficiente RAM).

Cuando se dispone de más memoria de GPU que el tamaño del modelo, esta es la diferencia entre cada opción:

• "auto" y "balanced" dividen el modelo de manera uniforme en todas las GPU disponibles, lo que te permite usar un tamaño de lote mayor que 1.
• "balanced_low_0" divide el modelo de manera uniforme en todas las GPUs excepto la primera, y solo coloca en la GPU 0 lo que no cabe en las demás. Esta opción es ideal cuando necesitas usar la GPU 0 para algún procesamiento de las salidas, como al usar la función generate para modelos de Transformers
• "sequential" ocupará lo que pueda en la GPU 0, luego pasará a la GPU 1 y así sucesivamente (por lo que no usará las últimas GPUs si no las necesita).

Las opciones "auto" y "balanced" producen los mismos resultados por ahora, pero el comportamiento de "auto" podría cambiar en el futuro si encontramos una estrategia que tenga más sentido, mientras que "balanced" se mantendrá estable.

En primer lugar, ten en cuenta que puedes limitar la memoria utilizada en cada GPU mediante el argumento max_memory (disponible en infer_auto_device_map() y en todas las funciones que lo utilizan). Al establecer max_memory, debes pasar un diccionario que contenga los identificadores de las GPUs (por ejemplo, 0, 1, etc.) y la clave "cpu" para la RAM máxima que deseas usar para el offload de CPU. Los valores pueden ser un entero (en bytes) o una cadena que represente un número con su unidad, como "10GiB" o "10GB".

A continuación se muestra un ejemplo en el que no queremos usar más de 10GiB en cada una de las dos GPUs y no más de 30GiB de RAM de CPU para los pesos del modelo:

from accelerate import infer_auto_device_map

device_map = infer_auto_device_map(my_model, max_memory={0: "10GiB", 1: "10GiB", "cpu": "30GiB"})

Cuando se produce la primera asignación en PyTorch, se cargan los kernels de CUDA, que ocupan aproximadamente 1-2 GB de memoria dependiendo de la GPU. Por lo tanto, siempre se dispone de menos memoria utilizable que el tamaño real de la GPU. Para ver cuánta memoria se utiliza realmente, ejecute torch.ones(1).cuda() y observe el uso de memoria.

Por lo tanto, al crear mapas de memoria con max_memory, asegúrate de ajustar la memoria disponible en consecuencia para evitar errores de memoria insuficiente.

Además, si realizas operaciones adicionales con tus salidas sin devolverlas a la CPU (por ejemplo, dentro del método generate de Transformers) y si colocaste tus entradas en una GPU, esa GPU consumirá más memoria que las demás (Accelerate siempre devuelve la salida al dispositivo de la entrada). Por lo tanto, si deseas optimizar el tamaño máximo del lote y tienes muchas GPUs, asigna menos memoria a la primera GPU. Por ejemplo, con BLOOM-176B en una configuración de 8x80 A100, el mapa casi ideal es:

max_memory = {0: "30GIB", 1: "46GIB", 2: "46GIB", 3: "46GIB", 4: "46GIB", 5: "46GIB", 6: "46GIB", 7: "46GIB"}

como pueden ver, asignamos a las 7 GPUs restantes aproximadamente un 50% más de memoria que a la GPU 0.

Si opta por diseñar completamente el device_map usted mismo, debe ser un diccionario cuyas claves sean los nombres de los módulos de su modelo y cuyos valores sean un identificador de dispositivo válido (por ejemplo, un entero para las GPU) o "cpu" para el offload a CPU, "disk" para el offload a disco. Las claves deben cubrir todo el modelo, y luego puede definir su mapa de dispositivos como desee: por ejemplo, si su modelo tiene dos bloques (digamos block1 y block2) que contienen cada uno tres capas lineales (digamos linear1, linear2 y linear3), un mapa de dispositivos válido puede ser:

device_map = {"block1": 0, "block2": 1}

otro que sea válido podría ser:

device_map = {"block1": 0, "block2.linear1": 0, "block2.linear2": 1, "block2.linear3": 1}

Por otro lado, esta no es válida porque no cubre todos los parámetros del modelo:

device_map = {"block1": 0, "block2.linear1": 1, "block2.linear2": 1}

Para ser lo más eficiente posible, asegúrate de que tu mapa de dispositivos coloque los parámetros en las GPUs de manera secuencial (por ejemplo, no pongas uno de los primeros pesos en la GPU 0, luego pesos en la GPU 1 y el último peso de vuelta a la GPU 0) para evitar muchas transferencias de datos entre las GPUs.

Solo offload de CPU

Si desea descargar su modelo en la CPU, puede usar cpu_offload(). Como resultado, todos los parámetros del modelo se descargarán y solo se mantendrá una copia del state dict del modelo. Durante la pasada hacia adelante, los parámetros se extraerán de ese state dict y se colocarán en el dispositivo de ejecución, se pasarán a medida que se necesiten y luego se descargarán nuevamente.

cpu_offload(model, execution_device)

También puede usar cpu_offload_with_hook(). Esta función descarga el modelo a la CPU y lo devuelve al dispositivo de ejecución cuando se ejecuta. La diferencia con cpu_offload() es que el modelo permanece en el dispositivo de ejecución después del forward y solo se descarga de nuevo cuando se llama al método offload del hook devuelto. Además, cpu_offload_with_hook() es más eficiente pero ahorra menos memoria. Es útil para pipelines que ejecutan un modelo en un bucle:

model_1, hook_1 = cpu_offload_with_hook(model_1, execution_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, execution_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, execution_device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

Solo offload en disco

Para realizar el descargado en disco, puede usar disk_offload(). Como resultado, todos los parámetros del modelo se descargarán como un array mapeado en memoria en una carpeta dada. Durante la pasada hacia adelante, los parámetros se accederán desde esa carpeta y se colocarán en el dispositivo de ejecución pasado a medida que se necesiten, para luego volver a descargarse.

disk_offload(model, offload_dir, execution_device)

Limitaciones y desarrollo futuro

Somos conscientes de las limitaciones actuales de la API:

• infer_auto_device_map() (o device_map="auto" en load_checkpoint_and_dispatch()) intenta maximizar la RAM de GPU y CPU que detecta como disponible en el momento de su ejecución. Si bien PyTorch es muy eficiente a la hora de gestionar la RAM de GPU (y liberarla cuando no se necesita), esto no ocurre del todo con Python y la RAM de CPU. Por lo tanto, un mapa de dispositivos calculado automáticamente podría resultar demasiado exigente para la CPU. Si experimentas errores por falta de RAM, mueve algunos módulos al dispositivo de disco.
• infer_auto_device_map() (o device_map="auto" en load_checkpoint_and_dispatch()) asigna dispositivos de forma secuencial (para evitar mover los datos de un lado a otro), por lo que si tu primera capa es más grande que el tamaño de la GPU que tienes, todo terminará en la CPU/Disk.
• load_checkpoint_and_dispatch() y load_checkpoint_in_model() no realizan ninguna comprobación de la corrección de su state dict en comparación con su modelo en este momento (esto se corregirá en una versión futura), por lo que es posible que obtenga algunos errores extraños si intenta cargar un checkpoint con claves que no coinciden o que faltan.
• El paralelismo de modelos utilizado cuando su modelo se divide en varias GPUs es ingenuo y no está optimizado, lo que significa que solo una GPU trabaja en un momento dado y las demás permanecen inactivas.
• Cuando los pesos se descargan en la CPU/disco duro, no hay pre-carga (aún, trabajaremos en esto para futuras versiones), lo que significa que los pesos se colocan en la GPU cuando se necesitan y no antes.
• La descarga en el disco duro puede ser muy lenta si el hardware en el que se ejecuta no tiene una comunicación rápida entre el disco y la CPU (como los NVMes).