Overmind: Reduzierung des ML-Modell-Ladevorgangs von 15s auf 0,2s

TL;DR Das Laden von ML-Modellen ist langsam, selbst mit einem aufgewärmten Linux-Seitencache. Daher haben wir eine Bibliothek entwickelt, um es schneller zu machen. Es gibt einige interessante technische Details, die wir teilen möchten, also haben wir diesen Blog geschrieben. Die Bibliothek hatte auch einen unerwarteten Einfluss, der am Ende besprochen wird.

Begründung

Alles begann vor 2 Jahren, als wir unseren ersten Versuch des Lowpoly-Generierungsmodus veröffentlichten. Der Lowpoly-Modus lief nicht gut, er liefert aus heutiger Sicht schlechte Ergebnisse, aber wir haben viel dafür bezahlt – eine dedizierte GPU verarbeitet nur einstellige Aufgaben pro Tag. Sie hat feinabgestimmte Gewichte, groß genug, um alle anderen Modellgewichte aus dem VRAM zu verdrängen. Schlimmer noch, wir haben vielleicht 3 solcher Modelle (ich kann mich nicht an die genaue Anzahl erinnern), sie bildeten einen bedeutenden Teil unserer Inferenz-Infrastruktur und hatten ein ziemlich unnachgiebiges Effizienzverhältnis. Und nein, wir können die Modelle nicht einfach naiv just-in-time laden, es kostet 30 Sekunden, mehr als die eigentliche Verarbeitungszeit.

Damals hatten wir keine dedizierten Pipeline-Ingenieure, unsere Algorithmus-Entwickler versuchten ihr Bestes, um dies zu umgehen. Tage später war unser Code mit this.to('cpu') und that.to('cuda') übersät. Dieser Ansatz funktioniert eine Weile, unterbricht jedoch von Zeit zu Zeit den Arbeitsfluss unserer Algorithmus-Entwickler. Was wäre, wenn die Dinge automatisch passieren könnten? Es ist Python, Dinge passieren automatisch in Python.

Wie definieren Sie 'automatisch'?

Versetzen wir uns in die Rolle eines Algorithmus-Entwicklers. Die Dinge sind ziemlich klar: Ich möchte mich nicht um die Laufzeitleistung außerhalb meines Kernalgorithmus kümmern, es sei denn, ich muss es unbedingt. Ich würde lieber nichts über das Ein- und Auswechseln von Modellen wissen.

Natürlich können wir das nicht erreichen, aber wir können versuchen, die Eingriffe, die wir in den Algorithmuscode einführen müssen, zu minimieren. Das erinnert mich an das Monkey-Patching der gevent-Bibliothek, sie patcht (hauptsächlich) die socket-Bibliothek und ersetzt sie durch gevent.socket, das zu anderen Greenlets wechseln kann, wenn IO blockieren würde, ähnlich wie ein Goroutine (tatsächlich ist gevent älter als Golang!).

Da wir zu der Zeit nur HuggingFace-Bibliotheken (transformers, diffusers) zum Laden von Modellen verwendeten, wurde das Ziel klar: Wir führen nur einen Monkey-Patch-Aufruf ein, und der Rest des Codes sollte unverändert bleiben, XXXPipeline.from_pretrained(...) sollte viel schneller sein.

Einige Fakten, offensichtliche Entscheidungen und Annahmen

Overmind ist eine Caching-Bibliothek, die Ergebnisse von Modellladeaufrufen im Systemspeicher zwischenspeichert und später schnell rekonstruiert.

Wir überspringen die Diskussion darüber, wie das Monkey-Patching implementiert ist, das ist ein nicht so interessantes Detail. Alles, was wir wissen müssen, ist, dass es alle XXXPipeline.from_pretrained(...)-Aufrufe auf overmind.api.load(XXXPipeline.from_pretrained, ...) umleitet.

Wir verwenden pickle, um unser Cache-Ergebnis zu serialisieren, da... wir keine Wahl haben, und torch.save selbst verwendet pickle, es wäre seltsam, es nicht zu verwenden.

Wir verwenden eine Client/Server-Architektur, da wir unseren Cache nicht ungültig machen wollen, wenn der Prozess beendet wird. Es gibt viele Unterprozessaufrufe, die davon profitieren könnten.

Wir nehmen an, dass XXXPipeline.from_pretrained-Parameter einfache hashbare Dinge sind (str und ähnliche Dinge) und andere Modelle von overmind geladen werden (später erklärt).

Der Name overmind ist, wie Sie vielleicht vermutet haben, aus Starcraft entlehnt.

Rekonstruiere es schnell!

Wir können nicht naiv das Ergebnis von pickle.loads im Speicher speichern und es einen Tag nennen. Schließlich hat der Linux-Seitencache in einem aufgewärmten Szenario seine Aufgabe erfüllt, Modelle auf der Festplatte zu cachen, und wir können immer noch eine Ladezeit im Bereich von mehreren zehn Sekunden sehen.

Die Ineffizienz kommt vom Speicherkopieren. In Python würde selbst das Erstellen von Millionen von Objekten nicht mehr als mehrere hundert Millisekunden kosten. Für eine Speicherkopie von 10 GiB würde es jedoch eine halbe Sekunde kosten. Wir müssen das Speicherkopieren so weit wie möglich vermeiden.

Glücklicherweise sind die meisten der großen Speicherblöcke Torch-Tensoren, wir können sicher nur sie adressieren und den Rest ignorieren.

Tatsächlich habe ich das Wissen über die interne Struktur eines Torch-Tensors im Reduktionscode erhalten, während ich den Mechanismus der Tensorfreigabe erforschte:

# Kopiert von torch.multiprocessing.reductions, der größte Teil des Codes wurde entfernt
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))

Ganz einfach: Ein Tensor besteht aus seinem Typ, seinen Metadaten und seinem zugrunde liegenden Speicher. Hier ist storage vom Typ TypedStorage, aber eigentlich ist TypedStorage nur ein einfacher Wrapper für UntypedStorage. UntypedStorage ist die Klasse, die tatsächlich alle Tensor-Daten hält.

Unsere Aufgabe wird jetzt spezifischer: Wie vermeiden wir das Kopieren von UntypedStorage? Können wir diesen Tensor-Speicher selbst verwalten und UntypedStorages konstruieren, indem wir auf den von uns verwalteten Speicher verweisen?

Die Antwort ist ja!

Beim Durchsehen des C++-Codes, wo UntypedStorage konstruiert wird, können wir leicht einen Codeausschnitt wie diesen finden:

// Kopiert von torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...nicht relevanter Code weggelassen...

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

Nicht nur können wir einen Zeiger verwenden, sondern die at::DataPtr-Klasse kann auch die Zerstörung handhaben, was das Lebenszyklus-Management viel einfacher macht.

Auf der Python-Seite wird ein Zeiger auf einen Speicherbereich durch ein memoryview-Objekt dargestellt, diese Objekte unterstützen das Pufferprotokoll. Wir können ein memoryview-Objekt aus vielen Dingen erhalten, bytes und mmap sind die 2 Hauptdinge, die es unterstützen, und sie sind auch das, worauf wir uns konzentrieren.

Schließlich wissen wir, was wir tun sollten: eine Funktion erstellen, die ein memoryview-Objekt akzeptiert und es ohne Kopieren in ein UntypedStorage umwandelt. Mit der Fähigkeit, UntypedStorage aus memoryview zu rekonstruieren, müssen die tatsächlichen Tensor-Daten nicht im Pickle-Stream sein, was die Datenmenge, die wir herumkopieren müssen, stark reduziert.

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

Das ist der Kernbaustein von overmind.

Teilen der Tensors!

Hinweis: Es gibt bereits einen Mechanismus zum Teilen von Tensors in PyTorch, aber er passt nicht zu unseren Bedürfnissen. Mehr dazu später.

Zuerst, Speicher zwischen Client und Server teilen

Wenn wir 'teilen' und 'Speicher' zusammen sehen, haben wir alle den Drang, shmget und seine Freunde zu verwenden. Es ist "entworfen", um als Mechanismus zur Speicherfreigabe verwendet zu werden, richtig? Aber es hat 2 große Schwächen:

• POSIX shm ist eine knappe Ressource, was Sie verwenden können, wird durch die Konfiguration des Systems durch den Systemadministrator bestimmt. Ein extremes, aber allgegenwärtiges Beispiel sind Docker-Container, standardmäßig haben Sie nur 64MiB POSIX shm nutzbar.
• POSIX shm lebt länger als Ihr Prozess, Sie müssen Ihr eigenes Management durchführen. Wenn der Managementprozess gewaltsam beendet wird oder nicht sorgfältig damit umgeht, könnte das shm-Objekt unbegrenzt im System verbleiben.

Wenn Sie genau hinschauen, ist Linux voll von interessanten Systemaufrufen. memfd_create ist einer, der uns interessiert: Er gibt Ihnen einen fd, der eine Zuweisung von anonymem Speicher darstellt. Sie können alle Arten von Dateioperationen darauf ausführen: lesen, schreiben und natürlich mmap. Wenn wir den fd teilen können, können wir den Speicher teilen.

Das Teilen eines fd hat eine 'standardisierte', aber geheimnisvolle Methode: sendmsg mit SCM_RIGHTS. Wir können Bibliotheken nutzen, um uns die entmutigenden Details des sendmsg-Prozesses zu ersparen, aber wir müssen immer noch unsere Koordination zwischen Server- und Client-Prozessen durchführen. Wir haben uns entschieden, hier einen Trick zu verwenden: Öffnen Sie einfach /proc/{pidof(server)}/fd/{memfd} auf der Client-Seite, während Sie den fd auf der overmind-Server-Seite niemals schließen. Die einzige benötigte Kommunikation ist ein (pid, fd)-Tupel. In unserem Fall funktioniert es perfekt.

Die obigen Worte lassen sich auf diese Zeilen reduzieren:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Wird auf der Serverseite aufgerufen
        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):
        # Wird auf der Client-Seite aufgerufen
        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):
        # Wird auf beiden Seiten aufgerufen
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integration mit Pickling

Wie wir zuvor besprochen haben, müssen wir den Pickling-Prozess von UntypedStorage modifizieren. Ähnlich wie in torch.multiprocessing.reductions implementiert, definieren wir unsere benutzerdefinierten Reduktionsfunktionen für pickle:

# Hoarder und borrower sind Wrapper für SharedMemory oben, enthalten
# langweilige Dinge wie Speicherarena usw.
def _reduce_storage(storage):
    # Wird vom Server aufgerufen
    device = storage.device
    storage = storage.cpu()

    # Inhalt im Shared Memory speichern
    # Der `frag` enthält die vollständigen Informationen, die benötigt werden, um den Inhalt zu lokalisieren.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Wird vom Client aufgerufen
    mv = borrower.borrow(frag)  # Erhalte eine memoryview aus dem 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)

Jetzt werden einfache OvermindPickler.dumps und OvermindPickler.loads Shared Memory nutzen, um die Geschwindigkeit zu erhöhen. Sie können hier aufhören zu lesen, wenn Sie bereits genug haben. Der Rest sind Details.

Teufel im Detail

Warum nicht PyTorchs internes Tensor-Sharing?

Mit 'internem Tensor-Sharing' meine ich torch.multiprocessing.reductions.

1. Auf hoher Ebene sind PyTorchs Methoden für das 'Übergeben von Tensoren an Subprozesse' konzipiert, was ähnlich erscheint, aber subtile Unterschiede aufweist.
2. PyTorch verwendet POSIX shm, um Speicher zu teilen, was den zuvor erwähnten Einschränkungen unterliegt.
3. Für jeden Tensor (oder UntypedStorage) allokiert PyTorch ein dediziertes POSIX shm-Objekt, selbst wenn es nur 4 Bytes enthält. Jedes Objekt verbraucht einen fd.
4. PyTorch deallokiert das POSIX shm, sobald sie entpackt werden, was es für unsere Bedürfnisse ungeeignet macht. Wir müssen denselben Pickle-Stream mehrfach deserialisieren.
5. Es gibt viel CUDA-bezogene Sharing-Logik, die für unseren Anwendungsfall reiner Lärm und Ärger ist.

Warum sagen Sie 'Tensor-Daten werden mehrfach kopiert'?

Für ein typisches On-Disk torch.load:

• Die On-Disk torch.save-Datei wird in den Speicher gelesen.
• Holen Sie sich die tatsächlichen torch.UntypedStorage-Daten als bytes durch Extraktion der Zip-Datei (torch.save erzeugt eine Zip-Datei).
• C++-Code kopiert die Daten in seinen eigenen verwalteten Speicher im torch.UntypedStorage-Konstruktor.

Für ein naives pickle.dumps und später pickle.loads:

• Der erzeugte Pickle-Stream bettet intern einen anderen Pickle-Stream ein, pickle.loads kopiert den inneren Stream in ein neues bytes.
• torch.UntypedStorage-Daten werden im inneren Pickle-Stream eingebettet, eine weitere Kopie erfolgt beim Erstellen von torch.UntypedStorage.
• C++-Code kopiert die Daten in seinen eigenen verwalteten Speicher im torch.UntypedStorage-Konstruktor.

diffusers haben ein dynamisches Modul

Modell-Repos können Python-Dateien enthalten, die zur Laufzeit in einen diffusers_modules-Namespace importiert werden. Der Client hat diese nicht in sys.path, was das Unpickling unterbricht. Glücklicherweise schreibt diffusers diese dynamischen Python-Dateien auf die Festplatte, sodass wir das Modul einfach importieren und den Tag genießen können.

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)

Unterstützung für bitsandbytes

Das Nervigste an der Unterstützung von bitsandbytes ist, dass der Quantisierungsprozess auf einer GPU stattfindet. Sobald wir CUDA und Torch im overmind-Server initialisiert haben, gibt es keinen einfachen Weg, es zu deinitialisieren, was Probleme für reale Arbeitslasten verursachen kann (hauptsächlich weniger nutzbarer VRAM). Daher haben wir unseren Server so modifiziert, dass er einen Unterprozess startet, ihn in den gemeinsamen Speicher lädt und beendet. Dies verbessert die Stabilität des overmind-Servers.

Die quantisierten Parameter sind spezielle Unterklassen, die von bitsandbytes bereitgestellt werden. Sie wurden nicht mit der Absicht der 'Picklability' entworfen, daher müssen wir das selbst erledigen.

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)

Quantisierte Modelle über bitsandbytes kommen mit Hooks und Monkey-Patches, die nicht pickelbar sind, wir müssen sie entfernen:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Entfernen von Warnungs-Monkeypatches
model.__dict__.pop('cuda', None)

Wir sind auch auf Probleme gestoßen, bei denen Funktionen innerhalb anderer Funktionen verschachtelt sind (anstatt auf oberster Ebene zu sein), was sie nicht pickelbar macht. Wir haben versucht, dies zu umgehen, aber ohne Erfolg. Wir mussten unser Pickle von dem von der Standardbibliothek bereitgestellten auf dill umstellen, um dies zu pickeln. dill ist viel leistungsfähiger, aber es ist eine reine Python-Implementierung, die viel langsamer ist als die Version der Standardbibliothek. Glücklicherweise wird dieser Aufwand nur einmal bezahlt, wenn wir das Modell zum ersten Mal laden (betrifft nur das Pickeln, nicht das Unpickeln).

Unterstützung für stable-fast

stable-fast erzeugt torch.compile-Ergebnisse, die nicht pickelbar sind. Aber mit torch.jit.save könnten wir die Ergebnisse als Zip-Datei speichern. Das klingt ineffizient, aber es ist besser als nichts.

Mit nur torch.jit.save ist es nicht ausreichend, stable-fast-Ergebnisse zu pickeln. stable-fast verwendet einen 'flatten'-Prozess, um das Torch-Modul nachvollziehbar zu machen. Wenn es auf etwas stößt, das es nicht erkennt (zum Beispiel die Klasse des dataclass), wird es nicht serialisiert, sondern nur eine Referenz auf die tatsächliche Klasse behalten. Wir haben die relevante Logik gepatcht, um tatsächlich eine gepickelte Klasse innerhalb des 'flatten'-Streams zu speichern.

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

Es gibt hier zwei weitere Tricks:

1. Wir packen die ZIP-Datei mit ZIP_STORED neu, sodass wir die ZIP-Datei nicht bei jedem nachfolgenden Laden dekomprimieren müssen.
2. Die torch.jit.load-Schnittstelle verursacht ebenfalls das Problem des Speicherkopierens, daher haben wir einen einfachen Wrapper geschrieben, um es über das Python-Buffer-Protokoll zu laden, genau wie 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);  // Kein Kopieren!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

Das vae=vae-Muster

Unser Codebase enthält so etwas, es versucht, ein Modell mit einem zuvor geladenen Modell als Argument zu laden:

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],  # und Hier!
    torch_dtype=torch.float16,
    safety_checker=None,
)

pipeline.to('cuda')

Wie wir bereits erwähnt haben, wird angenommen, dass die Funktionsargumente einfache, leicht pickelbare Objekte sind, aber dieses Muster bricht diese Annahme. Um dies zu handhaben, haben wir spezielle Logik hinzugefügt: Jedes zwischengespeicherte Ergebnis erhält eine ID angehängt. Wenn dieses Objekt als Argument in einem anderen Aufruf verwendet wird, ersetzt der Client es durch seine ID, und der Server kann dann das tatsächliche Objekt basierend auf der ID wiederherstellen.

Das resultierende pipeline-Modell wird eine Referenz zu vae enthalten. Der Einfachheit halber pickeln wir es hier direkt. Wenn wir jedoch das tatsächliche UntypedStorage in den gemeinsamen Speicher verschieben, deduplizieren wir alle wiederholten Daten.

Wir hätten das persistent_id-Mechanismus von Pickle verwenden können, aber ich habe diesen Weg nicht ausprobiert. Das ist ein bisschen schade.

Benchmarking

Und nun zu dem Teil, den jeder gerne sieht.

Wir verwenden das VAE-Muster-Skript des letzten Abschnitts, um unseren Test durchzuführen.

Wie Sie sehen können, dauert das initiale Laden mit overmind 24,2 Sekunden, was deutlich länger ist im Vergleich zum Laden ohne es. Allerdings ist bei den nachfolgenden Ladevorgängen nur noch die .to('cuda')-Kosten vorhanden.

Addiert man die Größen aller serialisierten Modell-Dateien, wird geschätzt, dass die gesamte Pipeline etwa 5808 Megabyte Speicher verwendet. Ein kurzer Benchmark liefert ein ähnliches Ergebnis.

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)

Getestet auf Intel i9-11900K + GeForce RTX 4090.

Unerwartete Nebeneffekte (Positiv!)

Unsere Hauptmotivation für den Bau von overmind war es, einen schnellen Wechsel der Modellgewichte während der Inferenz zu ermöglichen. Während es seinen Zweck erfüllte, entdeckten wir auf dem Weg mehrere zusätzliche Vorteile.

Wir setzen mehrere Instanzen unserer Anwendung ein, eine für jede GPU. Somit gibt es 8 Prozesse pro Knoten. Nachdem wir overmind implementiert hatten, wurde der Systemspeicherverbrauch drastisch reduziert. Wir hatten zwar keinen Mangel an Systemspeicher, aber wenn wir einen gehabt hätten, wäre dies ein großer Gewinn gewesen.

Später stellten wir fest, dass es einen großen Schub für unsere Algorithmus- und Pipeline-Entwickler darstellte. Für jede Änderungs-Überprüfungs-Schleife konnten wir 10 bis 20 Sekunden Ladezeit sparen, was sich zu einer großen Zahl summieren könnte. Noch wichtiger ist, dass die gesparten Sekunden die Entwickler im Fluss halten konnten.

Github

Wir machen es auf Github als Open Source verfügbar, wir würden uns freuen, wenn es hilft.