Загрузка моделей

Transformers предоставляет множество предобученных моделей, готовых к использованию в одну строку кода. Для этого требуется класс модели и метод from_pretrained().

Вызовите from_pretrained(), чтобы скачать и загрузить веса и конфигурацию модели, хранящиеся на Hub Hugging Face.

The from_pretrained() method loads weights stored in the safetensors file format if they’re available. Traditionally, PyTorch model weights are serialized with the pickle utility which is known to be insecure. Safetensor files are more secure and faster to load.

from transformers import AutoModelForCausalLM

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

В этом руководстве объясняется, как загружаются модели, какие способы загрузки модели существуют, как решить проблемы с памятью при работе с очень большими моделями и как загружать собственные модели.

Модели и конфигурации

Все модели имеют файл configuration.py с конкретными атрибутами, такими как количество скрытых слоев, размер словаря, функция активации и другие. Вы также найдете файл modeling.py, который определяет слои и математические операции, выполняемые внутри каждого слоя. Файл modeling.py принимает атрибуты модели из configuration.py и строит модель соответственно. На этом этапе у вас есть модель со случайными весами, которую необходимо обучить, чтобы она выдавала осмысленные результаты.

Термин архитектура обозначает каркас модели, а чекпоинт — веса модели для данной архитектуры. Например, BERT — это архитектура, а google-bert/bert-base-uncased — чекпоинт. Термин модель используется как синоним архитектуры и чекпоинта.

Существует два основных типа моделей, которые можно загрузить:

1. Модель без дополнительных слоёв, например AutoModel или LlamaModel, которая выводит скрытые состояния.
2. Модель с подключённой конкретной head, например AutoModelForCausalLM или LlamaForCausalLM, для выполнения определённых задач.

Классы моделей

Чтобы получить предобученную модель, необходимо загрузить веса в модель. Это делается с помощью вызова from_pretrained(), который принимает веса из Hugging Face Hub или из локального каталога.

Существуют два класса моделей: класс AutoModel и класс, специфичный для конкретной модели.

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")

Большие модели

Загрузка крупных предобученных моделей требует значительного объема памяти. Процесс загрузки включает:

1. создание модели со случайными весами
2. загрузка предобученных весов
3. размещение предобученных весов на модели

Для хранения двух копий весов модели (случайной и предобученной) требуется достаточный объём памяти, что может быть невозможно в зависимости от вашего оборудования. В средах распределённого обучения задача усложняется ещё больше, поскольку каждый процесс загружает предобученную модель.

Transformers снижает часть этих проблем, связанных с памятью, за счёт быстрой инициализации, шардированных контрольных точек, функции Big Model Inference в Accelerate и поддержки типов данных с меньшим количеством бит.

Шардированные контрольные точки

save_pretrained() автоматически разбивает на части контрольные точки размером более 50 ГБ. Это позволяет удерживать количество частей на низком уровне для больших моделей и упрощает управление файлами.

Параметры загружаются параллельно, а пиковое потребление памяти зависит только от размера модели. Используйте max_shard_size в save_pretrained(), чтобы задать максимальный размер контрольной точки перед разбиением на шарды.

Объём используемой памяти для моделей, требующих динамического преобразования весов, зависит от размера модели и размера наибольших параметров в рамках одного преобразования. Как правило, это касается моделей типа mixture-of-experts (MoE), где объём используемой памяти равен размеру модели плюс количеству экспертов на одном слое. Подробнее о том, как загружаются модели, см. в руководстве по динамическому загрузчику весов.

save_pretrained() также создаёт файл индекса, сопоставляющий имена параметров с их файлами шардов. Индекс содержит два ключа: metadata и 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 хранит общий размер модели.

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

weight_map сопоставляет каждый параметр с его файлом шарда.

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',
 ...
}

Вывод больших моделей

Убедитесь, что у вас установлены Accelerate v0.9.0 и PyTorch v1.9.0 или более поздние версии, чтобы использовать эту функцию!

from_pretrained() is supercharged with Accelerate’s Big Model Inference feature.

Big Model Inference создаёт каркас модели на устройстве meta PyTorch. Устройство meta не хранит реальные данные, только метаданные.

Случайно инициализированные веса создаются только при загрузке предобученных весов, чтобы избежать одновременного хранения двух копий модели в памяти. Максимальное потребление памяти равно размеру модели.

Подробнее о размещении на устройствах — в разделе Проектирование карты устройств.

Вторая особенность Big Model Inference связана с тем, как веса загружаются и распределяются в каркасе модели. Веса модели распределяются по всем доступным устройствам, начиная с самого быстрого (обычно это GPU), а оставшиеся веса выгружаются на более медленные устройства (CPU и жесткий диск).

Совместное использование обеих функций снижает потребление памяти и время загрузки для больших предобученных моделей.

Установите для device_map значение "auto", чтобы включить Big Model Inference.

from transformers import AutoModelForCausalLM

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

Вы также можете вручную назначить слои устройствам в device_map. Он должен отображать все параметры модели на устройства, но вам не нужно детально указывать, на каком устройстве находятся все подмодули слоя, если весь слой размещен на одном и том же устройстве.

Обратитесь к атрибуту hf_device_map, чтобы увидеть, как модель распределена по устройствам.

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

Разгрузка на диск

Если модель слишком велика для суммарного объёма памяти ваших GPU и CPU, Big Model Inference выгружает оставшиеся веса на диск. Укажите offload_folder как каталог, если хотя бы одна карта весов указывает на "disk", поскольку именно там Transformers хранит выгруженные веса.

from transformers import AutoModelForCausalLM

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

Используйте max_memory, чтобы ограничить объём данных, хранимых на каждом устройстве. Transformers заполняет устройства в порядке их скорости и переносит на диск всё, что не помещается в память. В примере ниже GPU ограничена 10 ГБ, а CPU — 30 ГБ, поэтому оставшиеся веса размещаются в offload_folder.

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

Установите max_memory={}, чтобы указать, что память GPU или CPU недоступна, и принудительно выгрузить всю модель на диск.

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

Разгрузка на диск жертвует скоростью ради экономии памяти. Чтение весов с диска при каждом прямом проходе медленнее, чем чтение из памяти CPU или GPU. Полная разгрузка на диск требует accelerate>1.14.0

Тип данных модели

Аргумент dtype определяет dtype PyTorch, используемый для инициализации весов модели. По умолчанию Transformers загружает веса с dtype или устаревшим значением torch_dtype из config.json. Если config.json не содержит ни одного из этих значений, Transformers использует dtype первого веса с плавающей точкой в контрольной точке.

Переопределите значение по умолчанию, передав конкретный dtype.

import torch
from transformers import AutoModelForCausalLM

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

AutoConfig также принимает dtype для моделей, создаваемых с нуля.

import torch
from transformers import AutoConfig, AutoModel

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

Пользовательские модели

Пользовательские модели строятся на классах конфигурации и моделирования Transformers, поддерживают API AutoClass и загружаются с помощью from_pretrained(). Отличие состоит в том, что код моделирования не входит в состав Transformers.

Будьте особенно осторожны при загрузке пользовательской модели. Хотя на Hub для каждого репозитория выполняется сканирование на вредоносное ПО, всё равно следует проявлять осторожность, чтобы случайно не выполнить вредоносный код.

Установите trust_remote_code=True в from_pretrained(), чтобы загрузить пользовательскую модель.

from transformers import AutoModelForImageClassification

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

В качестве дополнительной меры безопасности загружайте пользовательскую модель из конкретного ревизия, чтобы избежать загрузки кода модели, который мог измениться. Хеш коммита можно скопировать из истории коммитов модели.

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

См. руководство Customize models для получения дополнительной информации.