Короткий зміст Завантаження моделей машинного навчання повільне, навіть з прогрітим кешем сторінок Linux. Тому ми створили бібліотеку, щоб зробити це швидким. Є деякі цікаві технічні деталі, якими ми хочемо поділитися, тому ми написали цей блог. Бібліотека також мала несподіваний вплив, про що йдеться в кінці.
Обґрунтування
Все почалося 2 роки тому, коли ми випустили нашу першу спробу режиму генерації lowpoly. Режим lowpoly не пішов добре, він видає погані результати з сьогоднішньої точки зору, але ми заплатили за це багато — виділений GPU обробляє лише одноцифрові завдання на день. Він має тонко налаштовані ваги, достатньо великі, щоб витіснити всі інші ваги моделі з VRAM. Гірше того, у нас, можливо, є 3 такі моделі (не можу згадати точну кількість), вони складали значну частину нашої інфраструктури інференсу, створюючи досить невблаганне співвідношення ефективності. І ні, ми не можемо наївно завантажувати моделі вчасно, це займає 30 секунд, більше, ніж фактичний час обробки.
У нас тоді не було спеціальних інженерів з конвеєрів, наші розробники алгоритмів робили все можливе, щоб обійти це. Через кілька днів наш код був засмічений this.to('cpu') та that.to('cuda'). Цей підхід працює деякий час, але час від часу порушує потік наших розробників алгоритмів. А що, якби все відбувалося автоматично? Це Python, речі дійсно відбуваються автоматично в Python.
Як ви визначаєте 'автоматично'?
Давайте перейдемо до ролі розробника алгоритмів. Все досить ясно: я не хочу турбуватися про продуктивність під час виконання поза моїм основним алгоритмом, якщо це абсолютно не потрібно. Я б вважав за краще не знати нічого про обмін моделями.
Звичайно, ми не можемо цього досягти, але можемо спробувати мінімізувати втручання, яке ми повинні ввести в код алгоритму. Це нагадує мені про монкі-патчинг бібліотеки gevent, яка патчить (в основному) бібліотеку socket, замінюючи її на gevent.socket, яка може переключатися на інші зелені потоки, коли IO блокується, дуже схоже на горутину (насправді gevent старший за Golang!).
Оскільки ми використовували лише бібліотеки HuggingFace (transformers, diffusers) для завантаження моделей на той час, мета стала ясною: ми вводимо лише виклик монкі-патч, і решта коду повинна залишатися незмінною, XXXPipeline.from_pretrained(...) має бути набагато швидшим.
Деякі факти, очевидні рішення та припущення
Overmind — це бібліотека кешування, вона кешує результати викликів завантаження моделей у пам'ять системи та пізніше швидко їх реконструює.
Ми пропускаємо обговорення того, як реалізовано монкі-патчинг, це не дуже цікава деталь. Все, що нам потрібно знати, це те, що він перенаправляє всі виклики XXXPipeline.from_pretrained(...) до overmind.api.load(XXXPipeline.from_pretrained, ...).
Ми використовуємо pickle для серіалізації результату кешу, оскільки... у нас немає вибору, і torch.save сам використовує pickle, дивно не використовувати його.
Ми використовуємо архітектуру клієнт/сервер, оскільки не хочемо анулювати наш кеш, коли процес завершується. Багато викликів підпроцесів можуть отримати від цього користь.
Ми припускаємо, що параметри XXXPipeline.from_pretrained є простими хешованими речами (str та подібні) та іншими моделями, завантаженими за допомогою overmind (пояснено пізніше).
Назва overmind запозичена з Starcraft, як ви могли здогадатися.
Реконструюйте це швидко!
Ми не можемо наївно зберегти результат pickle.loads у пам'яті та вважати, що все зроблено. Зрештою, у прогрітому сценарії кеш сторінок Linux виконує свою роботу, кешуючи моделі на диску, і ми все ще можемо бачити час завантаження, що вимірюється десятками секунд.
Неефективність виникає через копіювання пам'яті. У Python навіть створення мільйонів об'єктів коштуватиме не більше кількох сотень мс. Однак для копіювання пам'яті об'ємом 10 ГіБ це коштуватиме пів секунди. Ми повинні уникати копіювання пам'яті наскільки це можливо.
На щастя, більшість великих шматків пам'яті — це тензори Torch, ми можемо безпечно адресувати лише їх і ігнорувати решту.
Насправді, я отримав знання про внутрішню структуру тензора Torch у коді редукції, досліджуючи механізм спільного використання тензорів:
# Скопійовано з torch.multiprocessing.reductions, більшість коду видалено
def reduce_tensor(tensor):
...
storage = tensor._typed_storage()
...
metadata = (
tensor.storage_offset(),
tensor.size(),
tensor.stride(),
tensor.requires_grad,
)
return (rebuild_tensor, (type(tensor), storage, metadata))Досить просто: тензор — це його тип, метадані та його базове сховище. Тут storage має тип TypedStorage, але насправді TypedStorage — це просто простий обгорток для UntypedStorage. UntypedStorage — це клас, який фактично зберігає всі дані тензора.
Наше завдання стало більш конкретним: як уникнути копіювання UntypedStorage? Чи можемо ми самостійно керувати цією пам'яттю тензора і створювати UntypedStorage, вказуючи на пам'ять, якою ми керуємо?
Відповідь — так!
Переглядаючи код на C++, де створюється UntypedStorage, ми можемо легко знайти такий фрагмент коду:
// Скопійовано з torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
// ...опущено нерелевантний код...
auto new_storage_impl = make_storage_impl(
c10::StorageImpl::use_byte_size_t(),
slicelength,
at::DataPtr(
static_cast<void*>(data + start),
old_storage_impl,
[](void* s) {
c10::raw::intrusive_ptr::decref(static_cast<at::StorageImpl*>(s));
},
old_storage_impl->device()),
old_storage_impl->allocator(),
/* resizable */ false,
device_opt);
PyObject* _ret =
THPStorage_NewWithStorage(Py_TYPE(self), std::move(new_storage_impl));
return _ret;
}Ми можемо не тільки використовувати вказівник, але й клас at::DataPtr може обробляти знищення, що значно спрощує управління життєвим циклом.
На стороні Python вказівник на область пам'яті представлений об'єктом memoryview, ці об'єкти підтримують протокол буфера. Ми можемо отримати об'єкт memoryview з багатьох речей, bytes і mmap — це дві основні речі, які його підтримують, і саме вони нас цікавлять.
Нарешті, ми знаємо, що повинні зробити: створити функцію, яка приймає об'єкт memoryview і перетворює його в UntypedStorage без копіювання. З можливістю реконструкції UntypedStorage з memoryview, фактичні дані тензора не повинні бути в потоці pickle, що значно зменшує обсяг даних, які ми повинні копіювати.
void initOvermindHelpers(py::module m) {
// ...
m.def("_make_untyped_storage", [](py::buffer b) {
auto info = new py::buffer_info(b.request());
return pybind11::reinterpret_steal<py::object>(THPStorage_NewWithStorage(
THPStorageClass,
c10::make_intrusive<at::StorageImpl>(
c10::StorageImpl::use_byte_size_t(),
info->size,
at::DataPtr(
info->ptr, info,
[](void* ptr) {
py::gil_scoped_acquire gil;
auto b = static_cast<py::buffer_info*>(ptr);
delete b;
},
at::DeviceType::CPU
),
/*allocator=*/nullptr,
/*resizable=*/false,
)
));
});
}Це основний будівельний блок overmind.
Спільне використання тензорів!
Примітка: У PyTorch вже є механізм спільного використання тензорів, але він не відповідає нашим потребам. Докладніше про це пізніше.
По-перше, спільне використання пам'яті між клієнтом і сервером
Коли ми бачимо, що 'share' і 'memory' йдуть разом, у нас виникає бажання використовувати shmget та його друзі. Це "розроблено" для використання як механізм спільного використання пам'яті, чи не так? Але у нього є 2 основні недоліки:
- POSIX shm є обмеженим ресурсом, те, що ви можете використовувати, визначається тим, як системний адміністратор налаштовує систему. Екстремальний, але повсюдний приклад — Docker-контейнери, за замовчуванням у вас є лише 64MiB POSIX shm, які можна використовувати.
- POSIX shm живе довше, ніж ваш процес, тому вам потрібно самостійно керувати ним. Якщо процес управління буде примусово завершено або не буде оброблено належним чином, об'єкт shm може залишитися в системі на невизначений час.
Якщо уважно придивитися, Linux сповнений цікавих системних викликів. memfd_create - це той, який нас цікавить: він надає вам fd, що представляє виділення анонімної пам'яті. Ви можете виконувати всілякі файлові операції з ним: читати, записувати і, звичайно, mmap. Якщо ми можемо поділитися fd, ми можемо поділитися пам'яттю.
Спільне використання fd має 'стандартний', але архаїчний спосіб: sendmsg з SCM_RIGHTS. Ми можемо використовувати бібліотеки, щоб допомогти нам приховати складні деталі процесу sendmsg, але нам все ще потрібно координувати процеси сервера і клієнта. Ми вирішили використати тут хак: просто відкрити /proc/{pidof(server)}/fd/{memfd} на стороні клієнта, ніколи не закриваючи fd на стороні сервера overmind. Єдине необхідне спілкування - це кортеж (pid, fd). Це працює ідеально у нашому випадку.
Вищезгадані слова зводяться до цих рядків:
class SharedMemory:
@classmethod
def create(cls, shift):
# Викликається на стороні сервера
libc = ctypes.CDLL(None)
name = _make_filename(shift).encode('utf-8')
fd = libc.memfd_create(name, os.O_RDWR)
os.ftruncate(fd, 1 << shift)
mem_id = (os.getpid(), fd)
return cls(fd=fd, mem_id=mem_id)
@classmethod
def rebuild(cls, mem_id):
# Викликається на стороні клієнта
pid, fd = mem_id
local_fd = os.open(f'/proc/{pid}/fd/{fd}', os.O_RDWR)
return cls(fd=local_fd, mem_id=mem_id)
def get_buffer(self):
# Викликається на обох сторонах
self._mmap = mmap.mmap(self._fd, size)
self._buf = memoryview(self._mmap)
return self._bufІнтеграція з серіалізацією
Як ми обговорювали раніше, нам потрібно змінити процес серіалізації UntypedStorage. Подібно до того, що було реалізовано в torch.multiprocessing.reductions, ми визначаємо наші власні функції зменшення для pickle:
# Hoarder і borrower - це обгортки до SharedMemory вище, містять
# нудні речі, такі як арена пам'яті і т.д.
def _reduce_storage(storage):
# Викликається сервером
device = storage.device
storage = storage.cpu()
# Зберігайте вміст у спільній пам'яті
# `frag` містить повну інформацію, необхідну для знаходження вмісту.
frag = hoarder.put(storage)
return (_rebuild_storage_on_client, (frag, device))
def _rebuild_storage_on_client(frag, device):
# Викликається клієнтом
mv = borrower.borrow(frag) # Отримати memoryview зі спільної пам'яті
storage = _make_untyped_storage(mv) # Zero-copy!
if device.type == 'cuda':
return storage.cuda(device.index)
return storage
class OvermindPickler(dill.Pickler):
...
OvermindPickler.register(torch.storage.UntypedStorage, _reduce_storage)Тепер прості OvermindPickler.dumps і OvermindPickler.loads будуть використовувати спільну пам'ять для прискорення. Ви можете припинити читання тут, якщо вже втомилися. Решта - це деталі.
Диявол у деталях
Чому не використовувати внутрішнє спільне використання тензорів PyTorch?
Під 'внутрішнім спільним використанням тензорів' я маю на увазі torch.multiprocessing.reductions.
- На високому рівні, методи PyTorch призначені для 'передачі тензора в підпроцес', здається, те ж саме, але з тонкими відмінностями.
- PyTorch використовує POSIX shm для спільного використання пам'яті, що підлягає обмеженням, згаданим раніше.
- Для кожного тензора (або
UntypedStorage) PyTorch виділяє окремий об'єкт POSIX shm, навіть якщо він містить лише 4 байти. Кожен об'єкт споживає fd. - PyTorch звільняє POSIX shm, як тільки вони розпаковуються, що робить його непридатним для наших потреб. Нам потрібно десеріалізувати один і той же потік pickle кілька разів.
- Є багато логіки, пов'язаної з обміном CUDA, яка є чистим шумом і проблемою для нашого випадку використання.
Чому ви кажете, що 'дані тензора копіюються кілька разів'?
Для типового зчитування з диска torch.load:
- Файл
torch.saveзчитується в пам'ять. - Отримайте фактичні дані
torch.UntypedStorageякbytesшляхом вилучення Zip-файлу (torch.saveгенерує zip-файл). - Код на C++ скопіює дані у власну керовану пам'ять у конструкторі
torch.UntypedStorage.
Для наївного pickle.dumps і пізніше pickle.loads:
- Згенерований потік pickle внутрішньо вбудовує інший потік pickle,
pickle.loadsскопіює внутрішній потік у новийbytes. - Дані
torch.UntypedStorageвбудовуються у внутрішній потік pickle, інше копіювання відбувається при створенніtorch.UntypedStorage. - Код на C++ скопіює дані у власну керовану пам'ять у конструкторі
torch.UntypedStorage.
diffusers мають динамічний модуль
Репозиторії моделей можуть включати файли Python, які імпортуються під час виконання в простір імен diffusers_modules. Клієнт не має їх у sys.path, що порушує розпакування. На щастя, diffusers записує ці динамічні файли Python на диск, тому ми можемо просто імпортувати модуль і закінчити на цьому.
def diffusers_dyn_module_workaround():
from diffusers.utils.constants import HF_MODULES_CACHE
modpath = Path(HF_MODULES_CACHE) / "diffusers_modules/__init__.py"
spec = importlib.util.spec_from_file_location("diffusers_modules", modpath)
sys.modules["diffusers_modules"] = importlib.util.module_from_spec(spec)Підтримка bitsandbytes
Найбільш дратівливим у підтримці bitsandbytes є те, що процес квантування відбувається на GPU. Після ініціалізації CUDA та torch у сервері overmind немає простого способу його деініціалізувати, що може викликати проблеми для реальних навантажень (в основному менше використовуваної VRAM). Тому ми модифікували наш сервер, щоб він створював підпроцес, завантажував його в спільну пам'ять і завершував. Це покращує стабільність сервера overmind.
Квантизовані параметри є спеціальними підкласами, наданими bitsandbytes. Вони не були розроблені з урахуванням 'збережуваності', тому ми повинні зробити це самостійно.
def _reduce_bnb_param(p):
dev = p._prev_device
assert p.quant_state
return (_rebuild_bnb_param, (type(p), p.data, p.quant_state.as_dict(packed=True), dev))
def _rebuild_bnb_param(typ, data, qs_dict, dev):
return typ.from_prequantized(data, qs_dict, device=dev)
def bitsandbytes_quirks():
try:
import bitsandbytes
except ImportError as e:
return
ForkingPickler.register(bitsandbytes.nn.modules.Params4bit, _reduce_bnb_param)
ForkingPickler.register(bitsandbytes.nn.modules.Int8Params, _reduce_bnb_param)Квантизовані моделі через bitsandbytes мають хуки та монкі-патчі, які не зберігаються, ми повинні їх видалити:
from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None) # Видалити попередження монкі-патчів
model.__dict__.pop('cuda', None)Ми також стикалися з проблемами, коли функції вкладені в інші функції (замість того, щоб бути на верхньому рівні), що робить їх незбережуваними. Ми намагалися обійти це, але без успіху. Ми змушені були переключити наш pickle з наданого стандартною бібліотекою на dill, щоб зберегти це. dill набагато потужніший, але це чисто Python-імплементація, яка значно повільніша, ніж версія стандартної бібліотеки. На щастя, ця вартість буде оплачена лише один раз, коли ми завантажуємо модель вперше (впливає лише на збереження, не на розпакування).
Підтримка stable-fast
stable-fast генерує результати torch.compile, які не можуть бути збережені. Але з torch.jit.save ми можемо зберегти результати як zip-файл. Це звучить неефективно, але краще, ніж нічого.
З одним лише torch.jit.save недостатньо, щоб зберегти результати stable-fast. stable-fast використовує процес 'flatten', щоб зробити модуль Torch трасуваним. Коли він стикається з чимось, що не впізнає (наприклад, клас dataclass), він не серіалізує це, а лише зберігає посилання на фактичний клас. Ми виправили відповідну логіку, щоб фактично зберігати збережений клас у 'flatten'ed потоці.
def stable_fast_quirks():
...
# pickle dataclass type instead of just put it into a container (which will not survive after torch.jit.save)
def flatten_dataclass(obj):
from sfast.utils.flat_tensors import flatten_bytes, flatten_dict
import dataclasses
d = dict((field.name, getattr(obj, field.name))
for field in dataclasses.fields(obj))
import pickle
pickled = pickle.dumps(obj.__class__)
return flatten_bytes(pickled) + flatten_dict(d)
def unflatten_dataclass(tensors, start):
from sfast.utils.flat_tensors import unflatten_bytes, unflatten_dict
import pickle
pickled, start = unflatten_bytes(tensors, start)
clz = pickle.loads(pickled)
content, start = unflatten_dict(tensors, start)
return clz(**content), start
sfast.utils.flat_tensors.flatten_dataclass = flatten_dataclass
sfast.utils.flat_tensors.unflatten_dataclass = unflatten_dataclassТут є ще два трюки:
- Ми перепаковуємо ZIP-файл з
ZIP_STORED, щоб не розпаковувати ZIP-файл при кожному наступному завантаженні. - Інтерфейс
torch.jit.loadтакож викликає проблему копіювання пам'яті, тому ми написали простий обгортковий метод для завантаження через протокол буфера Python, так само якUntypedStorage.
void initOvermindHelpers(py::module m) {
// ...
m.def("import_ir_module_from_buffer_0copy",
[](std::shared_ptr<torch::jit::CompilationUnit> cu, py::buffer buffer) {
auto info = buffer.request();
imemstream in((char*)info.ptr, info.size); // Без копіювання!
return import_ir_module(std::move(cu), in, ...);
}
);
}Шаблон vae=vae
У нашій кодовій базі є щось подібне, що намагається завантажити модель з попередньо завантаженою моделлю як аргументом:
import overmind.api
overmind.api.monkey_patch_all()
import torch
from diffusers.models import AutoencoderKL
from diffusers import (
ControlNetModel,
StableDiffusionControlNetPipeline,
)
vae = AutoencoderKL.from_pretrained(
"lemon2431/ChineseInkComicStrip_v10",
subfolder="vae",
torch_dtype=torch.float16,
)
controlnet_depth = ControlNetModel.from_pretrained(
"lllyasviel/control_v11f1p_sd15_depth",
torch_dtype=torch.float16,
variant="fp16",
)
controlnet_edge = ControlNetModel.from_pretrained(
"lllyasviel/control_v11p_sd15_softedge",
torch_dtype=torch.float16,
variant="fp16",
)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
"lemon2431/ChineseInkComicStrip_v10",
vae=vae, # Тут!
controlnet=[controlnet_edge, controlnet_depth], # і тут!
torch_dtype=torch.float16,
safety_checker=None,
)
pipeline.to('cuda')Як ми вже згадували раніше, аргументи функції вважаються простими, легко серіалізованими об'єктами, але цей шаблон порушує це припущення. Щоб вирішити цю проблему, ми додали спеціальну логіку: кожен кешований результат отримує прикріплений ID. Якщо цей об'єкт використовується як аргумент в іншому виклику, клієнт замінює його на його ID, і сервер може відновити фактичний об'єкт на основі ID.
Отримана модель pipeline міститиме посилання на vae. Для простоти, ми просто серіалізуємо її безпосередньо тут. Однак, при переміщенні фактичного UntypedStorage в спільну пам'ять, ми усуваємо будь-які повторювані дані.
Ми могли б використати механізм persistent_id з pickle, але я не пробував цей шлях. Це трохи шкода.
Бенчмаркінг
І тепер частина, яку всі люблять бачити.
Ми використовуємо скрипт шаблону VAE з останнього розділу для нашого тесту.
| Тест | vae | depth | edge | pipeline | to('cuda') | Всього |
|---|---|---|---|---|---|---|
| без, 1-й | 1.18 | 0.98 | 1.41 | 1.65 | 0.91 | 6.16 |
| без, 2-й | 1.15 | 0.96 | 0.97 | 1.65 | 0.89 | 5.66 |
| без, 3-й | 1.15 | 0.96 | 0.98 | 1.61 | 0.91 | 5.65 |
| w/o, 4th | 1.42 | 1.10 | 1.11 | 1.72 | 0.88 | 6.27 |
| w/o, 5th | 1.28 | 1.08 | 1.10 | 1.72 | 0.92 | 6.13 |
| w/, 1st | 5.44 | 5.17 | 5.41 | 7.29 | 0.86 | 24.20 |
| w/, 2nd | 0.00 | 0.01 | 0.01 | 0.20 | 0.87 | 1.12 |
| w/, 3rd | 0.01 | 0.01 | 0.01 | 0.21 | 0.86 | 1.12 |
| w/, 4th | 0.01 | 0.01 | 0.01 | 0.20 | 0.90 | 1.15 |
| w/, 5th | 0.01 | 0.01 | 0.01 | 0.21 | 0.86 | 1.13 |
Як ви можете бачити, початкове завантаження з overmind займає 24.2 секунди, що значно довше в порівнянні з завантаженням без нього. Проте, при наступних завантаженнях залишається лише вартість .to('cuda').
Додавши розміри всіх серіалізованих файлів моделі, вся система оцінюється у використанні близько 5808 мегабайт пам'яті. Швидкий бенчмарк дає схожий результат.
In [1]: t = torch.ones((5808, 1024, 1024), dtype=torch.uint8)
In [2]: %time a = t.cuda()
CPU times: user 976 ms, sys: 874 μs, total: 977 ms
Wall time: 976 ms
In [3]: %timeit a = t.cuda()
1.01 s ± 56.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)Тестовано на Intel i9-11900K + GeForce RTX 4090.
Неочікувані побічні ефекти (Позитивні!)
Нашою основною мотивацією для створення overmind було забезпечення швидкого переключення ваг моделі під час інференсу. Хоча він виконує свою задачу, ми виявили кілька додаткових переваг.
Ми розгортаємо кілька екземплярів нашого застосунку, по одному для кожного GPU. Таким чином, на кожному вузлі буде 8 процесів. Після розгортання overmind використання системної пам'яті значно зменшилося. Ми не страждали від нестачі системної пам'яті, але якби це було так, це було б великою перемогою.
Пізніше ми виявили, що це значно підвищило продуктивність наших розробників алгоритмів та конвеєра. Для кожного циклу зміни-перевірки ми могли заощадити 10-20 секунд часу завантаження, що могло скласти величезну кількість. Більш важливо, що збережені секунди могли утримати розробників у потоці.
Github
Ми відкриваємо вихідний код на Github, будемо раді, якщо це допоможе.
