Overmind: Minska ML-modellens laddningstid från 15s till 0,2s

TL;DR ML-modellinläsning är långsam, även med uppvärmd Linux-sidcache. Så vi byggde ett bibliotek för att göra det snabbt. Det finns några intressanta tekniska detaljer vi vill dela med oss av, så vi skrev denna blogg. Biblioteket hade också en oväntad påverkan, som diskuteras i slutet.

Motivering

Allt började för 2 år sedan, när vi lanserade vårt första försök med lowpoly-genereringsläge. Lowpoly-läget gick inte bra, det ger dåliga resultat ur dagens perspektiv, men vi betalade mycket för det -- en dedikerad GPU bearbetar endast ensiffriga uppgifter per dag. Den har finjusterade vikter, tillräckligt stora för att driva ut alla andra modellvikter ur VRAM. Värre, vi har kanske 3 sådana modeller (kan inte minnas det exakta antalet), de utgjorde en betydande del av vår inferensinfrastruktur, vilket skapade en ganska oförlåtande effektivitetskvot. Och nej, vi kan inte naivt ladda modellerna just-in-time, det kostar 30 sekunder, större än den faktiska bearbetningstiden.

Vi hade inga dedikerade pipeline-ingenjörer då, våra algoritmutvecklare gjorde sitt bästa för att kringgå detta. Dagar senare var vår kodbas full av this.to('cpu') och that.to('cuda'). Denna metod fungerar ett tag, men bryter flödet för våra algoritmutvecklare då och då. Tänk om saker kan hända automagiskt? Det är Python, saker händer automagiskt i Python.

Hur definierar du 'automagiskt'?

Låt oss hoppa in i rollen som en algoritmutvecklare. Saker är ganska klara: Jag vill inte bry mig om prestandan utanför min kärnalgoritm om jag absolut inte måste. Jag skulle hellre inte veta något om modellbyte in och ut.

Naturligtvis kan vi inte uppnå det, men vi kan försöka minimera intrånget vi måste introducera i algoritmkoden. Detta påminner mig om monkey-patching av gevent-biblioteket, det patchar (främst) socket-biblioteket, ersätter det med gevent.socket som kan byta till andra greenlets när IO skulle blockera, ungefär som en goroutine (faktiskt är gevent äldre än Golang!).

Eftersom vi bara använde HuggingFace-bibliotek (transformers, diffusers) för att ladda modeller vid den tiden, blev målet klart: Vi introducerar bara ett monkey-patch-anrop, och resten av koden bör förbli oförändrad, XXXPipeline.from_pretrained(...) bör vara mycket snabbare.

Några Fakta, Självklara Beslut och Antaganden

Overmind är ett cache-bibliotek, det cachelagrar modellinläsningsanropens resultat i systemminnet och rekonstruerar det senare snabbt.

Vi hoppar över diskussionen om hur monkey-patching implementeras, det är en inte så intressant detalj. Allt vi behöver veta är att det omdirigerar alla XXXPipeline.from_pretrained(...)-anrop till overmind.api.load(XXXPipeline.from_pretrained, ...).

Vi använder pickle för att serialisera vårt cache-resultat eftersom... vi har inget val, och torch.save använder själv pickle, det är konstigt att inte använda det.

Vi använder en klient/server-arkitektur eftersom vi inte vill ogiltigförklara vår cache när processen avslutas. Det finns många subprocess-anrop som kan dra nytta av det.

Vi antar att XXXPipeline.from_pretrained-parametrar är enkla hashbara saker (str och liknande) och andra modeller laddade av overmind (förklaras senare).

Namnet overmind är lånat från Starcraft, som du kanske har gissat.

Rekonstruera det snabbt!

Vi kan inte naivt spara pickle.loads-resultatet i minnet och kalla det en dag. Trots allt, i ett uppvärmt scenario, gjorde Linux-sidcachen sitt jobb med att cachelagra modeller på disk och vi kan fortfarande se en laddningstid mätt i tiotals sekunder.

Ineffektiviteten kommer från minneskopiering. I Python, även att skapa miljontals objekt skulle kosta högst flera hundra ms. Men för en minneskopiering av 10GiB, skulle det kosta en halv sekund. Vi måste undvika minneskopiering så mycket som möjligt.

Lyckligtvis är de flesta av de stora minnesblocken Torch-tensorer, vi kan säkert adressera bara dem och ignorera resten.

Faktiskt fick jag kunskapen om den interna strukturen av en Torch-tensor i reduktionskoden medan jag forskade om tensor-delningsmekanismen:

# Kopierad från torch.multiprocessing.reductions, det mesta av koden är borttagen
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))

Ganska enkelt: en tensor är dess typ, dess metadata och dess underliggande lagring. Här är storage av typen TypedStorage, men faktiskt är TypedStorage bara ett enkelt omslag till UntypedStorage. UntypedStorage är klassen som faktiskt håller all tensor-data.

Vår uppgift blir nu mer specifik: Hur undviker vi att kopiera UntypedStorage? Kan vi hantera dessa tensor-minnen själva och konstruera UntypedStorage genom att peka på det minne vi hanterar?

Svaret är ja!

Genom att snabbt gå igenom C++-koden där UntypedStorage konstrueras, kan vi enkelt hitta en kodsnutt som denna:

// Kopierad från torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...utelämnar orelaterad kod...

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

Inte bara kan vi använda en pekare, men klassen at::DataPtr kan också hantera destruktion, vilket gör livscykelhanteringen mycket enklare.

På Python-sidan representeras en pekare till en minnesregion av ett memoryview-objekt, dessa objekt stöder bufferprotokoll. Vi kan få ett memoryview-objekt från många saker, bytes och mmap är de 2 huvudsakliga sakerna som stöder det, och de är också vad vi bryr oss om.

Slutligen vet vi vad vi ska göra: skapa en funktion som accepterar ett memoryview-objekt och omvandlar det till en UntypedStorage utan att kopiera. Med förmågan att rekonstruera UntypedStorage från memoryview, behöver den faktiska tensor-datan inte vara i pickle-strömmen, vilket kraftigt minskar datamängden vi måste kopiera runt.

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

Det är kärnbyggstenen i overmind.

Dela tensorerna!

Notera: Det finns redan en mekanism för att dela tensorer i PyTorch, men den passar inte våra behov. Mer om detta senare.

Först, dela minne mellan klient och server

När vi ser 'dela' och 'minne' tillsammans, har vi alla en lust att använda shmget och dess vänner. Det är "designat" för att användas som en mekanism för minnesdelning, eller hur? Men det har 2 stora brister:

• POSIX shm är en knapp resurs, vad du kan använda bestäms av hur systemadministratören konfigurerar systemet. Ett extremt men allmänt exempel är Docker-containers, som standard har du bara 64MiB POSIX shm användbart.
• POSIX shm lever längre än din process, du måste hantera det själv. Om hanteringsprocessen tvingas avslutas eller inte hanteras noggrant, kan shm-objektet lämnas kvar på systemet på obestämd tid.

Om du tittar noga, är Linux fullt av intressanta systemanrop. memfd_create är ett vi är intresserade av: Det ger dig en fd som representerar en allokering av anonymt minne. Du kan utföra alla typer av filoperationer på det: läsa, skriva och, naturligtvis, mmap. Om vi kan dela fd, kan vi dela minnet.

Att dela en fd har ett 'standard' men arkaiskt sätt att göra det: sendmsg med SCM_RIGHTS. Vi kan utnyttja bibliotek för att hjälpa oss att dölja de skrämmande detaljerna i sendmsg-processen, men vi måste fortfarande koordinera mellan server- och klientprocesser. Vi bestämde oss för att använda en hack här: Öppna bara /proc/{pidof(server)}/fd/{memfd} på klientsidan, medan fd aldrig stängs på overmind-serversidan. Den enda kommunikation som behövs är en (pid, fd)-tupl. Det fungerar perfekt i vårt fall.

Ovanstående ord kokar ner till dessa rader:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Called on server side
        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):
        # Called on client side
        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):
        # Called on both side
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integrera med pickling

Som vi diskuterade tidigare, behöver vi modifiera pickling-processen för UntypedStorage. Liknande vad som implementerades i torch.multiprocessing.reductions, definierar vi våra egna reduce-funktioner för pickle:

# Hoarder och borrower är en wrapper till SharedMemory ovan, innehåller
# tråkiga saker som minnesarena, etc.
def _reduce_storage(storage):
    # Called by server
    device = storage.device
    storage = storage.cpu()

    # Store content in shared memory
    # The `frag` contains the complete information needed to locate the content.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Called by client
    mv = borrower.borrow(frag)  # Get a memoryview from shared memory
    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 kommer enkla OvermindPickler.dumps och OvermindPickler.loads att använda delat minne för att snabba upp. Du kan sluta läsa här om du redan är trött. Resten är detaljer.

Djävulen i detaljerna

Varför inte PyTorchs inbyggda tensor-delning?

För 'inbyggd tensor-delning' menar jag torch.multiprocessing.reductions.

1. På hög nivå är PyTorchs metoder designade för 'överföra tensor till subprocess', verkar samma men med subtila skillnader.
2. PyTorch använder POSIX shm för att dela minne, vilket är föremål för den tidigare nämnda begränsningen.
3. För varje tensor (eller UntypedStorage) allokerar PyTorch ett dedikerat POSIX shm-objekt för det, även om det bara innehåller 4 byte. Varje objekt förbrukar en fd.
4. PyTorch avallokerar POSIX shm när de avkodas, vilket gör det olämpligt för våra behov. Vi behöver deserialisera samma pickle-ström flera gånger.
5. Det finns mycket CUDA-relaterad delningslogik, vilket är ren brus och problem för vårt användningsfall.

Varför säger du 'tensor data kopieras flera gånger'?

För en typisk torch.load från disk:

• Filen torch.save på disk läses in i minnet.
• Hämta den faktiska torch.UntypedStorage-datan som bytes genom att extrahera Zip-filen (torch.save genererar en zip-fil).
• C++-koden kommer att kopiera datan till sitt eget hanterade minne i torch.UntypedStorage-konstruktorn.

För en naiv pickle.dumps och senare pickle.loads:

• Den genererade pickle-strömmen bäddar internt in en annan pickle-ström, pickle.loads kommer att kopiera den inre strömmen till en ny bytes.
• torch.UntypedStorage-data bäddas in i den inre pickle-strömmen, en annan kopia sker vid konstruktionen av torch.UntypedStorage.
• C++-koden kommer att kopiera datan till sitt eget hanterade minne i torch.UntypedStorage-konstruktorn.

diffusers har en dynamisk modul

Modell-repos kan inkludera Python-filer som importeras vid körning till ett diffusers_modules-namnrymd. Klienten har inte dessa i sys.path, vilket bryter uppackningen. Lyckligtvis kommer diffusers att skriva dessa dynamiska Python-filer på disk, så vi kan bara importera modulen och kalla det en dag.

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)

Stöd för bitsandbytes

Det mest irriterande med att stödja bitsandbytes är att kvantiseringsprocessen sker på en GPU. När vi har initierat CUDA och torch i overmind-servern, finns det inget enkelt sätt att avinitiera det, vilket kan orsaka problem för verkliga arbetsbelastningar (främst mindre användbart VRAM). Därför modifierade vi vår server för att starta en subprocess, ladda den i delat minne och avsluta. Detta råkar förbättra stabiliteten hos overmind-servern.

De kvantiserade parametrarna är speciella underklasser som tillhandahålls av bitsandbytes. De var inte designade med 'picklability' i åtanke, så vi måste göra det själva.

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)

Kvantiserade modeller via bitsandbytes kommer med hooks och monkey-patches som inte pickle, vi måste ta bort dem:

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

Vi har också stött på problem där funktioner är inbäddade i andra funktioner (snarare än att vara på toppnivå), vilket gör dem inte picklable. Vi försökte kringgå detta, men utan framgång. Vi var tvungna att byta vår pickle från den standardbibliotek som tillhandahålls till dill för att pickle detta. dill är mycket kraftfullare, men det är en ren Python-implementation, vilket är mycket långsammare än standardbiblioteksversionen. Lyckligtvis kommer denna kostnad endast att betalas en gång när vi laddar modellen första gången (påverkar endast pickling, inte unpickling).

Stöd för stable-fast

stable-fast genererar torch.compile-resultat, som inte kan picklas. Men med torch.jit.save kunde vi spara resultaten som en zip-fil. Detta låter ineffektivt, men det är bättre än ingenting.

Med endast torch.jit.save är det inte tillräckligt för att pickle stable-fast-resultat. stable-fast använder en 'flatten'-process för att göra Torch-modulen spårbar. När den stöter på något den inte känner igen (till exempel dataclass's klass), kommer den inte att serialisera det, utan bara behålla en referens till den faktiska klassen. Vi har patchat den relevanta logiken för att faktiskt lagra en picklad klass inom den 'flatten'ade strömmen.

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

Det finns två ytterligare knep här:

1. Vi packar om ZIP-filen med ZIP_STORED, så vi slipper dekomprimera ZIP-filen vid varje efterföljande laddning.
2. torch.jit.load-gränssnittet medför också problemet med minneskopiering, så vi skrev en enkel omslag för att ladda den via Python-buffertprotokollet, precis som 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);  // Ingen kopiering!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

Mönstret vae=vae

Vår kodbas har något liknande detta, den försöker ladda en modell med en tidigare laddad modell som dess 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,  # Här!
    controlnet=[controlnet_edge, controlnet_depth],  # och Här!
    torch_dtype=torch.float16,
    safety_checker=None,
)

pipeline.to('cuda')

Som vi nämnde tidigare, antas funktionsargumenten vara enkla, lätt picklebara objekt, men detta mönster bryter den antagelsen. För att hantera detta lade vi till speciell logik: varje cachelagrat resultat får ett ID bifogat. Om det objektet används som ett argument i ett annat anrop, ersätter klienten det med dess ID, och servern kan sedan återställa det faktiska objektet baserat på ID:t.

Den resulterande pipeline-modellen kommer att innehålla en referens till vae. För enkelhetens skull picklar vi det direkt här. Men när vi flyttar den faktiska UntypedStorage till delat minne, deduplicerar vi all upprepad data.

Vi kan ha använt pickles persistent_id-mekanism, men jag provade inte denna väg. Det är lite synd.

Prestandamätning

Och nu till den del som alla älskar att se.

Vi använder VAE-mönsterskriptet från den sista sektionen för att göra vårt test.

Som du kan se, tar den initiala laddningen med overmind 24,2 sekunder, vilket är avsevärt längre jämfört med att ladda utan den. Dock, vid efterföljande laddningar, är endast kostnaden för .to('cuda') fortfarande närvarande.

Genom att summera storlekarna av alla de serialiserade modelfilerna, uppskattas hela pipeline att använda runt 5808 megabyte minne. Ett snabbt benchmark ger ett liknande resultat.

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)

Testad på Intel i9-11900K + GeForce RTX 4090.

Oväntade Bieffekter (Positiva!)

Vår primära motivation för att bygga overmind var att möjliggöra snabb växling av modellvikter under inferens. Medan det tjänade sitt syfte, upptäckte vi flera ytterligare fördelar längs vägen.

Vi distribuerar flera instanser av vår applikation, en för varje GPU. Således kommer det att finnas 8 processer per nod. Efter att vi distribuerade overmind, reducerades systemminnesanvändningen dramatiskt. Vi led inte av brist på systemminne, men om vi hade gjort det, skulle detta ha varit en stor vinst.

Senare fann vi att det var en stor boost för våra algoritm- och pipelinedevelopers. För varje modifiera-verifiera-loop kunde vi spara 10 till 20 sekunder av laddningstid, detta kunde lägga upp till ett stort antal. Ännu viktigare, sekunderna som sparades kunde hålla utvecklarna i flödet.

Github

Vi öppnar källkoden på Github, vi blir glada om det hjälpte.