Overmind: Verminder het laden van ML-modellen van 15s naar 0,2s

TL;DR Het laden van ML-modellen is traag, zelfs met een warme Linux-pagina cache. Dus hebben we een bibliotheek gebouwd om het snel te maken. Er zijn enkele interessante technische details die we willen delen, dus hebben we deze blog geschreven. De bibliotheek had ook een onverwachte impact, besproken aan het einde.

Reden

Het begon allemaal 2 jaar geleden, toen we onze eerste poging tot lowpoly generatie modus lanceerden. De lowpoly modus ging niet goed, het levert slechte resultaten op vanuit het perspectief van vandaag, maar we hebben er veel voor betaald -- een speciale GPU verwerkt slechts enkele taken per dag. Het heeft fijn-afgestelde gewichten, groot genoeg om alle andere modelgewichten uit VRAM te verdrijven. Erger nog, we hebben misschien 3 van zulke modellen (ik kan me het exacte aantal niet herinneren), ze vormden een aanzienlijk deel van onze inferentie-infrastructuur, wat zorgde voor een behoorlijk onverbiddelijke efficiëntieratio. En nee, we kunnen de modellen niet simpelweg just-in-time laden, het kost 30 seconden, groter dan de daadwerkelijke verwerkingstijd.

We hadden toen geen toegewijde pipeline-ingenieurs, onze algoritmeontwikkelaars deden hun best om dit te omzeilen. Dagen later was onze codebase bezaaid met this.to('cpu') en that.to('cuda'). Deze aanpak werkt een tijdje, maar verstoort de flow van onze algoritmeontwikkelaars van tijd tot tijd. Wat als dingen automatisch kunnen gebeuren? Het is Python, dingen gebeuren automatisch in Python.

Hoe definieer je 'automatisch'?

Laten we in de rol van een algoritmeontwikkelaar stappen. De zaken zijn vrij duidelijk: ik wil me niet druk maken over de runtime-prestaties buiten mijn kernalgoritme, tenzij het absoluut noodzakelijk is. Ik zou liever niets weten over het in- en uitwisselen van modellen.

Natuurlijk kunnen we dat niet bereiken, maar we kunnen proberen de inbreuk die we moeten introduceren in de algoritmecode te minimaliseren. Dit doet me denken aan de monkey-patching van de gevent bibliotheek, die (voornamelijk) de socket bibliotheek patcht, en deze vervangt door gevent.socket die kan overschakelen naar andere greenlets wanneer IO zou blokkeren, net als een goroutine (eigenlijk is gevent ouder dan Golang!).

Aangezien we destijds alleen HuggingFace bibliotheken (transformers, diffusers) gebruikten om modellen te laden, werd het doel duidelijk: We introduceren slechts een monkey-patch aanroep, en de rest van de code zou ongewijzigd moeten blijven, XXXPipeline.from_pretrained(...) zou veel sneller moeten zijn.

Enkele Feiten, Voor de Hand Liggende Beslissingen en Aannames

Overmind is een caching bibliotheek, het cachet de resultaten van model laadoproepen in het systeemgeheugen en reconstrueert het later snel.

We slaan de discussie over hoe monkey-patching is geïmplementeerd over, dat is een niet-zo-interessant detail. Alles wat we moeten weten is dat het alle XXXPipeline.from_pretrained(...) oproepen omleidt naar overmind.api.load(XXXPipeline.from_pretrained, ...).

We gebruiken pickle om ons cache resultaat te serialiseren omdat... we geen keuze hebben, en torch.save zelf gebruikt pickle, het is vreemd om het niet te gebruiken.

We gebruiken een client/server architectuur omdat we onze cache niet ongeldig willen maken wanneer het proces eindigt. Er zijn veel subprocess-aanroepen die hiervan kunnen profiteren.

We gaan ervan uit dat XXXPipeline.from_pretrained parameters eenvoudige hashbare dingen zijn (str en dergelijke) en andere modellen geladen door overmind (later uitgelegd).

De naam overmind is geleend van Starcraft, zoals je misschien al geraden had.

Herbouw het snel!

We kunnen niet simpelweg het resultaat van pickle.loads in het geheugen opslaan en het een dag noemen. Uiteindelijk, in een opgewarmd scenario, deed de Linux-pagina cache zijn werk door modellen op schijf te cachen en we kunnen nog steeds een laadtijd zien die in tientallen seconden wordt gemeten.

De inefficiëntie komt van het kopiëren van geheugen. In Python zou zelfs het creëren van miljoenen objecten niet meer dan enkele honderden ms kosten. Echter, voor een geheugenkopie van 10GiB, zou het een halve seconde kosten. We moeten geheugen kopiëren zoveel mogelijk vermijden.

Gelukkig zijn de meeste grote geheugenblokken Torch tensors, we kunnen ze veilig alleen adresseren en de rest negeren.

Eigenlijk kreeg ik de kennis van de interne structuur van een Torch tensor in de reductiecode tijdens het onderzoeken van het tensor deelmechanisme:

# Gekopieerd van torch.multiprocessing.reductions, het grootste deel van de code is verwijderd
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))

Heel eenvoudig: een tensor is zijn type, zijn metadata en zijn onderliggende opslag. Hier is storage van het type TypedStorage, maar eigenlijk is TypedStorage slechts een eenvoudige wrapper voor UntypedStorage. UntypedStorage is de klasse die daadwerkelijk alle tensorgegevens bevat.

Onze taak wordt nu specifieker: Hoe vermijden we het kopiëren van UntypedStorage? Kunnen we deze tensorgeheugens zelf beheren en UntypedStorages construeren door te verwijzen naar het geheugen dat we beheren?

Het antwoord is ja!

Als we door de C++-code bladeren waar UntypedStorage wordt geconstrueerd, kunnen we gemakkelijk een codefragment zoals dit vinden:

// Gekopieerd van torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...onbelangrijke code weggelaten...

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

Niet alleen kunnen we een pointer gebruiken, maar de at::DataPtr klasse kan ook destructie afhandelen, waardoor het beheer van de levensduur veel eenvoudiger wordt.

Aan de Python-kant wordt een pointer naar een geheugenregio weergegeven door een memoryview object, deze objecten ondersteunen het bufferprotocol. We kunnen een memoryview object krijgen van veel dingen, bytes en mmap zijn de 2 belangrijkste dingen die het ondersteunen, en dat is ook waar we ons om bekommeren.

Ten slotte weten we wat we moeten doen: een functie maken die een memoryview object accepteert en het omzet in een UntypedStorage zonder te kopiëren. Met de mogelijkheid om UntypedStorage te reconstrueren vanuit memoryview, hoeven de daadwerkelijke tensorgegevens niet in de pickle-stroom te zitten, wat de hoeveelheid gegevens die we moeten kopiëren aanzienlijk vermindert.

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

Dat is het kernbouwblok van overmind.

Het delen van de tensors!

Opmerking: Er is al een tensor-deelmechanisme in PyTorch, maar het voldoet niet aan onze behoeften. Meer hierover later.

Eerst, geheugen delen tussen client en server

Wanneer we 'delen' en 'geheugen' samen zien, hebben we allemaal de neiging om shmget en zijn vrienden te gebruiken. Het is "ontworpen" om te worden gebruikt als een mechanisme voor geheugendeling, toch? Maar het heeft 2 grote nadelen:

• POSIX shm is een schaars hulpmiddel, wat je kunt gebruiken wordt bepaald door hoe de systeembeheerder het systeem configureert. Een extreem maar alomtegenwoordig voorbeeld zijn Docker-containers, standaard heb je slechts 64MiB POSIX shm bruikbaar.
• POSIX shm leeft langer dan je proces, je moet je eigen beheer uitvoeren. Als het beheersproces gedwongen wordt beëindigd, of het niet zorgvuldig wordt afgehandeld, kan het shm-object voor onbepaalde tijd op het systeem achterblijven.

Als je goed kijkt, zit Linux vol met interessante systeemaanroepen. memfd_create is er een die onze interesse wekt: Het geeft je een fd dat een toewijzing van anoniem geheugen vertegenwoordigt. Je kunt allerlei bestandsbewerkingen erop uitvoeren: lezen, schrijven en natuurlijk mmap. Als we de fd kunnen delen, kunnen we het geheugen delen.

Het delen van een fd heeft een 'standaard' maar mysterieuze manier om het te doen: sendmsg met SCM_RIGHTS. We kunnen bibliotheken gebruiken om ons te helpen de ontmoedigende details van het sendmsg-proces te verbergen, maar we moeten nog steeds onze coördinatie tussen server- en clientprocessen uitvoeren. We besloten hier een hack te gebruiken: Open gewoon /proc/{pidof(server)}/fd/{memfd} aan de clientzijde, terwijl je de fd aan de overmind serverzijde nooit sluit. De enige communicatie die nodig is, is een (pid, fd)-tupel. Het werkt perfect in ons geval.

De bovenstaande woorden komen neer op deze regels:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Aangeroepen aan de serverzijde
        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):
        # Aangeroepen aan de clientzijde
        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):
        # Aangeroepen aan beide zijden
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integreren met pickling

Zoals we eerder bespraken, moeten we het pickling-proces van UntypedStorage aanpassen. Vergelijkbaar met wat werd geïmplementeerd in torch.multiprocessing.reductions, definiëren we onze aangepaste reduceerfuncties voor pickle:

# Hoarder en borrower zijn een wrapper voor SharedMemory hierboven, bevat
# saaie dingen zoals geheugenarena, enz.
def _reduce_storage(storage):
    # Aangeroepen door server
    device = storage.device
    storage = storage.cpu()

    # Inhoud opslaan in gedeeld geheugen
    # De `frag` bevat de volledige informatie die nodig is om de inhoud te lokaliseren.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Aangeroepen door client
    mv = borrower.borrow(frag)  # Haal een memoryview uit gedeeld geheugen
    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)

Nu zullen eenvoudige OvermindPickler.dumps en OvermindPickler.loads gebruik maken van gedeeld geheugen om te versnellen. Je kunt hier stoppen met lezen als je er al genoeg van hebt. De rest zijn details.

Duivels in de Details

Waarom niet PyTorch's eigen tensor delen?

Met 'eigen tensor delen' bedoel ik torch.multiprocessing.reductions.

1. Op hoog niveau zijn PyTorch's methoden ontworpen voor 'het doorgeven van tensor aan subprocess', lijkt hetzelfde maar met subtiele verschillen.
2. PyTorch gebruikt POSIX shm om geheugen te delen, onderhevig aan de eerder genoemde limiet.
3. Voor elke tensor (of UntypedStorage) wijst PyTorch een toegewijd POSIX shm-object toe, zelfs als het slechts 4 bytes bevat. Elk object verbruikt een fd.
4. PyTorch dealloceert de POSIX shm zodra ze worden uitgepakt, wat het ongeschikt maakt voor onze behoeften. We moeten dezelfde pickle-stroom meerdere keren deserialiseren.
5. Er is veel CUDA-gerelateerde deelingslogica, wat pure ruis en problemen zijn voor ons gebruik.

Waarom zeg je 'tensordata meerdere keren gekopieerd'?

Voor een typische on-disk torch.load:

• Het on-disk torch.save bestand wordt in het geheugen gelezen.
• Verkrijg de daadwerkelijke torch.UntypedStorage data als bytes door Zip-bestand extractie (torch.save genereert een zip-bestand).
• C++ code zal de data kopiëren naar zijn eigen beheerd geheugen in de torch.UntypedStorage constructor.

Voor een naïeve pickle.dumps en later pickle.loads:

• De gegenereerde pickle stream embed intern een andere pickle stream, pickle.loads zal de interne stream kopiëren naar een nieuwe bytes.
• torch.UntypedStorage data wordt ingebed in de interne pickle stream, een andere kopie vindt plaats bij de constructie van torch.UntypedStorage.
• C++ code zal de data kopiëren naar zijn eigen beheerd geheugen in de torch.UntypedStorage constructor.

diffusers hebben een dynamische module

Model repositories kunnen Python-bestanden bevatten die tijdens runtime worden geïmporteerd in een diffusers_modules namespace. De client heeft deze niet in sys.path, wat het unpicklen breekt. Gelukkig schrijft diffusers deze dynamische Python-bestanden op schijf, zodat we gewoon de module kunnen importeren en klaar zijn.

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)

Ondersteuning voor bitsandbytes

Het meest vervelende aan de ondersteuning van bitsandbytes is dat het kwantisatieproces op een GPU plaatsvindt. Zodra we CUDA en torch in de overmind server hebben geïnitialiseerd, is er geen gemakkelijke manier om het te de-initialiseren, wat problemen kan veroorzaken voor echte workloads (voornamelijk minder bruikbare VRAM). Daarom hebben we onze server aangepast om een subprocess te starten, het in gedeeld geheugen te laden en te beëindigen. Dit verbetert de stabiliteit van de overmind server.

De gekwantiseerde parameters zijn speciale subklassen die door bitsandbytes worden geleverd. Ze zijn niet ontworpen met 'picklability' in gedachten, dus moeten we het zelf doen.

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)

Gekwantiseerde modellen via bitsandbytes komen met hooks en monkey-patches die niet picklen, we moeten ze verwijderen:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Verwijder waarschuwings-monkeypatches
model.__dict__.pop('cuda', None)

We zijn ook problemen tegengekomen waarbij functies genest zijn binnen andere functies (in plaats van op het hoogste niveau te staan), waardoor ze niet picklable zijn. We hebben geprobeerd dit te omzeilen, maar zonder succes. We moesten onze pickle wisselen van de stdlib-provided naar dill om dit te picklen. dill is veel krachtiger, maar het is een pure Python-implementatie, die veel trager is dan de standaard bibliotheekversie. Gelukkig wordt deze kost alleen eenmalig betaald wanneer we het model voor de eerste keer laden (alleen van invloed op picklen, niet op unpicklen).

Ondersteuning voor stable-fast

stable-fast genereert torch.compile resultaten, die niet gepickled kunnen worden. Maar met torch.jit.save, kunnen we de resultaten opslaan als een zip-bestand. Dit klinkt inefficiënt, maar het is beter dan niets.

Met alleen torch.jit.save is het niet voldoende om stable-fast resultaten te picklen. stable-fast gebruikt een 'flatten' proces om de Torch module traceerbaar te maken. Wanneer het iets tegenkomt dat het niet herkent (bijvoorbeeld, de dataclass's klasse), zal het het niet serialiseren, maar alleen een referentie naar de daadwerkelijke klasse behouden. We hebben de relevante logica gepatcht om daadwerkelijk een gepickled klasse binnen de 'flatten'ed stream op te slaan.

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

Er zijn hier nog twee trucs:

1. We pakken het ZIP-bestand opnieuw in met ZIP_STORED, zodat we het ZIP-bestand niet bij elke volgende laadactie hoeven te decomprimeren.
2. De torch.jit.load interface veroorzaakt ook het probleem van geheugen kopiëren, dus we hebben een eenvoudige wrapper geschreven om het te laden via het Python buffer protocol, net zoals 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);  // Geen kopie!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

Het vae=vae patroon

Onze codebase heeft iets zoals dit, het probeert een model te laden met een eerder geladen model als argument:

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

pipeline.to('cuda')

Zoals we eerder vermeldden, wordt aangenomen dat de functieargumenten eenvoudige, gemakkelijk te pickelen objecten zijn, maar dit patroon doorbreekt die aanname. Om dit aan te pakken, hebben we speciale logica toegevoegd: elk gecachet resultaat krijgt een ID toegewezen. Als dat object als argument in een andere oproep wordt gebruikt, vervangt de client het door zijn ID, en de server kan dan het daadwerkelijke object herstellen op basis van de ID.

Het resulterende pipeline model zal een referentie naar vae bevatten. Voor de eenvoud pickelen we het hier direct. Echter, bij het verplaatsen van de daadwerkelijke UntypedStorage naar gedeeld geheugen, dedupliceren we alle herhaalde data.

We hadden het persistent_id mechanisme van pickle kunnen gebruiken, maar ik heb deze route niet geprobeerd. Dat is een beetje jammer.

Benchmarking

En nu voor het deel dat iedereen graag wil zien.

We gebruiken het VAE-patroonscript van de laatste sectie om onze test uit te voeren.

Zoals je kunt zien, duurt de initiële laadtijd met overmind 24,2 seconden, wat aanzienlijk langer is vergeleken met laden zonder. Echter, bij volgende ladingen is alleen de .to('cuda') kost nog steeds aanwezig.

Als we de groottes van alle geserialiseerde modelbestanden optellen, wordt geschat dat de hele pijplijn ongeveer 5808 megabytes aan geheugen gebruikt. Een snelle benchmark geeft een vergelijkbaar resultaat.

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)

Getest op Intel i9-11900K + GeForce RTX 4090.

Onverwachte Neveneffecten (Positief!)

Onze primaire motivatie voor het bouwen van overmind was om snel schakelen van modelgewichten tijdens inferentie mogelijk te maken. Terwijl het zijn doel diende, ontdekten we onderweg verschillende extra voordelen.

We implementeren meerdere instanties van onze applicatie, één voor elke GPU. Dus er zullen 8 processen per node zijn. Nadat we overmind hadden geïmplementeerd, werd het systeemgeheugengebruik drastisch verminderd. We hadden geen last van een tekort aan systeemgeheugen, maar als dat wel het geval was geweest, zou dit een grote winst zijn geweest.

Later ontdekten we dat het een geweldige boost was voor onze algoritme- en pijplijnontwikkelaars. Voor elke wijzig-verifieer-lus konden we 10 tot 20 seconden laadtijd besparen, wat kon oplopen tot een groot aantal. Belangrijker nog, de bespaarde seconden konden ontwikkelaars in de flow houden.

Github

We maken het open-source op Github, we zouden blij zijn als het helpt.