Загрузка больших моделей в память

При загрузке предобученной модели в PyTorch типичный рабочий процесс выглядит следующим образом:

import torch

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

Проще говоря, эти шаги выглядят так:

1. Создайте модель со случайно инициализированными весами
2. Загрузите веса модели (обычно в виде словаря, называемого state dict) с диска
3. Загрузите эти веса в модель

Хотя этот подход отлично работает для моделей стандартного размера, у него есть очевидные ограничения при работе с очень большими моделями: на шаге 1 мы загружаем полную версию модели в оперативную память и тратим время на случайную инициализацию весов (которые будут отброшены на шаге 3). На шаге 2 мы загружаем в оперативную память еще одну полную версию модели, но уже с предобученными весами. Если вы загружаете модель с 6 миллиардами параметров, это означает, что для каждой копии модели потребуется 24 ГБ оперативной памяти, то есть 48 ГБ в общей сложности (половина из них — для загрузки модели в формате FP16).

Этот API довольно новый и всё ещё находится на экспериментальной стадии. Мы стремимся обеспечить стабильность API, но в будущем возможны небольшие изменения в отдельных частях публичного API.

Как работает процесс: краткий обзор

Как работает процесс: работа с кодом

Создание пустой модели

The first tool Accelerate introduces to help with big models is a context manager init_empty_weights() that helps you initialize a model without using any RAM so that step 1 can be done on models of any size. Here is how it works:

from accelerate import init_empty_weights

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

Например:

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

инициализирует пустую модель с чуть более чем 100 млрд параметров. За кулисами это опирается на мета-устройство, представленное в PyTorch 1.9. Во время инициализации в контекстном менеджере каждый раз при создании параметра он немедленно перемещается на это устройство.

Нельзя напрямую переместить модель, инициализированную таким образом, на CPU или другое устройство, так как она не содержит данных. Кроме того, с высокой вероятностью прямой проход с этой пустой моделью завершится ошибкой, поскольку не все операции поддерживаются на meta-устройстве.

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

Возможно, ваша модель настолько велика, что даже одна её копия не помещается в оперативную память. Это не означает, что её нельзя загрузить: если у вас есть один или несколько GPU, это дополнительный объём памяти для хранения модели. В этом случае лучше, если ваш checkpoint разделён на несколько меньших файлов, которые мы называем фрагментами checkpoint.

Accelerate будет работать с разбитыми на части контрольными точками, если вы придерживаетесь следующего формата: контрольная точка должна находиться в папке, содержащей несколько файлов с частичными state dict, а также должен присутствовать индекс в формате JSON, который содержит словарь, сопоставляющий имена параметров с файлами, содержащими их веса. Вы можете легко разбить модель на части с помощью save_model(). Например, папка может содержать:

first_state_dict.bin
index.json
second_state_dict.bin

где index.json — это следующий файл:

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

и first_state_dict.bin, содержащий веса для "linear1.weight" и "linear1.bias", а second_state_dict.bin — для "linear2.weight" и "linear2.bias"

Загрузка весов

Второй инструмент, который представляет Accelerate, — это функция load_checkpoint_and_dispatch(), позволяющая загрузить контрольную точку в вашу пустую модель. Поддерживаются как полные контрольные точки (один файл, содержащий весь state dict), так и разбитые на части (sharded). Функция также автоматически распределяет эти веса по доступным вам устройствам (GPU, оперативная память CPU), поэтому при загрузке разбитой контрольной точки максимальное потребление оперативной памяти будет равно размеру крупнейшей части.

Если вы хотите использовать инференс больших моделей с моделями Transformers, ознакомьтесь с этой документацией.

Вот как можно использовать это для загрузки модели GPT2-1.5B.

Скачаем шардированную версию этой модели.

pip install huggingface_hub

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

Для инициализации модели мы будем использовать библиотеку 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)

Затем загрузите только что скачанный контрольный пункт с помощью:

from accelerate import load_checkpoint_and_dispatch

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

Передавая device_map="auto", мы указываем Accelerate автоматически определять, на каком устройстве разместить каждый слой модели в зависимости от доступных ресурсов:

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

no_split_module_classes

Этот параметр указывает, что некоторые модули с именем "Block" не следует распределять по разным устройствам. Здесь следует указать все блоки, содержащие остаточное соединение любого типа.

Параметр device_map

Вы можете увидеть device_map, который выбрал Accelerate, обратившись к атрибуту hf_device_map вашей модели:

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}

Полностью возможно создать собственную карту устройств для слоёв, указав используемое устройство GPU (номер), "cpu" или "disk", и передать её:

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
)

Запуск модели

Теперь, когда мы это сделали, наша модель распределена по нескольким устройствам, а возможно, и на жёстком диске. Но её по-прежнему можно использовать как обычную модель PyTorch:

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

Под капотом Accelerate добавляет хуки к модели, чтобы:

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

Таким образом, ваша модель может выполнять инференс, даже если она не помещается в память одного из GPU или в оперативную память CPU!

Эта функция поддерживает только инференс вашей модели, но не обучение. Большая часть вычислений выполняется внутри контекстных менеджеров torch.no_grad(), чтобы избежать расхода GPU-памяти на промежуточные активации.

Проектирование карты устройств

Вы можете поручить Accelerate вычисление карты устройств, установив device_map в одно из поддерживаемых значений ("auto", "balanced", "balanced_low_0", "sequential"), либо создать её самостоятельно, если хотите точнее контролировать, куда должна попасть каждый слой.

Вы можете вывести все размеры модели (а значит, вычислить device_map) для модели, находящейся на meta-устройстве.

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

Когда доступной памяти GPU больше, чем размер модели, вот в чём разница между каждым вариантом:

• "auto" и "balanced" равномерно распределяют модель по всем доступным GPU, что позволяет использовать размер батча, больший 1.
• "balanced_low_0" равномерно распределяет модель по всем GPU, кроме первого, и размещает на GPU 0 только то, что не помещается на остальных. Этот вариант отлично подходит, когда нужно использовать GPU 0 для обработки результатов, например, при использовании функции generate для моделей Transformers
• "sequential" размещает на GPU 0 всё, что помещается, затем переходит на GPU 1 и так далее (поэтому последние GPU не используются, если в них нет необходимости).

Варианты "auto" и "balanced" на данный момент дают одинаковые результаты, однако поведение "auto" может измениться в будущем, если мы найдём более разумную стратегию, тогда как "balanced" останется стабильным.

Сначала обратите внимание, что вы можете ограничить объем памяти, используемой на каждой GPU, с помощью аргумента max_memory (доступного в infer_auto_device_map() и во всех функциях, использующих его). При установке max_memory необходимо передать словарь, содержащий идентификаторы GPU (например, 0, 1 и т. д.) и ключ "cpu" для указания максимального объема ОЗУ, который вы хотите использовать для CPU offload. Значения могут быть целым числом (в байтах) или строкой, представляющей число с единицей измерения, например "10GiB" или "10GB".

Вот пример, в котором мы не хотим использовать более 10GiB на каждом из двух GPU и не более 30GiB оперативной памяти CPU для весов модели:

from accelerate import infer_auto_device_map

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

При первом выделении памяти в PyTorch загружаются CUDA-ядра, которые занимают около 1–2 ГБ памяти в зависимости от GPU. Поэтому доступная память всегда меньше фактического размера GPU. Чтобы увидеть, сколько памяти используется на самом деле, выполните torch.ones(1).cuda() и посмотрите на использование памяти.

Поэтому при создании карт памяти с max_memory убедитесь, что доступная память скорректирована соответствующим образом, чтобы избежать ошибок нехватки памяти.

Кроме того, если вы выполняете дополнительные операции с результатами, не возвращая их на CPU (например, внутри метода generate в Transformers), и при этом ваши входные данные размещены на GPU, то этот GPU будет потреблять больше памяти, чем остальные (Accelerate всегда возвращает результат на устройство, на котором находились входные данные). Поэтому, если вы хотите оптимизировать максимальный размер батча и у вас много GPU, выделите первому GPU меньше памяти. Например, для BLOOM-176B на конфигурации 8x80 A100 близкое к идеальному распределение выглядит так:

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

как видно, мы выделили оставшимся 7 GPU примерно на 50% больше памяти, чем GPU 0.

Если вы решите полностью спроектировать device_map самостоятельно, он должен представлять собой словарь, где ключами являются имена модулей вашей модели, а значениями — допустимый идентификатор устройства (например, целое число для GPU) или "cpu" для выгрузки на CPU, "disk" для выгрузки на диск. Ключи должны покрывать всю модель, после чего вы можете определить карту устройств по своему усмотрению: например, если ваша модель содержит два блока (назовем их block1 и block2), каждый из которых включает три линейных слоя (назовем их linear1, linear2 и linear3), допустимая карта устройств может выглядеть так:

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

ещё один допустимый вариант:

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

С другой стороны, этот вариант не является корректным, так как он не охватывает все параметры модели:

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

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

Только выгрузка на CPU

Если вы хотите выгрузить модель на CPU, используйте cpu_offload(). В результате все параметры модели будут выгружены, и будет сохранена только одна копия state dict модели. Во время прямого прохода параметры извлекаются из этого state dict, переносятся на устройство выполнения и передаются по мере необходимости, после чего снова выгружаются.

cpu_offload(model, execution_device)

You can also use cpu_offload_with_hook(). This function will offloads a model on the CPU and puts it back to an execution device when executed. The difference with cpu_offload() is that the model stays on the execution device after the forward and is only offloaded again when the offload method of the returned hook is called. Furthermore, cpu_offload_with_hook() is more performant but less memory saving. It is useful for pipelines running a model in a loop:

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

Только выгрузка на диск

Для выполнения disk offload можно использовать disk_offload(). В результате все параметры модели будут выгружены в виде memory-mapped массива в указанную папку. Во время forward pass параметры будут считываться из этой папки и по мере необходимости перемещаться на устройство выполнения, переданное в качестве аргумента, после чего снова выгружаться.

disk_offload(model, offload_dir, execution_device)

Ограничения и дальнейшее развитие

Нам о текущих ограничениях API:

• infer_auto_device_map() (или device_map="auto" в load_checkpoint_and_dispatch()) пытается максимально использовать доступную память GPU и CPU в момент выполнения. PyTorch очень эффективно управляет памятью GPU (и освобождает её, когда она не нужна), но то же самое нельзя сказать о Python и памятью CPU. Поэтому автоматически вычисленная карта устройств может чрезмерно нагружать CPU. Если возникают сбои из-за нехватки RAM, переместите несколько модулей на устройство disk.
• infer_auto_device_map() (или device_map="auto" в load_checkpoint_and_dispatch()) назначает устройства последовательно (чтобы избежать перемещения данных туда-сюда), поэтому, если ваш первый слой больше размера вашей GPU, всё окажется на CPU/Disk.
• load_checkpoint_and_dispatch() и load_checkpoint_in_model() в настоящее время не проверяют корректность вашего state dict по сравнению с моделью (это будет исправлено в будущей версии), поэтому при попытке загрузить checkpoint с несовпадающими или отсутствующими ключами могут возникать странные ошибки.
• Модельный параллелизм, используемый при разбиении модели на несколько GPU, является наивным и не оптимизированным, что означает, что в каждый момент времени работает только один GPU, а остальные простаивают.
• Когда веса выгружаются на CPU/жёсткий диск, предварительная загрузка (pre-fetching) не выполняется (пока что, мы будем работать над этим в будущих версиях), что означает, что веса переносятся на GPU только в момент их необходимости, а не заранее.
• Разгрузка на жёсткий диск может быть очень медленной, если ваше оборудование не обеспечивает быструю связь между диском и процессором (например, через NVMe).