Overmind: Skrócenie czasu ładowania modelu ML z 15s do 0.2s

TL;DR Ładowanie modelu ML jest wolne, nawet przy rozgrzanej pamięci podręcznej stron Linuxa. Dlatego stworzyliśmy bibliotekę, aby to przyspieszyć. Istnieją pewne interesujące szczegóły techniczne, którymi chcemy się podzielić, więc napisaliśmy tego bloga. Biblioteka miała także niespodziewany wpływ, o którym mówimy na końcu.

Uzasadnienie

Wszystko zaczęło się 2 lata temu, kiedy wprowadziliśmy naszą pierwszą próbę trybu generowania lowpoly. Tryb lowpoly nie poszedł dobrze, z dzisiejszej perspektywy generuje słabe wyniki, ale dużo za to zapłaciliśmy -- dedykowany GPU przetwarzał tylko jednocyfrową liczbę zadań dziennie. Ma dostrojone wagi, wystarczająco duże, aby wypchnąć wszystkie inne wagi modelu z VRAM. Co gorsza, mamy może 3 takie modele (nie pamiętam dokładnej liczby), stanowiły one znaczną część naszej infrastruktury inferencyjnej, tworząc dość nieprzebaczalny współczynnik efektywności. I nie, nie możemy naiwnie ładować modeli na żądanie, to kosztuje 30 sekund, więcej niż faktyczny czas przetwarzania.

Nie mieliśmy wtedy dedykowanych inżynierów ds. pipeline, nasi deweloperzy algorytmów starali się jak mogli, aby to obejść. Kilka dni później, nasza baza kodu była zaśmiecona this.to('cpu') i that.to('cuda'). To podejście działało przez jakiś czas, ale od czasu do czasu przerywało pracę naszych deweloperów algorytmów. Co jeśli rzeczy mogą dziać się automagicznie? To Python, rzeczy dzieją się automagicznie w Pythonie.

Jak definiujesz 'automagicznie'?

Przejdźmy do roli dewelopera algorytmów. Rzeczy są dość jasne: nie chcę się martwić o wydajność czasu wykonania poza moim głównym algorytmem, chyba że naprawdę muszę. Wolałbym nie wiedzieć nic o zamianie modelu.

Oczywiście nie możemy tego osiągnąć, ale możemy spróbować zminimalizować ingerencję, którą musimy wprowadzić do kodu algorytmu. Przypomina mi to monkey-patching biblioteki gevent, która łata (głównie) bibliotekę socket, zastępując ją gevent.socket, która może przełączać się na inne greenlety, gdy IO blokowałoby, podobnie jak goroutine (w rzeczywistości gevent jest starszy niż Golang!).

Ponieważ używaliśmy tylko bibliotek HuggingFace (transformers, diffusers) do ładowania modeli w tamtym czasie, cel stał się jasny: Wprowadzamy tylko wywołanie monkey-patch, a reszta kodu powinna pozostać niezmieniona, XXXPipeline.from_pretrained(...) powinno być znacznie szybsze.

Kilka Faktów, Oczywistych Decyzji i Założeń

Overmind to biblioteka buforująca, która buforuje wyniki wywołań ładowania modelu do pamięci systemowej i później szybko je rekonstruuje.

Pomijamy dyskusję o tym, jak implementowany jest monkey-patching, to nie jest zbyt interesujący szczegół. Wszystko, co musimy wiedzieć, to to, że przekierowuje wszystkie wywołania XXXPipeline.from_pretrained(...) do overmind.api.load(XXXPipeline.from_pretrained, ...).

Używamy pickle do serializacji naszego wyniku buforowania, ponieważ... nie mamy wyboru, a torch.save sam używa pickle, byłoby dziwne go nie używać.

Używamy architektury klient/serwer, ponieważ nie chcemy unieważniać naszego bufora, gdy proces się kończy. Wiele wywołań podprocesów mogłoby na tym skorzystać.

Zakładamy, że parametry XXXPipeline.from_pretrained to proste hashowalne rzeczy (str i podobne) oraz inne modele ładowane przez overmind (wyjaśnione później).

Nazwa overmind jest zapożyczona ze Starcrafta, jak można się domyślić.

Szybko to zrekonstruuj!

Nie możemy naiwnie zapisać wyniku pickle.loads w pamięci i uznać to za zakończone. W końcu w rozgrzanym scenariuszu, pamięć podręczna stron Linuxa wykonuje swoją pracę buforując modele na dysku i nadal możemy zobaczyć czas ładowania mierzony w dziesiątkach sekund.

Nieskuteczność wynika z kopiowania pamięci. W Pythonie, nawet tworzenie milionów obiektów kosztowałoby nie więcej niż kilkaset ms. Jednak dla kopiowania pamięci o wielkości 10GiB, kosztowałoby to pół sekundy. Musimy unikać kopiowania pamięci tak bardzo, jak to możliwe.

Na szczęście, większość dużych fragmentów pamięci to tensory Torch, możemy bezpiecznie adresować tylko je i ignorować resztę.

Właściwie, zdobyłem wiedzę o wewnętrznej strukturze tensora Torch w kodzie redukcji podczas badania mechanizmu współdzielenia tensorów:

# Skopiowane z torch.multiprocessing.reductions, większość kodu została usunięta
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))

Całkiem proste: tensor to jego typ, metadane i jego podstawowa pamięć. Tutaj storage jest typu TypedStorage, ale w rzeczywistości TypedStorage to tylko prosty wrapper do UntypedStorage. UntypedStorage to klasa, która faktycznie przechowuje wszystkie dane tensora.

Nasze zadanie staje się teraz bardziej szczegółowe: Jak uniknąć kopiowania UntypedStorage? Czy możemy sami zarządzać tą pamięcią tensora i konstruować UntypedStorage wskazując na pamięć, którą zarządzamy?

Odpowiedź brzmi tak!

Przeglądając kod C++ tam, gdzie konstruowany jest UntypedStorage, łatwo możemy znaleźć taki fragment kodu:

// Skopiowane z torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...pomijanie niepowiązanego kodu...

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

Nie tylko możemy użyć wskaźnika, ale klasa at::DataPtr może również obsługiwać destrukcję, co znacznie upraszcza zarządzanie czasem życia.

Po stronie Pythona, wskaźnik do regionu pamięci jest reprezentowany przez obiekt memoryview, te obiekty obsługują protokół bufora. Możemy uzyskać obiekt memoryview z wielu rzeczy, bytes i mmap to 2 główne rzeczy, które to wspierają, i to właśnie one nas interesują.

W końcu wiemy, co powinniśmy zrobić: stworzyć funkcję, która akceptuje obiekt memoryview i zamienia go w UntypedStorage bez kopiowania. Dzięki zdolności do rekonstrukcji UntypedStorage z memoryview, rzeczywiste dane tensora nie muszą być w strumieniu pickle, co znacznie zmniejsza rozmiar danych, które musimy kopiować.

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 jest podstawowy element konstrukcyjny overmind.

Udostępnianie tensorów!

Uwaga: W PyTorch istnieje już mechanizm udostępniania tensorów, ale nie spełnia on naszych potrzeb. Więcej na ten temat później.

Najpierw, udostępnianie pamięci między klientem a serwerem

Kiedy widzimy, że 'udostępnianie' i 'pamięć' pojawiają się razem, wszyscy mamy ochotę użyć shmget i jego przyjaciół. Jest to "zaprojektowane" jako mechanizm udostępniania pamięci, prawda? Ale ma 2 główne wady:

• POSIX shm to zasób ograniczony, to, co możesz użyć, zależy od tego, jak administrator systemu skonfiguruje system. Ekstremalnym, ale powszechnym przykładem są kontenery Docker, domyślnie masz tylko 64MiB POSIX shm do wykorzystania.
• POSIX shm żyje dłużej niż twój proces, musisz zarządzać nim samodzielnie. Jeśli proces zarządzający zostanie siłą zakończony lub nie obsłuży go ostrożnie, obiekt shm może pozostać w systemie na czas nieokreślony.

Jeśli przyjrzysz się uważnie, Linux jest pełen interesujących wywołań systemowych. memfd_create to jedno z nich, które nas interesuje: Daje ci deskryptor pliku (fd), który reprezentuje alokację anonimowej pamięci. Możesz wykonywać na nim wszelkiego rodzaju operacje plikowe: czytać, pisać i, oczywiście, mapować pamięć (mmap). Jeśli możemy udostępnić fd, możemy udostępnić pamięć.

Udostępnianie fd ma 'standardowy', ale tajemniczy sposób: sendmsg z SCM_RIGHTS. Możemy wykorzystać biblioteki, aby pomogły nam ukryć zniechęcające szczegóły procesu sendmsg, ale nadal musimy koordynować procesy serwera i klienta. Zdecydowaliśmy się na pewien trik: Po prostu otwieramy /proc/{pidof(server)}/fd/{memfd} po stronie klienta, jednocześnie nigdy nie zamykając fd po stronie serwera overmind. Jedyną potrzebną komunikacją jest krotka (pid, fd). W naszym przypadku działa to doskonale.

Powyższe słowa sprowadzają się do tych linii:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Wywołane po stronie serwera
        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):
        # Wywołane po stronie 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):
        # Wywołane po obu stronach
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integracja z picklingiem

Jak omówiliśmy wcześniej, musimy zmodyfikować proces picklingu UntypedStorage. Podobnie jak w torch.multiprocessing.reductions, definiujemy nasze własne funkcje redukcji dla pickle:

# Hoarder i borrower to opakowania dla SharedMemory powyżej, zawierają
# nudne rzeczy jak arena pamięci, itp.
def _reduce_storage(storage):
    # Wywołane przez serwer
    device = storage.device
    storage = storage.cpu()

    # Przechowaj zawartość w pamięci współdzielonej
    # `frag` zawiera pełne informacje potrzebne do zlokalizowania zawartości.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Wywołane przez klienta
    mv = borrower.borrow(frag)  # Pobierz memoryview z pamięci współdzielonej
    storage = _make_untyped_storage(mv)  # Zero-kopia!
    if device.type == 'cuda':
        return storage.cuda(device.index)
    return storage

class OvermindPickler(dill.Pickler):
    ...

OvermindPickler.register(torch.storage.UntypedStorage, _reduce_storage)

Teraz, proste OvermindPickler.dumps i OvermindPickler.loads będą wykorzystywać pamięć współdzieloną do przyspieszenia. Możesz przestać czytać tutaj, jeśli już masz dość. Reszta to szczegóły.

Diabeł tkwi w szczegółach

Dlaczego nie użyć wewnętrznego udostępniania tensorów PyTorch?

Przez 'wewnętrzne udostępnianie tensorów' rozumiem torch.multiprocessing.reductions.

1. Na wysokim poziomie, metoda PyTorch jest zaprojektowana do 'przekazywania tensora do podprocesu', wydaje się podobna, ale z subtelnymi różnicami.
2. PyTorch używa POSIX shm do udostępniania pamięci, co podlega wcześniej wspomnianym ograniczeniom.
3. Dla każdego tensora (lub UntypedStorage), PyTorch alokuje dedykowany obiekt POSIX shm, nawet jeśli zawiera tylko 4 bajty. Każdy obiekt zużywa fd.
4. PyTorch dealokuje POSIX shm, gdy tylko zostaną one odczytane, co czyni go nieodpowiednim dla naszych potrzeb. Musimy deserializować ten sam strumień pickle wielokrotnie.
5. Istnieje wiele logiki związanej z udostępnianiem CUDA, która jest czystym hałasem i problemem dla naszego przypadku użycia.

Dlaczego mówisz, że 'dane tensora są kopiowane wielokrotnie'?

Dla typowego torch.load z dysku:

• Plik torch.save z dysku jest wczytywany do pamięci.
• Uzyskaj rzeczywiste dane torch.UntypedStorage jako bytes poprzez ekstrakcję pliku Zip (funkcja torch.save generuje plik zip).
• Kod C++ skopiuje dane do własnej zarządzanej pamięci w konstruktorze torch.UntypedStorage.

Dla naiwnego pickle.dumps i późniejszego pickle.loads:

• Generowany strumień pickle wewnętrznie osadza inny strumień pickle, pickle.loads skopiuje wewnętrzny strumień do nowych bytes.
• Dane torch.UntypedStorage są osadzone w wewnętrznym strumieniu pickle, kolejna kopia następuje podczas konstrukcji torch.UntypedStorage.
• Kod C++ skopiuje dane do własnej zarządzanej pamięci w konstruktorze torch.UntypedStorage.

diffusers mają dynamiczny moduł

Repozytoria modeli mogą zawierać pliki Python, które są importowane w czasie rzeczywistym do przestrzeni nazw diffusers_modules. Klient nie ma ich w sys.path, co powoduje problemy z rozpakowywaniem. Na szczęście, diffusers zapisze te dynamiczne pliki Python na dysku, więc możemy po prostu zaimportować moduł i gotowe.

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)

Wsparcie dla bitsandbytes

Najbardziej irytującą rzeczą przy wspieraniu bitsandbytes jest to, że proces kwantyzacji odbywa się na GPU. Po zainicjowaniu CUDA i torch na serwerze overmind, nie ma łatwego sposobu na jego odinicjowanie, co może powodować problemy dla rzeczywistych obciążeń (głównie mniej użytecznego VRAM). Dlatego zmodyfikowaliśmy nasz serwer, aby uruchamiał podproces, ładował go do pamięci współdzielonej i kończył. To poprawia stabilność serwera overmind.

Skwantyzowane parametry są specjalnymi podklasami dostarczanymi przez bitsandbytes. Nie zostały zaprojektowane z myślą o 'picklability', więc musimy to zrobić 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)

Skwantyzowane modele za pomocą bitsandbytes mają haki i monkey-patche, które nie są picklowalne, musimy je usunąć:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Usuń monkey-patche ostrzegawcze
model.__dict__.pop('cuda', None)

Napotkaliśmy również problemy, gdy funkcje są zagnieżdżone w innych funkcjach (zamiast być na najwyższym poziomie), co sprawia, że nie są picklowalne. Próbowaliśmy obejść ten problem, ale bez powodzenia. Musieliśmy przełączyć nasz pickle z dostarczonego przez stdlib na dill, aby to picklować. dill jest znacznie potężniejszy, ale to czysta implementacja w Pythonie, która jest znacznie wolniejsza niż wersja standardowej biblioteki. Na szczęście ten koszt zostanie poniesiony tylko raz, gdy ładujemy model po raz pierwszy (dotyczy tylko picklowania, nie rozpakowywania).

Wsparcie dla stable-fast

stable-fast generuje wyniki torch.compile, które nie mogą być picklowane. Ale z torch.jit.save, możemy zapisać wyniki jako plik zip. To brzmi nieefektywnie, ale lepsze to niż nic.

Samo torch.jit.save nie wystarcza do picklowania wyników stable-fast. stable-fast używa procesu 'flatten', aby uczynić moduł Torch śledzonym. Gdy napotka coś, czego nie rozpoznaje (na przykład klasę dataclass), nie zserializuje tego, ale zachowa tylko odniesienie do rzeczywistej klasy. Zaktualizowaliśmy odpowiednią logikę, aby faktycznie przechowywać picklowaną klasę w strumieniu '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

Są tutaj dwa dodatkowe triki:

1. Przepakowujemy plik ZIP z ZIP_STORED, więc nie musimy dekompresować pliku ZIP przy każdym kolejnym ładowaniu.
2. Interfejs torch.jit.load również powoduje problem z kopiowaniem pamięci, więc napisaliśmy prostą nakładkę, aby ładować go za pomocą protokołu bufora Pythona, podobnie jak 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);  // Bez kopiowania!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

Wzorzec vae=vae

Nasz kod zawiera coś takiego, próbuje załadować model z wcześniej załadowanym modelem jako jego 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,  # Tutaj!
    controlnet=[controlnet_edge, controlnet_depth],  # i Tutaj!
    torch_dtype=torch.float16,
    safety_checker=None,
)

pipeline.to('cuda')

Jak wspomnieliśmy wcześniej, argumenty funkcji są zakładane jako proste, łatwe do serializacji obiekty, ale ten wzorzec łamie to założenie. Aby sobie z tym poradzić, dodaliśmy specjalną logikę: każdy wynik z cache otrzymuje przypisane ID. Jeśli ten obiekt jest używany jako argument w innym wywołaniu, klient zastępuje go jego ID, a serwer może wtedy odzyskać rzeczywisty obiekt na podstawie ID.

Ostateczny model pipeline będzie zawierał odniesienie do vae. Dla uproszczenia, po prostu serializujemy go bezpośrednio tutaj. Jednak przy przenoszeniu rzeczywistego UntypedStorage do pamięci współdzielonej, deduplikujemy wszelkie powtarzające się dane.

Mogliśmy użyć mechanizmu persistent_id z pickle, ale nie próbowałem tej drogi. To trochę szkoda.

Benchmarking

A teraz część, którą wszyscy lubią oglądać.

Używamy skryptu wzorca VAE z ostatniej sekcji do naszego testu.

Jak widać, początkowe ładowanie z overmind zajmuje 24,2 sekundy, co jest znacznie dłużej w porównaniu do ładowania bez niego. Jednak przy kolejnych ładowaniach jedynie koszt .to('cuda') nadal występuje.

Dodając rozmiary wszystkich zserializowanych plików modelu, szacuje się, że cały pipeline używa około 5808 megabajtów pamięci. Szybki benchmark daje podobny wynik.

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)

Testowane na Intel i9-11900K + GeForce RTX 4090.

Nieoczekiwane Efekty Uboczne (Pozytywne!)

Naszą główną motywacją do stworzenia overmind było umożliwienie szybkiego przełączania wag modelu podczas wnioskowania. Chociaż spełnił swoje zadanie, odkryliśmy po drodze kilka dodatkowych zalet.

Wdrażamy wiele instancji naszej aplikacji, po jednej na każdy GPU. W związku z tym na każdym węźle będzie 8 procesów. Po wdrożeniu overmind zużycie pamięci systemowej zostało dramatycznie zmniejszone. Nie cierpieliśmy na brak pamięci systemowej, ale gdybyśmy cierpieli, byłoby to wielkie zwycięstwo.

Później odkryliśmy, że jest to świetny impuls dla naszych programistów algorytmów i pipeline'ów. Dla każdej pętli modyfikacji i weryfikacji mogliśmy zaoszczędzić 10 do 20 sekund czasu ładowania, co mogło się sumować do dużej liczby. Co ważniejsze, zaoszczędzone sekundy mogły utrzymać programistów w rytmie pracy.

Github

Otwieramy kod na Github, będziemy szczęśliwi, jeśli to pomoże.