Overmind: Bawasan ang Paglo-load ng ML Model mula 15s hanggang 0.2s

TL;DR Ang pag-load ng ML model ay mabagal, kahit na may warmed Linux page cache. Kaya gumawa kami ng library para pabilisin ito. May ilang kawili-wiling teknikal na detalye na nais naming ibahagi, kaya isinulat namin ang blog na ito. Ang library ay nagkaroon din ng hindi inaasahang epekto, na tinalakay sa dulo.

Rationale

Nagsimula ito 2 taon na ang nakalipas, nang ilabas namin ang unang subok ng lowpoly generation mode. Ang lowpoly mode ay hindi naging maganda, naglalabas ito ng mahihinang resulta mula sa pananaw ngayon, pero malaki ang ginastos namin para dito -- isang dedicated GPU lamang ang nagpoproseso ng single digit na mga gawain kada araw. Ito ay may fine-tuned na weights, sapat na malaki para itulak ang lahat ng iba pang model weights palabas ng VRAM. Mas masahol pa, mayroon kaming siguro 3 ganitong mga modelo (hindi ko matandaan ang eksaktong bilang), sila ay bumubuo ng mahalagang bahagi ng aming inference infra, na nagdulot ng hindi mapagpatawad na efficiency ratio. At hindi, hindi namin maaaring basta-basta i-load ang mga modelo just-in-time, ito ay nagkakahalaga ng 30s, mas malaki kaysa sa aktwal na oras ng pagproseso.

Wala kaming dedicated pipeline engineers noon, sinubukan ng aming algorithm devs ang kanilang makakaya para i-workaround ito. Ilang araw ang lumipas, ang aming codebase ay puno ng this.to('cpu') at that.to('cuda'). Ang pamamaraang ito ay gumagana sa ilang sandali, pero paminsan-minsan ay nakakasira sa daloy ng aming algo devs. Paano kung ang mga bagay ay maaaring mangyari nang automagically? Ito ay Python, ang mga bagay ay talagang nangyayari nang automagically sa Python.

Paano mo idedefine ang 'automagically'?

Tayo'y pumasok sa papel ng isang algorithm developer. Ang mga bagay ay medyo malinaw: Ayokong mag-alala tungkol sa runtime performance sa labas ng aking core algorithm maliban kung talagang kailangan ko. Mas gusto kong hindi malaman ang anumang bagay tungkol sa model swap in at out.

Siyempre hindi natin maabot iyon, pero maaari nating subukan na bawasan ang intrusion na kailangan nating ipakilala sa algorithm code. Naalala ko ang monkey-patching ng gevent library, ito ay nag-patch (pangunahing) sa socket library, pinapalitan ito ng gevent.socket na maaaring lumipat sa ibang greenlets kapag ang IO ay magba-block, katulad ng isang goroutine (sa totoo lang, ang gevent ay mas matanda kaysa sa Golang!).

Dahil gumagamit lamang kami ng HuggingFace libs (transformers, diffusers) para mag-load ng mga modelo noon, naging malinaw ang target: Magpapakilala lamang kami ng isang monkey-patch call, at ang natitirang bahagi ng code ay dapat manatiling hindi nabago, ang XXXPipeline.from_pretrained(...) ay dapat na mas mabilis.

Ilang Katotohanan, Obvious na Desisyon at Mga Palagay

Ang Overmind ay isang caching library, ito ay nagka-cache ng mga resulta ng tawag sa pag-load ng modelo sa system memory at pagkatapos ay muling binubuo ito nang mabilis.

Laktawan natin ang talakayan tungkol sa kung paano naipatupad ang monkey-patching, iyon ay isang hindi gaanong kawili-wiling detalye. Ang kailangan lang nating malaman ay, ito ay nagre-redirect ng lahat ng XXXPipeline.from_pretrained(...) na tawag sa overmind.api.load(XXXPipeline.from_pretrained, ...).

Gumagamit kami ng pickle para i-serialize ang aming cache result dahil... wala kaming pagpipilian, at ang torch.save mismo ay gumagamit ng pickle, kakaiba kung hindi ito gagamitin.

Gumagamit kami ng client/server architecture dahil ayaw naming i-invalidate ang aming cache kapag nag-terminate ang proseso. Maraming subprocess calls ang makikinabang dito.

Ina-assume namin na ang XXXPipeline.from_pretrained na mga parameter ay simple at hashable na mga bagay (str at mga katulad nito) at iba pang mga modelo na na-load ng overmind (ipinaliwanag mamaya).

Ang pangalang overmind ay hiniram mula sa Starcraft, gaya ng maaaring nahulaan mo.

Muling Buuin ito nang Mabilis!

Hindi natin maaaring basta-basta i-save ang pickle.loads na resulta sa memory at tawagin itong tapos na. Pagkatapos ng lahat, sa isang warmed up na senaryo, ginawa ng Linux page cache ang trabaho nito sa pag-cache ng mga on-disk na modelo at maaari pa rin nating makita ang oras ng pag-load na sinusukat sa sampu-sampung segundo.

Ang inefficiency ay nagmumula sa memory copying. Sa Python, kahit na ang paglikha ng milyun-milyong mga object ay hindi aabot ng higit sa ilang daang ms. Gayunpaman, para sa isang memory copy ng 10GiB, ito ay aabot ng kalahating segundo. Dapat nating iwasan ang memory copy hangga't maaari.

Sa kabutihang palad, karamihan sa mga malalaking memory chunks ay Torch tensors, maaari nating ligtas na i-address lamang ang mga ito at balewalain ang iba.

Sa totoo lang, nakuha ko ang kaalaman tungkol sa internal structure ng isang Torch tensor sa reduction code habang nagsasaliksik sa tensor sharing mechanism:

# Kinopya mula sa torch.multiprocessing.reductions, karamihan sa code ay tinanggal
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))

Medyo simple: ang tensor ay ang kanyang uri, metadata at ang kanyang underlying storage. Dito ang storage ay ng uri TypedStorage, ngunit sa totoo lang ang TypedStorage ay isang simpleng balot lamang sa UntypedStorage. Ang UntypedStorage ang klase na talagang naglalaman ng lahat ng data ng tensor.

Ang ating gawain ay nagiging mas tiyak ngayon: Paano natin maiiwasan ang pagkopya ng UntypedStorage? Maaari ba nating pamahalaan ang mga memorya ng tensor at bumuo ng mga UntypedStorage sa pamamagitan ng pagturo sa memorya na ating pinamamahalaan?

Ang sagot ay oo!

Sa pag-skim sa C++ code kung saan ang UntypedStorage ay nabuo, madali nating makikita ang isang snippet ng code na ganito:

// Kinopya mula sa torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...omitting unrelated code...

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

Hindi lamang natin magagamit ang isang pointer, kundi ang klase na at::DataPtr ay maaari ring humawak ng destruction, na ginagawang mas simple ang pamamahala sa buhay ng object.

Sa panig ng Python, ang isang pointer sa isang rehiyon ng memorya ay kinakatawan ng isang memoryview object, ang mga object na ito ay sumusuporta sa buffer protocol. Makakakuha tayo ng isang memoryview object mula sa maraming bagay, bytes at mmap ang 2 pangunahing bagay na sumusuporta dito, at sila rin ang ating pinapahalagahan.

Sa wakas, alam natin kung ano ang dapat nating gawin: lumikha ng isang function na tumatanggap ng isang memoryview object at ginagawang isang UntypedStorage nang hindi kinokopya. Sa kakayahang muling buuin ang UntypedStorage mula sa memoryview, ang aktwal na data ng tensor ay hindi kailangang nasa pickle stream, na lubos na nabawasan ang laki ng data na kailangan nating kopyahin.

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

Iyan ang pangunahing bahagi ng overmind.

Pagbabahagi ng mga tensor!

Tandaan: Mayroon nang mekanismo ng pagbabahagi ng tensor sa PyTorch, ngunit hindi ito akma sa aming mga pangangailangan. Higit pa tungkol dito mamaya.

Una, pagbabahagi ng memorya sa pagitan ng kliyente at server

Kapag nakita natin ang 'bahagi' at 'memorya' na magkasama, lahat tayo ay may pagnanasa na gamitin ang shmget at ang mga kaibigan nito. Ito ay "dinisenyo" upang magamit bilang isang mekanismo ng pagbabahagi ng memorya, tama? Ngunit mayroon itong 2 pangunahing kapintasan:

• Ang POSIX shm ay isang limitadong mapagkukunan, kung ano ang maaari mong gamitin ay tinutukoy ng kung paano i-configure ng sysadmin ang sistema. Isang ekstremong ngunit pangkaraniwang halimbawa ay ang mga Docker container, sa default mayroon ka lamang 64MiB POSIX shm na magagamit.
• Ang POSIX shm ay mas matagal ang buhay kaysa sa iyong proseso, kailangan mong gawin ang sarili mong pamamahala. Kung ang proseso ng pamamahala ay sapilitang pinatay, o hindi ito maingat na hinawakan, ang shm object ay maaaring maiwan sa sistema nang walang hanggan.

Kung titingnan mong mabuti, puno ang Linux ng mga kawili-wiling system calls. Ang memfd_create ay isa sa mga interesado tayo: Binibigyan ka nito ng fd na kumakatawan sa isang alokasyon ng anonymous memory. Maaari mong gawin ang lahat ng uri ng file operations dito: read, write, at, siyempre, mmap. Kung maibabahagi natin ang fd, maibabahagi natin ang memorya.

Ang pagbabahagi ng fd ay may 'standard' ngunit mahiwagang paraan upang gawin ito: sendmsg na may SCM_RIGHTS. Maaari nating gamitin ang mga library upang matulungan tayong itago ang nakakatakot na mga detalye ng proseso ng sendmsg, ngunit kailangan pa rin nating gawin ang ating koordinasyon sa pagitan ng mga proseso ng server at kliyente. Nagpasya kaming gumamit ng isang hack dito: Buksan lang ang /proc/{pidof(server)}/fd/{memfd} sa panig ng kliyente, habang hindi kailanman isinasara ang fd sa panig ng overmind server. Ang tanging komunikasyon na kailangan ay isang (pid, fd) tuple. Gumagana ito nang perpekto sa aming kaso.

Ang mga salitang nasa itaas ay nauuwi sa mga linyang ito:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Tinawag sa panig ng server
        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):
        # Tinawag sa panig ng kliyente
        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):
        # Tinawag sa parehong panig
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integrasyon sa pickling

Tulad ng napag-usapan natin dati, kailangan nating baguhin ang proseso ng pickling ng UntypedStorage. Katulad ng kung ano ang ipinatupad sa torch.multiprocessing.reductions, tinutukoy natin ang ating custom na mga reduce function para sa pickle:

# Ang Hoarder at borrower ay isang wrapper sa SharedMemory sa itaas, naglalaman
# ng mga boring na bagay tulad ng memory arena, atbp.
def _reduce_storage(storage):
    # Tinawag ng server
    device = storage.device
    storage = storage.cpu()

    # Iimbak ang nilalaman sa shared memory
    # Ang `frag` ay naglalaman ng kumpletong impormasyon na kailangan upang mahanap ang nilalaman.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Tinawag ng kliyente
    mv = borrower.borrow(frag)  # Kumuha ng memoryview mula sa 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)

Ngayon, ang simpleng OvermindPickler.dumps at OvermindPickler.loads ay gagamit ng shared memory upang mapabilis. Maaari mong itigil ang pagbabasa dito kung ikaw ay sawa na. Ang natitira ay mga detalye.

Mga Detalye ng Diyablo

Bakit hindi PyTorch's in-house tensor sharing?

Para sa 'in-house tensor sharing', ang ibig kong sabihin ay torch.multiprocessing.reductions.

1. Sa mataas na antas, ang pamamaraan ng PyTorch ay idinisenyo para sa 'pagpapasa ng tensor sa subprocess', tila pareho ngunit may banayad na pagkakaiba.
2. Gumagamit ang PyTorch ng POSIX shm upang magbahagi ng memorya, na napapailalim sa limitasyon na nabanggit kanina.
3. Para sa bawat tensor (o UntypedStorage), ang PyTorch ay naglalaan ng isang dedikadong POSIX shm object para dito, kahit na naglalaman ito ng 4 bytes lamang. Ang bawat object ay kumokonsumo ng isang fd.
4. Ang PyTorch ay nagde-deallocate ng POSIX shm sa sandaling ito ay na-unpickle, na ginagawa itong hindi angkop para sa aming mga pangangailangan. Kailangan naming i-deserialize ang parehong pickle stream nang maraming beses.
5. Maraming CUDA na may kaugnayan sa pagbabahagi ng lohika, na puro ingay at problema para sa aming kaso ng paggamit.

Bakit mo sinasabing 'tensor data copied multiple times'?

Para sa isang tipikal na on-disk torch.load:

• Ang on-disk torch.save file ay binabasa sa memorya.
• Kunin ang aktwal na torch.UntypedStorage data bilang bytes sa pamamagitan ng Zip file extraction (torch.save ay bumubuo ng zip file).
• Ang C++ code ay kokopya ng data sa sarili nitong managed memory sa torch.UntypedStorage constructor.

Para sa isang simpleng pickle.dumps at kalaunan pickle.loads:

• Ang nabuo na pickle stream ay naglalaman ng isa pang pickle stream sa loob, pickle.loads ay kokopya ng inner stream sa isang bagong bytes.
• Ang torch.UntypedStorage data ay nakapaloob sa inner pickle stream, isa pang kopya ang nangyayari sa konstruksyon ng torch.UntypedStorage.
• Ang C++ code ay kokopya ng data sa sarili nitong managed memory sa torch.UntypedStorage constructor.

Ang diffusers ay may dynamic module

Maaaring maglaman ang mga Model repos ng mga Python file na na-import sa runtime sa isang diffusers_modules namespace. Ang client ay walang mga ito sa sys.path, na nagiging sanhi ng pagkasira ng unpickling. Sa kabutihang palad, ang diffusers ay magsusulat ng mga dynamic na Python file sa disk, kaya maaari nating i-import lamang ang module at tapusin na ang araw.

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)

Suporta para sa bitsandbytes

Ang pinaka-nakakainis na bagay tungkol sa pagsuporta sa bitsandbytes ay ang proseso ng quantization ay nangyayari sa isang GPU. Kapag na-initialize na namin ang CUDA at torch sa overmind server, walang madaling paraan upang i-uninitialize ito, na maaaring magdulot ng mga problema para sa mga tunay na workload (pangunahing mas kaunting magagamit na VRAM). Samakatuwid, binago namin ang aming server upang mag-spawn ng isang subprocess, i-load ito sa shared memory, at i-terminate. Ito ay nagkataon na nagpapabuti sa katatagan ng overmind server.

Ang mga quantized parameters ay mga espesyal na subclass na ibinigay ng bitsandbytes. Hindi sila dinisenyo na may 'picklability' sa isip, kaya kailangan naming gawin ito mismo.

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)

Ang mga quantized models sa pamamagitan ng bitsandbytes ay may kasamang mga hook at monkey-patches na hindi nagpi-pickle, kailangan nating alisin ang mga ito:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Alisin ang warning monkeypatches
model.__dict__.pop('cuda', None)

Nakaranas din kami ng mga isyu kung saan ang mga function ay nested sa loob ng iba pang mga function (sa halip na nasa top level), na nagiging sanhi ng hindi pagiging picklable. Sinubukan naming i-workaround ito, ngunit walang swerte. Kinailangan naming baguhin ang aming pickle mula sa stdlib na ibinigay sa dill upang i-pickle ito. Ang dill ay mas makapangyarihan, ngunit ito ay isang purong Python implementation, na mas mabagal kaysa sa bersyon ng standard library. Sa kabutihang palad, ang gastos na ito ay babayaran lamang ng isang beses kapag naglo-load kami ng model sa unang pagkakataon (nakakaapekto lamang sa pag-pickle, hindi sa pag-unpickle).

Suporta para sa stable-fast

Ang stable-fast ay bumubuo ng torch.compile na mga resulta, na hindi maaaring i-pickle. Ngunit sa torch.jit.save, maaari naming i-save ang mga resulta bilang isang zip file. Mukhang hindi ito epektibo, ngunit mas mabuti kaysa sa wala.

Sa torch.jit.save lamang ay hindi sapat upang i-pickle ang mga resulta ng stable-fast. Ang stable-fast ay gumagamit ng isang 'flatten' na proseso upang gawing traceable ang Torch module. Kapag nakatagpo ito ng isang bagay na hindi nito nakikilala (halimbawa, ang dataclass's class), hindi nito ise-serialize ito, ngunit mag-iingat lamang ng reference sa aktwal na klase. Na-patch namin ang kaugnay na logic upang aktwal na mag-imbak ng isang pickled class sa loob ng 'flatten'ed stream.

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

May dalawang karagdagang trick dito:

1. I-repack namin ang ZIP file gamit ang ZIP_STORED, kaya hindi na namin kailangang i-decompress ang ZIP file para sa bawat susunod na load.
2. Ang torch.jit.load interface ay nagdudulot din ng isyu sa memory copy, kaya nagsulat kami ng simpleng wrapper para i-load ito sa pamamagitan ng Python buffer protocol, katulad ng 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);  // No copy!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

Ang vae=vae pattern

Ang aming codebase ay may ganito, sinusubukan nitong i-load ang isang model na may naunang na-load na model bilang argumento nito:

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

pipeline.to('cuda')

Gaya ng nabanggit namin kanina, ang mga argumento ng function ay inaasahang simple, madaling ma-pickle na mga object, ngunit ang pattern na ito ay sumisira sa assumption na iyon. Upang harapin ito, nagdagdag kami ng espesyal na lohika: bawat naka-cache na resulta ay nakakakuha ng ID na nakakabit. Kung ang object na iyon ay ginagamit bilang argumento sa isa pang tawag, pinapalitan ito ng client ng ID nito, at ang server ay maaaring mabawi ang aktwal na object batay sa ID.

Ang nagresultang pipeline model ay maglalaman ng reference sa vae. Para sa pagiging simple, i-pickle lang namin ito nang direkta dito. Gayunpaman, kapag inilipat ang aktwal na UntypedStorage sa shared memory, dinideduplicate namin ang anumang paulit-ulit na data.

Maaaring ginamit namin ang persistent_id na mekanismo ng pickle, ngunit hindi ko sinubukan ang rutang ito. Medyo sayang iyon.

Benchmarking

At ngayon para sa bahagi na gustong makita ng lahat.

Ginagamit namin ang VAE pattern script ng huling seksyon para gawin ang aming pagsusulit.

Tulad ng nakikita mo, ang initial load gamit ang overmind ay tumatagal ng 24.2 segundo, na mas matagal kumpara sa pag-load nang wala ito. Gayunpaman, sa mga susunod na pag-load, ang tanging gastos na natitira ay ang .to('cuda').

Kapag pinagsama-sama ang mga laki ng lahat ng serialized na mga file ng modelo, tinatayang gagamit ang buong pipeline ng humigit-kumulang 5808 megabytes ng memorya. Ang isang mabilis na benchmark ay nagbibigay ng katulad na resulta.

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)

Nasubukan sa Intel i9-11900K + GeForce RTX 4090.

Hindi Inaasahang Mga Epekto (Positibo!)

Ang aming pangunahing motibasyon para sa pagbuo ng overmind ay upang paganahin ang mabilis na pagpapalit ng mga timbang ng modelo sa panahon ng inference. Habang natupad nito ang layunin nito, natuklasan namin ang ilang karagdagang mga benepisyo sa daan.

Nag-deploy kami ng maraming instance ng aming aplikasyon, isa para sa bawat GPU. Kaya, magkakaroon ng 8 proseso bawat node. Matapos naming i-deploy ang overmind, ang paggamit ng system memory ay nabawasan nang malaki. Hindi kami nagkukulang sa system memory, ngunit kung nagkataon, ito ay magiging isang malaking tagumpay.

Kalaunan, natagpuan namin itong isang mahusay na tulong sa aming mga developer ng algorithm at pipeline. Para sa bawat modify-verify loop, makakatipid kami ng 10 hanggang 20 segundo ng oras ng pag-load, na maaaring umabot sa malaking bilang. Higit sa lahat, ang mga segundong natipid ay makakatulong sa mga developer na manatili sa kanilang momentum.

Github

Ibinubukas namin ito sa Github, magiging masaya kami kung ito ay makakatulong.