Overmind: Zkrácení načítání ML modelu z 15s na 0.2s

TL;DR Načítání modelu ML je pomalé, i když je Linux page cache zahřátá. Takže jsme vytvořili knihovnu, aby to bylo rychlé. Existují některé zajímavé technické detaily, které chceme sdílet, takže jsme napsali tento blog. Knihovna měla také neočekávaný dopad, o kterém se diskutuje na konci.

Důvod

Všechno to začalo před 2 lety, když jsme vydali naši první zkoušku režimu generování lowpoly. Režim lowpoly nešel dobře, z dnešního pohledu vydává špatné výsledky, ale zaplatili jsme za to hodně - dedikovaná GPU zpracovává pouze jednociferné úkoly denně. Má jemně vyladěné váhy, dostatečně velké na to, aby vytlačily všechny ostatní váhy modelu z VRAM. Horší je, že máme možná 3 takové modely (nemohu si vzpomenout na přesný počet), tvořily významnou část naší inferenční infrastruktury, což vedlo k poměrně neodpouštějícímu poměru efektivity. A ne, nemůžeme naivně načítat modely právě včas, trvá to 30 sekund, více než samotný čas zpracování.

Tehdy jsme neměli dedikované pipeline inženýry, naši vývojáři algoritmů se snažili, co mohli, aby to obešli. O několik dní později byl náš kódový základ posetý this.to('cpu') a that.to('cuda'). Tento přístup funguje nějakou dobu, ale čas od času narušuje tok našich vývojářů algoritmů. Co kdyby se věci mohly dít automagicky? Je to Python, věci se v Pythonu dějí automagicky.

Jak definujete 'automagicky'?

Pojďme se vžít do role vývojáře algoritmů. Věci jsou docela jasné: Nechci se starat o výkon za běhu mimo svůj hlavní algoritmus, pokud to není absolutně nutné. Raději bych nevěděl nic o výměně modelů dovnitř a ven.

Samozřejmě toho nemůžeme dosáhnout, ale můžeme se pokusit minimalizovat zásahy, které musíme zavést do kódu algoritmu. To mi připomíná monkey-patching knihovny gevent, která patchuje (primárně) knihovnu socket, nahrazuje ji gevent.socket, která může přepínat na jiné greenlety, když by IO blokovalo, podobně jako goroutine (ve skutečnosti je gevent starší než Golang!).

Protože jsme tehdy používali pouze knihovny HuggingFace (transformers, diffusers) k načítání modelů, cíl se stal jasným: Zavedeme pouze volání monkey-patch a zbytek kódu by měl zůstat nezměněn, XXXPipeline.from_pretrained(...) by měl být mnohem rychlejší.

Některá fakta, zjevná rozhodnutí a předpoklady

Overmind je cacheovací knihovna, která ukládá výsledky volání načítání modelu do paměti systému a později je rychle rekonstruuje.

Přeskakujeme diskusi o tom, jak je monkey-patching implementován, to není tak zajímavý detail. Vše, co potřebujeme vědět, je, že přesměrovává všechna volání XXXPipeline.from_pretrained(...) na overmind.api.load(XXXPipeline.from_pretrained, ...).

Používáme pickle k serializaci našeho výsledku cache, protože... nemáme na výběr, a torch.save sám používá pickle, bylo by divné ho nepoužít.

Používáme architekturu klient/server, protože nechceme invalidovat naši cache, když proces skončí. Existuje mnoho volání subprocess, která by z toho mohla mít prospěch.

Předpokládáme, že parametry XXXPipeline.from_pretrained jsou jednoduché hashovatelné věci (str a podobné věci) a další modely načítané overmind (vysvětleno později).

Název overmind je převzat ze Starcraftu, jak jste možná uhodli.

Rychle to rekonstruujte!

Nemůžeme naivně uložit výsledek pickle.loads do paměti a říkat tomu den. Koneckonců, v zahřátém scénáři, Linux page cache odvedla svou práci při cachování modelů na disku a stále můžeme vidět dobu načítání měřenou v desítkách sekund.

Neefektivita pochází z kopírování paměti. V Pythonu by i vytvoření milionů objektů stálo ne více než několik stovek ms. Nicméně, pro kopírování paměti o velikosti 10GiB by to stálo půl sekundy. Musíme se vyhnout kopírování paměti, jak jen to je možné.

Naštěstí většina velkých paměťových bloků jsou Torch tensory, můžeme bezpečně adresovat pouze je a ignorovat zbytek.

Ve skutečnosti jsem získal znalosti o vnitřní struktuře Torch tensoru v redukčním kódu při zkoumání mechanismu sdílení tensorů:

# Kopírováno z torch.multiprocessing.reductions, většina kódu je odstraněna
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))

Docela jednoduché: tenzor je jeho typ, metadata a základní úložiště. Zde storage je typu TypedStorage, ale ve skutečnosti je TypedStorage jen jednoduchý obal pro UntypedStorage. UntypedStorage je třída, která skutečně drží všechna data tenzoru.

Náš úkol se nyní stává konkrétnějším: Jak se vyhnout kopírování UntypedStorage? Můžeme spravovat tuto paměť tenzoru sami a konstruovat UntypedStorage tím, že ukážeme na paměť, kterou spravujeme?

Odpověď je ano!

Při procházení C++ kódu, kde je UntypedStorage konstruován, můžeme snadno najít úryvek kódu jako tento:

// Zkopírováno z torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...vynechání nesouvisejícího kódu...

    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;
}

Nejenže můžeme použít ukazatel, ale třída at::DataPtr může také zpracovávat destrukci, což značně zjednodušuje správu životního cyklu.

Na straně Pythonu je ukazatel na paměťovou oblast reprezentován objektem memoryview, tyto objekty podporují buffer protokol. Můžeme získat objekt memoryview z mnoha věcí, bytes a mmap jsou dvě hlavní věci, které to podporují, a také nás zajímají.

Nakonec víme, co bychom měli udělat: vytvořit funkci, která přijímá objekt memoryview a přemění ho na UntypedStorage bez kopírování. S možností rekonstruovat UntypedStorage z memoryview nemusí být skutečná data tenzoru v pickle streamu, což značně snižuje velikost dat, kterou musíme kopírovat.

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,
            )
        ));
    });
}

To je základní stavební kámen overmind.

Sdílení tenzorů!

Poznámka: V PyTorch již existuje mechanismus sdílení tenzorů, ale nevyhovuje našim potřebám. Více o tom později.

Nejprve, sdílení paměti mezi klientem a serverem

Když vidíme, že se objevují slova 'sdílet' a 'paměť', máme všichni nutkání použít shmget a jeho přátele. Je to "navrženo" jako mechanismus pro sdílení paměti, že? Ale má to 2 hlavní nedostatky:

• POSIX shm je vzácný zdroj, co můžete použít, je určeno tím, jak sysadmin konfiguruje systém. Extrémním, ale všudypřítomným příkladem jsou Docker kontejnery, ve výchozím nastavení máte k dispozici pouze 64MiB POSIX shm.
• POSIX shm žije déle než váš proces, musíte si jej spravovat sami. Pokud je proces správy násilně ukončen nebo nebyl pečlivě zpracován, shm objekt může zůstat v systému na neurčito.

Pokud se podíváte pozorně, Linux je plný zajímavých systémových volání. memfd_create je jedno z těch, které nás zajímá: Dává vám fd, který představuje alokaci anonymní paměti. Můžete na něm provádět všechny druhy souborových operací: čtení, zápis a samozřejmě mmap. Pokud můžeme sdílet fd, můžeme sdílet paměť.

Sdílení fd má 'standardní', ale tajemný způsob, jak to udělat: sendmsg s SCM_RIGHTS. Můžeme využít knihovny, které nám pomohou skrýt děsivé detaily procesu sendmsg, ale stále musíme provádět koordinaci mezi serverovými a klientskými procesy. Rozhodli jsme se zde použít trik: Prostě otevřete /proc/{pidof(server)}/fd/{memfd} na straně klienta, zatímco nikdy nezavřete fd na straně serveru overmind. Jediná potřebná komunikace je (pid, fd) dvojice. V našem případě to funguje perfektně.

Výše uvedená slova se dají shrnout do těchto řádků:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Voláno na straně serveru
        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):
        # Voláno na straně klienta
        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):
        # Voláno na obou stranách
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integrace s picklingem

Jak jsme diskutovali dříve, potřebujeme upravit proces picklingu UntypedStorage. Podobně jako to bylo implementováno v torch.multiprocessing.reductions, definujeme naše vlastní funkce redukce pro pickle:

# Hoarder a borrower jsou obaly pro výše uvedenou SharedMemory, obsahují
# nudné věci jako paměťovou arénu atd.
def _reduce_storage(storage):
    # Voláno serverem
    device = storage.device
    storage = storage.cpu()

    # Uložení obsahu do sdílené paměti
    # `frag` obsahuje kompletní informace potřebné k lokalizaci obsahu.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Voláno klientem
    mv = borrower.borrow(frag)  # Získání memoryview ze sdílené paměti
    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)

Nyní, jednoduché OvermindPickler.dumps a OvermindPickler.loads využijí sdílenou paměť ke zrychlení. Můžete přestat číst zde, pokud jste již unaveni. Zbytek jsou detaily.

Ďáblové v detailu

Proč ne PyTorchovo interní sdílení tensorů?

Pro 'interní sdílení tensorů' mám na mysli torch.multiprocessing.reductions.

1. Na vysoké úrovni jsou PyTorchovy metody navrženy pro 'předávání tensoru do podprocesu', což se zdá být stejné, ale s jemným rozdílem.
2. PyTorch používá POSIX shm ke sdílení paměti, což podléhá limitu zmíněnému dříve.
3. Pro každý tensor (nebo UntypedStorage) PyTorch alokuje dedikovaný POSIX shm objekt, i když obsahuje pouze 4 bajty. Každý objekt spotřebovává fd.
4. PyTorch deallocuje POSIX shm, jakmile jsou rozbaleny, což je nevhodné pro naše potřeby. Potřebujeme deserializovat stejný pickle stream vícekrát.
5. Existuje mnoho logiky sdílení související s CUDA, která je pro náš případ použití čistým šumem a problémem.

Proč říkáte 'data tensoru jsou kopírována několikrát'?

Pro typické torch.load z disku:

• Soubor torch.save na disku je načten do paměti.
• Získejte skutečná data torch.UntypedStorage jako bytes extrakcí zip souboru (torch.save generuje zip soubor).
• C++ kód zkopíruje data do vlastní spravované paměti v konstruktoru torch.UntypedStorage.

Pro naivní pickle.dumps a pozdější pickle.loads:

• Generovaný pickle stream interně vkládá jiný pickle stream, pickle.loads zkopíruje vnitřní stream do nového bytes.
• Data torch.UntypedStorage jsou vložena ve vnitřním pickle streamu, další kopie se provádí při konstrukci torch.UntypedStorage.
• C++ kód zkopíruje data do vlastní spravované paměti v konstruktoru torch.UntypedStorage.

diffusers mají dynamický modul

Modelové repozitáře mohou zahrnovat Python soubory, které se importují za běhu do jmenného prostoru diffusers_modules. Klient je nemá v sys.path, což narušuje unpickling. Naštěstí diffusers zapíší tyto dynamické Python soubory na disk, takže můžeme jednoduše importovat modul a je to.

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)

Podpora pro bitsandbytes

Nejotravnější věcí na podpoře bitsandbytes je, že proces kvantizace probíhá na GPU. Jakmile inicializujeme CUDA a torch na serveru overmind, není snadný způsob, jak je deinicializovat, což může způsobit problémy pro skutečné pracovní zátěže (hlavně méně použitelné VRAM). Proto jsme upravili náš server, aby spouštěl podproces, načetl ho do sdílené paměti a ukončil. To zlepšuje stabilitu serveru overmind.

Kvantizované parametry jsou speciální podtřídy poskytované bitsandbytes. Nebyly navrženy s ohledem na 'picklability', takže to musíme udělat sami.

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)

Kvantizované modely přes bitsandbytes přicházejí s hooky a monkey-patche, které nelze picklovat, musíme je odstranit:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Odstraňte varovné monkeypatches
model.__dict__.pop('cuda', None)

Také jsme narazili na problémy, kdy jsou funkce vnořeny do jiných funkcí (namísto toho, aby byly na nejvyšší úrovni), což je činí nepicklovatelnými. Snažili jsme se to obejít, ale bez úspěchu. Museli jsme přepnout náš pickle z toho poskytovaného standardní knihovnou na dill, abychom to mohli picklovat. dill je mnohem výkonnější, ale je to čistě Python implementace, která je mnohem pomalejší než verze standardní knihovny. Naštěstí tento náklad bude zaplacen pouze jednou, když model načítáme poprvé (ovlivňuje pouze picklování, ne unpicklování).

Podpora pro stable-fast

stable-fast generuje výsledky torch.compile, které nelze picklovat. Ale s torch.jit.save bychom mohli výsledky uložit jako zip soubor. To zní neefektivně, ale je to lepší než nic.

Pouze s torch.jit.save není dostatečné picklovat výsledky stable-fast. stable-fast používá proces 'flatten', aby byl Torch modul sledovatelný. Když narazí na něco, co nerozpozná (například třídu dataclass), nebude ji serializovat, ale pouze uchová odkaz na skutečnou třídu. Opravili jsme příslušnou logiku, aby skutečně uložila picklovanou třídu ve 'flatten'ed streamu.

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

Existují zde dva další triky:

1. Znovu zabalíme ZIP soubor s ZIP_STORED, takže nemusíme dekomprimovat ZIP soubor při každém následném načtení.
2. Rozhraní torch.jit.load také způsobuje problém s kopírováním paměti, takže jsme napsali jednoduchý wrapper, který načítá pomocí Python buffer protokolu, stejně jako 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);  // Žádná kopie!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

Vzor vae=vae

Náš kódový základ má něco takového, pokouší se načíst model s dříve načteným modelem jako jeho argumentem:

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,  # Tady!
    controlnet=[controlnet_edge, controlnet_depth],  # a tady!
    torch_dtype=torch.float16,
    safety_checker=None,
)

pipeline.to('cuda')

Jak jsme zmínili dříve, argumenty funkcí jsou považovány za jednoduché, snadno serializovatelné objekty, ale tento vzor tuto předpoklad narušuje. Abychom to zvládli, přidali jsme speciální logiku: každému uloženému výsledku je přiřazeno ID. Pokud je tento objekt použit jako argument v jiném volání, klient jej nahradí jeho ID a server pak může obnovit skutečný objekt na základě ID.

Výsledný model pipeline bude obsahovat referenci na vae. Pro jednoduchost jej zde přímo serializujeme. Nicméně, při přesunu skutečného UntypedStorage do sdílené paměti, deduplikujeme jakákoliv opakovaná data.

Možná jsme mohli použít mechanismus persistent_id z pickle, ale tuto cestu jsem nezkoušel. To je trochu škoda.

Benchmarking

A nyní k části, kterou každý rád vidí.

Používáme skript vzoru VAE z poslední sekce k provedení našeho testu.

Jak můžete vidět, počáteční načítání s overmind trvá 24,2 sekundy, což je výrazně déle ve srovnání s načítáním bez něj. Nicméně, při následných načítáních je stále přítomna pouze cena .to('cuda').

Sečtením velikostí všech serializovaných souborů modelu se odhaduje, že celá pipeline používá přibližně 5808 megabajtů paměti. Rychlý benchmark dává podobný výsledek.

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)

Testováno na Intel i9-11900K + GeForce RTX 4090.

Neočekávané vedlejší účinky (Pozitivní!)

Naší hlavní motivací pro vytvoření overmind bylo umožnit rychlé přepínání vah modelu během inference. Zatímco to splnilo svůj účel, objevili jsme několik dalších výhod po cestě.

Nasazujeme více instancí naší aplikace, jednu pro každou GPU. Takže na každém uzlu bude 8 procesů. Po nasazení overmind se dramaticky snížilo využití systémové paměti. Netrpěli jsme nedostatkem systémové paměti, ale pokud bychom trpěli, bylo by to velké vítězství.

Později jsme zjistili, že je to skvělý přínos pro naše vývojáře algoritmů a pipeline. Pro každou smyčku modifikace-ověření jsme mohli ušetřit 10 až 20 sekund načítacího času, což by se mohlo nasčítat do velkého čísla. Co je důležitější, ušetřené sekundy mohly udržet vývojáře v proudu.

Github

Otevíráme to jako open-source na Github, budeme rádi, pokud to pomůže.