Overmind: Сокращение времени загрузки модели ML с 15 секунд до 0.2 секунд

Краткое содержание Загрузка моделей машинного обучения медленная, даже с прогретым кэшем страниц 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 — это 2 основные вещи, которые его поддерживают, и это также то, что нас интересует.

Наконец, мы знаем, что должны сделать: создать функцию, которая принимает объект 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 — это тот, который нас интересует: он предоставляет вам файловый дескриптор, представляющий выделение анонимной памяти. Вы можете выполнять все виды файловых операций с ним: читать, писать и, конечно же, использовать mmap. Если мы можем поделиться файловым дескриптором, мы можем поделиться памятью.

Обмен файловым дескриптором имеет 'стандартный', но загадочный способ сделать это: sendmsg с SCM_RIGHTS. Мы можем использовать библиотеки, чтобы помочь нам скрыть пугающие детали процесса sendmsg, но нам все равно нужно координировать процессы между сервером и клиентом. Мы решили использовать здесь хак: просто открыть /proc/{pidof(server)}/fd/{memfd} на стороне клиента, при этом никогда не закрывая файловый дескриптор на стороне сервера 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

Интеграция с pickling

Как мы обсуждали ранее, нам нужно изменить процесс pickling для 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)  # Нулевая копия!
    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.

1. На высоком уровне методы PyTorch предназначены для 'передачи тензора в подпроцесс', кажется, то же самое, но с тонкими различиями.
2. PyTorch использует POSIX shm для совместного использования памяти, что подвержено ограничению, упомянутому ранее.
3. Для каждого тензора (или UntypedStorage) PyTorch выделяет отдельный объект POSIX shm, даже если он содержит всего 4 байта. Каждый объект потребляет файловый дескриптор.
4. PyTorch освобождает POSIX shm, как только они распаковываются, что делает его неподходящим для наших нужд. Нам нужно десериализовать один и тот же поток pickle несколько раз.
5. Существует много логики, связанной с совместным использованием 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 нет простого способа деинициализировать их, что может вызвать проблемы для реальных рабочих нагрузок (в основном, меньше доступной видеопамяти). Поэтому мы модифицировали наш сервер, чтобы он порождал подпроцесс, загружал его в общую память и завершал работу. Это улучшает стабильность сервера overmind.

Квантованные параметры являются специальными подклассами, предоставляемыми bitsandbytes. Они не были разработаны с учетом 'picklability', поэтому мы должны сделать это сами.

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 содержат хуки и monkey-патчи, которые не сериализуются, их необходимо удалить:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Удалить предупреждающие monkey-патчи
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' потоке.

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

Здесь есть еще два трюка:

1. Мы перепаковываем ZIP-файл с помощью ZIP_STORED, чтобы нам не приходилось распаковывать ZIP-файл при каждой последующей загрузке.
2. Интерфейс 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')

Как мы упоминали ранее, аргументы функций предполагаются простыми, легко сериализуемыми объектами, но этот паттерн нарушает это предположение. Чтобы справиться с этим, мы добавили специальную логику: каждому кэшированному результату присваивается идентификатор. Если этот объект используется в качестве аргумента в другом вызове, клиент заменяет его на его идентификатор, и сервер может восстановить фактический объект на основе идентификатора.

Полученная модель pipeline будет содержать ссылку на vae. Для простоты мы просто сериализуем его напрямую здесь. Однако при перемещении фактического UntypedStorage в общую память мы устраняем дублирование любых повторяющихся данных.

Мы могли бы использовать механизм persistent_id в pickle, но я не пробовал этот путь. Это немного жаль.

Бенчмаркинг

А теперь та часть, которую все любят видеть.

Мы используем скрипт с паттерном VAE из последнего раздела для нашего теста.

Как видно, начальная загрузка с 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, будем рады, если это поможет.