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

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

В этом руководстве показано, как использовать Big Model Inference в Accelerate и экосистеме Hugging Face.

Accelerate

Типичный рабочий процесс загрузки модели PyTorch показан ниже. ModelClass — это модель, которая превышает объём памяти GPU вашего устройства (mps или cuda или xpu).

import torch

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

В Big Model Inference первым шагом является инициализация пустого каркаса модели с помощью контекстного менеджера init_empty_weights. Это не требует памяти, так как my_model является «беспараметрическим».

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

Далее веса загружаются в модель для инференса.

Метод load_checkpoint_and_dispatch() загружает контрольную точку в вашу пустую модель и распределяет веса каждого слоя по всем доступным устройствам, начиная с самых быстрых (GPU, MPS, XPU, NPU, MLU, SDAA, MUSA), а затем переходя к более медленным (CPU и жесткий диск).

Установка device_map="auto" автоматически заполняет сначала всё доступное пространство на GPU, затем на CPU и, наконец, на жёстком диске (самый медленный вариант), если памяти всё ещё недостаточно.

См. руководство Проектирование карты устройств, чтобы узнать больше о том, как спроектировать собственную карту устройств.

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

Если есть определённые «блоки» слоёв, которые не следует разделять, передайте их в no_split_module_classes (подробнее — здесь).

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

Теперь, когда модель полностью распределена, можно выполнять инференс.

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Каждый раз, когда входные данные проходят через слой, они передаются с CPU на GPU (или с диска на CPU, а затем на GPU), вычисляется выход, после чего слой удаляется с GPU и возвращается обратно в очередь. Хотя это добавляет некоторый накладной расход на инференс, такой подход позволяет запускать модели любого размера на вашей системе, при условии, что самый большой слой помещается на вашем GPU.

Можно использовать несколько GPU (так называемый «модельный параллелизм»), но в каждый момент времени активен только один GPU. Это заставляет GPU ждать, пока предыдущий GPU передаст ему вывод. Скрипт следует запускать обычным способом с помощью Python, а не с использованием других инструментов, таких как torchrun и accelerate launch.

Ниже приведён полный пример Big Model Inference.

import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

with init_empty_weights():
    model = MyModel(...)

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Экосистема Hugging Face

Другие библиотеки экосистемы Hugging Face, такие как Transformers или Diffusers, поддерживают Big Model Inference в своих конструкторах from_pretrained.

Достаточно добавить device_map="auto" в from_pretrained, чтобы включить Big Model Inference.

Например, загрузите модель Big Sciences T0pp с 11 миллиардами параметров с помощью Big Model Inference.

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")

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

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)

Дальнейшие шаги

Для более подробного объяснения Big Model Inference обязательно ознакомьтесь с концептуальным руководством!