Carga de modelos

Transformers ofrece muchos modelos preentrenados que están listos para usarse con una sola línea de código. Requiere una clase de modelo y el método from_pretrained().

Llama a from_pretrained() para descargar y cargar los pesos y la configuración de un modelo almacenados en el Hub de Hugging Face.

El método from_pretrained() carga los pesos almacenados en el formato de archivo safetensors si están disponibles. Tradicionalmente, los pesos de los modelos de PyTorch se serializan con la utilidad pickle, que se conoce por ser insegura. Los archivos Safetensor son más seguros y se cargan más rápido.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", device_map="auto")

Esta guía explica cómo se cargan los modelos, las diferentes formas de cargar un modelo, cómo superar problemas de memoria para modelos muy grandes y cómo cargar modelos personalizados.

Modelos y configuraciones

Todos los modelos tienen un archivo configuration.py con atributos específicos, como el número de capas ocultas, el tamaño del vocabulario, la función de activación, entre otros. También encontrarás un archivo modeling.py que define las capas y las operaciones matemáticas que se realizan dentro de cada capa. El archivo modeling.py toma los atributos del modelo en configuration.py y construye el modelo en consecuencia. En este punto, tienes un modelo con pesos aleatorios que necesita ser entrenado para producir resultados significativos.

Una arquitectura se refiere al esqueleto del modelo y un punto de control se refiere a los pesos del modelo para una arquitectura dada. Por ejemplo, BERT es una arquitectura, mientras que google-bert/bert-base-uncased es un punto de control. Verás que el término modelo se utiliza indistintamente para referirse a la arquitectura y al punto de control.

Existen dos tipos generales de modelos que se pueden cargar:

1. Un modelo básico, como AutoModel o LlamaModel, que genera estados ocultos.
2. Un modelo con un head específico adjunto, como AutoModelForCausalLM o LlamaForCausalLM, para realizar tareas específicas.

Clases de modelos

Para obtener un modelo preentrenado, es necesario cargar los pesos en el modelo. Esto se hace llamando a from_pretrained(), que acepta pesos desde el Hugging Face Hub o desde un directorio local.

Hay dos clases de modelo: la clase AutoModel y una clase específica del modelo.

from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoModelForQuestionAnswering

# use the same API for 3 different tasks
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForSequenceClassification.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForQuestionAnswering.from_pretrained("meta-llama/Llama-2-7b-hf")

from transformers import AutoModelForCausalLM

# use the same API to load 3 different models
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b")

from transformers import LlamaModel, LlamaForCausalLM

model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

Modelos grandes

Los modelos preentrenados de gran tamaño requieren mucha memoria para cargarse. El proceso de carga implica:

1. crear un modelo con pesos aleatorios
2. cargar los pesos preentrenados
3. colocar los pesos preentrenados en el modelo

Se necesita suficiente memoria para mantener dos copias de los pesos del modelo (aleatorios y preentrenados), lo cual puede no ser posible dependiendo del hardware. En entornos de entrenamiento distribuido, esto es aún más desafiante, ya que cada proceso carga un modelo preentrenado.

Transformers reduce algunos de estos desafíos relacionados con la memoria mediante la inicialización rápida, los puntos de control fragmentados, la función Big Model Inference de Accelerate y el soporte para tipos de datos de menor número de bits.

Puntos de control fragmentados

save_pretrained() divide automáticamente en fragmentos los puntos de control que superan los 50 GB. Esto mantiene bajo el número de fragmentos en modelos grandes y simplifica la gestión de archivos.

Los parámetros se cargan en paralelo y la memoria máxima depende únicamente del tamaño del modelo. Utilice max_shard_size en save_pretrained() para establecer el tamaño máximo del punto de control antes de la partición.

El uso de memoria para los modelos que requieren una conversión dinámica de pesos depende del tamaño del modelo y del tamaño de los parámetros más grandes en una sola conversión. Esto suele aplicarse a los modelos de mezcla de expertos (MoE), donde el uso de memoria es el tamaño del modelo más el número de expertos en una capa. Consulte la guía de dynamic weight loader para obtener más información sobre cómo se cargan los modelos.

save_pretrained() también crea un archivo de índice que mapea los nombres de los parámetros a sus archivos de fragmentos. El índice contiene dos claves, metadata y weight_map.

import json
import os
import tempfile

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="50GB")
    with open(os.path.join(tmp_dir, "model.safetensors.index.json"), "r") as f:
        index = json.load(f)

print(index.keys())

metadata almacena el tamaño total del modelo.

index["metadata"]
{'total_size': 28966928384}

weight_map asigna cada parámetro a su archivo de fragmento.

index["weight_map"]
{'lm_head.weight': 'model-00006-of-00006.safetensors',
 'model.embed_tokens.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.input_layernorm.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.mlp.down_proj.weight': 'model-00001-of-00006.safetensors',
 ...
}

Inferencia de modelos grandes

Asegúrate de tener instalados Accelerate v0.9.0 y PyTorch v1.9.0 o una versión posterior para usar esta función.

from_pretrained() se potencia con la función Big Model Inference de Accelerate.

Big Model Inference crea un esqueleto del modelo en el dispositivo meta de PyTorch. El dispositivo meta no almacena datos reales, solo metadatos.

Los pesos inicializados aleatoriamente solo se crean cuando se cargan los pesos preentrenados, para evitar mantener dos copias del modelo en memoria al mismo tiempo. El uso máximo de memoria es solo el tamaño del modelo.

Obtenga más información sobre la colocación de dispositivos en Diseño de un mapa de dispositivos.

La segunda característica de Big Model Inference se relaciona con la forma en que se cargan y distribuyen los pesos en la estructura del modelo. Los pesos del modelo se distribuyen entre todos los dispositivos disponibles, comenzando por el dispositivo más rápido (normalmente la GPU) y luego descargando los pesos restantes a dispositivos más lentos (CPU y disco duro).

La combinación de ambas características reduce el uso de memoria y los tiempos de carga para los modelos preentrenados de gran tamaño.

Establezca device_map en "auto" para habilitar Big Model Inference.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

También puedes asignar manualmente las capas a un dispositivo en device_map. Debe mapear todos los parámetros del modelo a un dispositivo, pero no es necesario detallar en qué dispositivo se encuentran todos los submódulos de una capa si la capa completa está en el mismo dispositivo.

Acceda al atributo hf_device_map para ver cómo se distribuye un modelo entre dispositivos.

device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}
model.hf_device_map

Descarga en disco

Cuando un modelo es demasiado grande para la memoria combinada de tus GPUs y la RAM de la CPU, Big Model Inference descarga los pesos restantes al disco. Establece offload_folder en un directorio siempre que algún peso se asigne a "disk", ya que es allí donde Transformers almacena los pesos descargados.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    offload_folder="./offload",
)

Use max_memory para limitar la cantidad de memoria que ocupa cada dispositivo. Transformers rellena los dispositivos en orden de velocidad y envía a disco lo que no quepa. El ejemplo siguiente limita la GPU a 10 GB y la CPU a 30 GB, de modo que los pesos restantes se almacenan en offload_folder.

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={0: "10GB", "cpu": "30GB"},
    offload_folder="./offload",
)

Establezca max_memory={} para especificar que no hay memoria de gpu ni de cpu disponible y para forzar el descargado del modelo completo al disco.

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-7b",
    device_map="auto",
    max_memory={},
    offload_folder="./offload",
)

El descargado en disco intercambia memoria por velocidad. Leer los pesos desde el disco en cada pasada hacia adelante es más lento que leerlos desde la memoria de la CPU o la GPU. El descargado completo en disco requiere accelerate>1.14.0

Tipo de datos del modelo

El argumento dtype controla el dtype de PyTorch utilizado para instanciar los pesos del modelo. De forma predeterminada, Transformers carga los pesos con el valor de dtype o el valor heredado de torch_dtype de config.json. Si config.json no incluye ninguno de estos valores, Transformers utiliza el dtype del primer peso de punto flotante en el checkpoint.

Sobrescribe el valor predeterminado pasando un dtype específico.

import torch
from transformers import AutoModelForCausalLM

# specific dtype
model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it", dtype=torch.float16)

AutoConfig también acepta dtype para modelos instanciados desde cero.

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", dtype=torch.float16)
model = AutoModel.from_config(my_config)

Modelos personalizados

Los modelos personalizados se basan en las clases de configuración y modelado de Transformers, admiten la API AutoClass y se cargan con from_pretrained(). La diferencia es que el código de modelado no proviene de Transformers.

Tome precauciones adicionales al cargar un modelo personalizado. Aunque el Hub incluye detección de malware en cada repositorio, debe tener cuidado para evitar ejecutar accidentalmente código malicioso.

Establezca trust_remote_code=True en from_pretrained() para cargar un modelo personalizado.

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

Como capa de seguridad adicional, cargue un modelo personalizado desde una revisión específica para evitar cargar código de modelo que podría haber cambiado. El hash del commit se puede copiar del historial de commits de los modelos.

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

Consulte la guía Personalizar modelos para obtener más información.