Overmind: Reduzir o Carregamento do Modelo de ML de 15s para 0.2s

TL;DR O carregamento de modelos de ML é lento, mesmo com a cache de páginas do Linux aquecida. Por isso, construímos uma biblioteca para torná-lo rápido. Existem alguns detalhes técnicos interessantes que queremos partilhar, por isso escrevemos este blog. A biblioteca também teve um impacto inesperado, discutido no final.

Justificação

Tudo começou há 2 anos, quando lançámos a nossa primeira tentativa de modo de geração lowpoly. O modo lowpoly não correu bem, emite resultados pobres da perspetiva de hoje, mas pagámos muito por isso -- uma GPU dedicada processa apenas tarefas de um dígito por dia. Tem pesos ajustados, grandes o suficiente para expulsar todos os outros pesos do modelo da VRAM. Pior, temos talvez 3 desses modelos (não me lembro do número exato), constituíam uma parte significativa da nossa infraestrutura de inferência, criando uma relação de eficiência bastante implacável. E não, não podemos carregar os modelos ingenuamente just-in-time, custa 30s, maior que o tempo de processamento real.

Na altura não tínhamos engenheiros de pipeline dedicados, os nossos desenvolvedores de algoritmos fizeram o melhor para contornar isso. Dias depois, a nossa base de código estava cheia de this.to('cpu') e that.to('cuda'). Esta abordagem funcionou por um tempo, mas interrompia o fluxo dos nossos desenvolvedores de algoritmos de tempos a tempos. E se as coisas pudessem acontecer automaticamente? É Python, as coisas acontecem automaticamente em Python.

Como defines 'automaticamente'?

Vamos assumir o papel de um desenvolvedor de algoritmos. As coisas são bastante claras: não quero preocupar-me com o desempenho em tempo de execução fora do meu algoritmo principal, a menos que realmente tenha de o fazer. Preferiria não saber nada sobre a troca de modelos.

Claro que não podemos alcançar isso, mas podemos tentar minimizar a intrusão que temos de introduzir no código do algoritmo. Isto faz-me lembrar a modificação de macaco da biblioteca gevent, que modifica (principalmente) a biblioteca socket, substituindo-a por gevent.socket que pode alternar para outros greenlets quando o IO bloquearia, muito parecido com uma goroutine (na verdade, gevent é mais antigo que Golang!).

Como estávamos apenas a usar as bibliotecas HuggingFace (transformers, diffusers) para carregar modelos na altura, o objetivo tornou-se claro: Apenas introduzimos uma chamada de modificação de macaco, e o resto do código deve permanecer inalterado, XXXPipeline.from_pretrained(...) deve ser muito mais rápido.

Alguns Factos, Decisões Óbvias e Suposições

Overmind é uma biblioteca de cache, armazena em cache os resultados das chamadas de carregamento de modelos na memória do sistema e depois reconstrói-os rapidamente.

Vamos pular a discussão sobre como a modificação de macaco é implementada, esse é um detalhe não tão interessante. Tudo o que precisamos saber é que redireciona todas as chamadas XXXPipeline.from_pretrained(...) para overmind.api.load(XXXPipeline.from_pretrained, ...).

Usamos pickle para serializar o nosso resultado de cache porque... não temos escolha, e torch.save em si usa pickle, seria estranho não usá-lo.

Usamos uma arquitetura cliente/servidor porque não queremos invalidar o nosso cache quando o processo termina. Existem muitas chamadas de subprocessos que poderiam beneficiar disso.

Assumimos que os parâmetros XXXPipeline.from_pretrained são coisas simples que podem ser hashadas (str e coisas semelhantes) e outros modelos carregados por overmind (explicado mais tarde).

O nome overmind é emprestado do Starcraft, como pode ter adivinhado.

Reconstrua-o rapidamente!

Não podemos simplesmente salvar o resultado de pickle.loads na memória e dar o dia por terminado. Afinal, num cenário aquecido, a cache de páginas do Linux fez o seu trabalho de cache dos modelos em disco e ainda podemos ver um tempo de carregamento medido em dezenas de segundos.

A ineficiência vem da cópia de memória. Em Python, mesmo criar milhões de objetos custaria não mais que várias centenas de ms. No entanto, para uma cópia de memória de 10GiB, custaria meio segundo. Devemos evitar a cópia de memória tanto quanto possível.

Felizmente, a maioria dos grandes blocos de memória são tensores do Torch, podemos endereçá-los com segurança e ignorar o resto.

Na verdade, obtive o conhecimento da estrutura interna de um tensor do Torch no código de redução enquanto pesquisava o mecanismo de partilha de tensores:

# Copiado de torch.multiprocessing.reductions, a maior parte do código foi removida
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))

Bastante simples: um tensor é o seu tipo, os seus metadados e o seu armazenamento subjacente. Aqui, storage é do tipo TypedStorage, mas na verdade TypedStorage é apenas um invólucro simples para UntypedStorage. UntypedStorage é a classe que realmente contém todos os dados do tensor.

A nossa tarefa torna-se agora mais específica: Como evitamos copiar UntypedStorage? Podemos gerir esta memória de tensor por nós mesmos e construir UntypedStorages apontando para a memória que gerimos?

A resposta é sim!

Ao folhear o código C++ de onde UntypedStorage é construído, podemos facilmente encontrar um trecho de código como este:

// Copiado de torch/csrc/Storage.cpp
static PyObject* THPStorage_get(THPStorage* self, PyObject* index) {
    // ...omitindo código não relacionado...

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

Não só podemos usar um ponteiro, mas a classe at::DataPtr também pode lidar com a destruição, tornando a gestão do tempo de vida muito mais simples.

Do lado do Python, um ponteiro para uma região de memória é representado por um objeto memoryview, estes objetos suportam o protocolo de buffer. Podemos obter um objeto memoryview de muitas coisas, bytes e mmap são as 2 principais coisas que o suportam, e são também o que nos interessa.

Finalmente, sabemos o que devemos fazer: criar uma função que aceite um objeto memoryview e o transforme num UntypedStorage sem copiar. Com a capacidade de reconstruir UntypedStorage a partir de memoryview, os dados reais do tensor não precisam estar no fluxo de pickle, reduzindo significativamente o tamanho dos dados que temos de copiar.

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

Esse é o bloco de construção principal do overmind.

Partilhando os tensores!

Nota: Já existe um mecanismo de partilha de tensores no PyTorch, mas não se adequa às nossas necessidades. Mais sobre isso mais tarde.

Primeiro, partilhar memória entre cliente e servidor

Quando vemos 'partilhar' e 'memória' juntos, todos temos uma vontade de usar shmget e seus amigos. Foi "desenhado" para ser usado como um mecanismo de partilha de memória, certo? Mas tem 2 falhas principais:

• O shm POSIX é um recurso escasso, o que podes usar é determinado por como o administrador do sistema configura o sistema. Um exemplo extremo mas ubíquo são os contentores Docker, por padrão tens apenas 64MiB de shm POSIX utilizável.
• O shm POSIX vive mais tempo do que o seu processo, você tem que fazer a sua própria gestão. Se o processo de gestão for forçado a terminar, ou não for manuseado cuidadosamente, o objeto shm pode ficar no sistema indefinidamente.

Se olhar com atenção, o Linux está cheio de chamadas de sistema interessantes. memfd_create é uma que nos interessa: Dá-lhe um fd que representa uma alocação de memória anónima. Pode fazer todo o tipo de operações de ficheiro nele: ler, escrever e, claro, mmap. Se conseguirmos partilhar o fd, podemos partilhar a memória.

Partilhar um fd tem uma forma 'padrão' mas arcana de o fazer: sendmsg com SCM_RIGHTS. Podemos aproveitar bibliotecas para nos ajudar a esconder os detalhes assustadores do processo sendmsg, mas ainda temos que fazer a nossa coordenação entre os processos do servidor e do cliente. Decidimos usar um truque aqui: Basta abrir /proc/{pidof(server)}/fd/{memfd} no lado do cliente, enquanto nunca fechamos o fd no lado do servidor overmind. A única comunicação necessária é um tuplo (pid, fd). Funciona perfeitamente no nosso caso.

As palavras acima resumem-se nestas linhas:

class SharedMemory:
    @classmethod
    def create(cls, shift):
        # Chamado no lado do servidor
        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):
        # Chamado no lado do cliente
        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):
        # Chamado em ambos os lados
        self._mmap = mmap.mmap(self._fd, size)
        self._buf = memoryview(self._mmap)
        return self._buf

Integrar com pickling

Como discutimos antes, precisamos modificar o processo de pickling de UntypedStorage. Similar ao que foi implementado em torch.multiprocessing.reductions, definimos as nossas funções de redução personalizadas para pickle:

# Hoarder e borrower são um wrapper para SharedMemory acima, contém
# coisas aborrecidas como arena de memória, etc.
def _reduce_storage(storage):
    # Chamado pelo servidor
    device = storage.device
    storage = storage.cpu()

    # Armazenar conteúdo na memória partilhada
    # O `frag` contém a informação completa necessária para localizar o conteúdo.
    frag = hoarder.put(storage)

    return (_rebuild_storage_on_client, (frag, device))

def _rebuild_storage_on_client(frag, device):
    # Chamado pelo cliente
    mv = borrower.borrow(frag)  # Obter uma memoryview da memória partilhada
    storage = _make_untyped_storage(mv)  # Zero-cópia!
    if device.type == 'cuda':
        return storage.cuda(device.index)
    return storage

class OvermindPickler(dill.Pickler):
    ...

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

Agora, simples OvermindPickler.dumps e OvermindPickler.loads irão utilizar a memória partilhada para acelerar. Pode parar de ler aqui se já estiver farto. O resto são detalhes.

Diabos nos Detalhes

Porque não o compartilhamento de tensor interno do PyTorch?

Por 'compartilhamento de tensor interno', refiro-me a torch.multiprocessing.reductions.

1. A um nível elevado, o método do PyTorch é desenhado para 'passar tensor para subprocesso', parece o mesmo mas com uma diferença subtil.
2. O PyTorch usa shm POSIX para partilhar memória, sujeito ao limite mencionado anteriormente.
3. Para cada tensor (ou UntypedStorage), o PyTorch aloca um objeto shm POSIX dedicado para ele, mesmo que contenha apenas 4 bytes. Cada objeto consome um fd.
4. O PyTorch desaloca o shm POSIX assim que são desempacotados, tornando-o inadequado para as nossas necessidades. Precisamos desserializar o mesmo fluxo de pickle várias vezes.
5. Há muita lógica de compartilhamento relacionada com CUDA, que são puro ruído e problema para o nosso caso de uso.

Porque diz que 'dados do tensor copiados várias vezes'?

Para um típico torch.load em disco:

• O ficheiro torch.save em disco é lido para a memória.
• Obtenha os dados reais de torch.UntypedStorage como bytes através da extração de ficheiros Zip (o torch.save gera um ficheiro zip).
• O código C++ copiará os dados para a sua própria memória gerida no construtor de torch.UntypedStorage.

Para um pickle.dumps ingênuo e posteriormente pickle.loads:

• O fluxo de pickle gerado embute internamente outro fluxo de pickle, o pickle.loads copiará o fluxo interno para um novo bytes.
• Os dados de torch.UntypedStorage são embutidos no fluxo de pickle interno, outra cópia acontece na construção de torch.UntypedStorage.
• O código C++ copiará os dados para a sua própria memória gerida no construtor de torch.UntypedStorage.

diffusers têm um módulo dinâmico

Os repositórios de modelos podem incluir ficheiros Python que são importados em tempo de execução para um namespace diffusers_modules. O cliente não tem estes ficheiros em sys.path, o que quebra o deserialização. Felizmente, diffusers escreverá estes ficheiros Python dinâmicos no disco, então podemos simplesmente importar o módulo e seguir em frente.

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)

Suporte para bitsandbytes

A coisa mais irritante sobre suportar bitsandbytes é que o processo de quantização acontece numa GPU. Uma vez que inicializámos CUDA e torch no servidor overmind, não há uma maneira fácil de desinicializá-lo, o que pode causar problemas para cargas de trabalho reais (principalmente menos VRAM utilizável). Portanto, modificámos o nosso servidor para gerar um subprocesso, carregá-lo em memória partilhada e terminar. Isto melhora a estabilidade do servidor overmind.

Os parâmetros quantizados são subclasses especiais fornecidas por bitsandbytes. Eles não foram projetados com 'capacidade de serialização' em mente, então temos que fazê-lo nós mesmos.

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)

Modelos quantizados via bitsandbytes vêm com hooks e monkey-patches que não são serializáveis, devemos removê-los:

from accelerate.hooks import remove_hook_from_module
remove_hook_from_module(model, True)
model.__dict__.pop('to', None)  # Remover monkeypatches de aviso
model.__dict__.pop('cuda', None)

Também encontramos problemas onde funções estão aninhadas dentro de outras funções (em vez de estarem no nível superior), o que as torna não serializáveis. Tentámos contornar isso, mas sem sucesso. Tivemos que mudar o nosso pickle do fornecido pela stdlib para dill para serializar isso. dill é muito mais poderoso, mas é uma implementação pura em Python, o que é muito mais lento do que a versão da biblioteca padrão. Felizmente, este custo só será pago uma vez quando estivermos a carregar o modelo pela primeira vez (afeta apenas a serialização, não a deserialização).

Suporte para stable-fast

stable-fast gera resultados de torch.compile, que não podem ser serializados. Mas com torch.jit.save, poderíamos salvar os resultados como um ficheiro zip. Isto parece ineficiente, mas é melhor do que nada.

Com apenas torch.jit.save não é suficiente para serializar os resultados de stable-fast. stable-fast usa um processo de 'flatten' para tornar o módulo Torch rastreável. Ao encontrar algo que não reconhece (por exemplo, a classe de dataclass), não o serializa, mas apenas mantém uma referência à classe real. Corrigimos a lógica relevante para realmente armazenar uma classe serializada dentro do fluxo '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

Existem mais dois truques aqui:

1. Reempacotamos o ficheiro ZIP com ZIP_STORED, para que não tenhamos de descomprimir o ficheiro ZIP em cada carregamento subsequente.
2. A interface torch.jit.load também incorre no problema de cópia de memória, por isso escrevemos um simples wrapper para carregá-lo através do protocolo de buffer do Python, tal como 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);  // Sem cópia!
            return import_ir_module(std::move(cu), in, ...);
        }
    );
}

O padrão vae=vae

A nossa base de código tem algo assim, tenta carregar um modelo com um modelo previamente carregado como seu argumento:

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

pipeline.to('cuda')

Como mencionámos anteriormente, os argumentos das funções são assumidos como objetos simples, facilmente serializáveis, mas este padrão quebra essa suposição. Para lidar com isso, adicionámos uma lógica especial: cada resultado em cache recebe um ID anexado. Se esse objeto for usado como argumento noutra chamada, o cliente substitui-o pelo seu ID, e o servidor pode então recuperar o objeto real com base no ID.

O modelo pipeline resultante conterá uma referência para vae. Para simplificar, apenas o serializamos diretamente aqui. No entanto, ao mover o UntypedStorage real para a memória partilhada, deduplicamos quaisquer dados repetidos.

Poderíamos ter usado o mecanismo persistent_id do pickle, mas não tentei essa rota. É uma pena.

Benchmarking

E agora para a parte que todos adoram ver.

Usamos o script do padrão VAE da última seção para fazer o nosso teste.

Como pode ver, o carregamento inicial com overmind demora 24,2 segundos, o que é significativamente mais longo em comparação com o carregamento sem ele. No entanto, em carregamentos subsequentes, apenas o custo de .to('cuda') ainda está presente.

Somando os tamanhos de todos os ficheiros de modelo serializados, estima-se que todo o pipeline utilize cerca de 5808 megabytes de memória. Um teste rápido dá um resultado semelhante.

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)

Testado em Intel i9-11900K + GeForce RTX 4090.

Efeitos Colaterais Inesperados (Positivos!)

A nossa principal motivação para construir o overmind foi permitir a troca rápida de pesos de modelo durante a inferência. Enquanto cumpria o seu propósito, descobrimos várias vantagens adicionais ao longo do caminho.

Desdobramos múltiplas instâncias da nossa aplicação, uma para cada GPU. Assim, haverá 8 processos por nó. Depois de implementarmos o overmind, o uso de memória do sistema foi reduzido drasticamente. Não estávamos a sofrer de falta de memória do sistema, mas se estivéssemos, isto teria sido uma grande vitória.

Mais tarde, descobrimos que foi um grande impulso para os nossos desenvolvedores de algoritmos e pipelines. Para cada ciclo de modificação-verificação, poderíamos poupar de 10 a 20 segundos de tempo de carregamento, o que poderia somar um grande número. Mais importante ainda, os segundos poupados poderiam manter os desenvolvedores no fluxo.

Github

Estamos a disponibilizá-lo em código aberto no Github, ficaremos felizes se ajudar.