Overmind: Reduzindo o Carregamento de Modelos de ML de 15s para 0.2s

TL;DR O carregamento de modelos de ML é lento, mesmo com o cache de página do Linux aquecido. Então, construímos uma biblioteca para torná-lo rápido. Existem alguns detalhes técnicos interessantes que queremos compartilhar, então escrevemos este blog. A biblioteca também teve um impacto inesperado, discutido no final.

Justificativa

Tudo começou há 2 anos, quando lançamos nossa primeira tentativa do modo de geração lowpoly. O modo lowpoly não foi bem, ele produz resultados ruins do ponto de vista atual, mas pagamos caro por isso -- uma GPU dedicada processa apenas tarefas de um dígito por dia. Ele tem pesos ajustados, grandes o suficiente para expulsar todos os outros pesos do modelo da VRAM. Pior, temos talvez 3 desses modelos (não consigo lembrar o número exato), eles constituíam uma parte significativa da nossa infraestrutura de inferência, tornando a relação de eficiência bastante implacável. E não, não podemos carregar os modelos ingenuamente no momento necessário, isso custa 30s, maior que o tempo de processamento real.

Não tínhamos engenheiros de pipeline dedicados na época, nossos desenvolvedores de algoritmos fizeram o melhor para contornar isso. Dias depois, nosso código estava cheio de this.to('cpu') e that.to('cuda'). Essa abordagem funcionou por um tempo, mas interrompia o fluxo dos nossos desenvolvedores de algoritmos de tempos em tempos. E se as coisas pudessem acontecer automaticamente? É Python, as coisas acontecem automaticamente em Python.

Como você define 'automaticamente'?

Vamos nos colocar no papel de um desenvolvedor de algoritmos. As coisas são bem claras: eu não quero me preocupar com o desempenho em tempo de execução fora do meu algoritmo principal, a menos que eu realmente precise. Eu 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 que introduzir no código do algoritmo. Isso me lembra o monkey-patching da biblioteca gevent, que faz patch (principalmente) na 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 usando apenas as bibliotecas HuggingFace (transformers, diffusers) para carregar modelos na época, o alvo ficou claro: introduzimos apenas uma chamada de monkey-patch, e o resto do código deve permanecer inalterado, XXXPipeline.from_pretrained(...) deve ser muito mais rápido.

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

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

Pulamos a discussão sobre como o monkey-patching é implementado, esse é um detalhe não tão interessante. Tudo o que precisamos saber é que ele redireciona todas as chamadas XXXPipeline.from_pretrained(...) para overmind.api.load(XXXPipeline.from_pretrained, ...).

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

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

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

O nome overmind é emprestado do Starcraft, como você deve ter adivinhado.

Reconstrua rapidamente!

Não podemos salvar ingenuamente o resultado de pickle.loads na memória e dar o dia por encerrado. Afinal, em um cenário aquecido, o cache de página do Linux fez seu trabalho armazenando em cache os 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 criando milhões de objetos custaria não mais que algumas 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, adquiri o conhecimento da estrutura interna de um tensor do Torch no código de redução enquanto pesquisava o mecanismo de compartilhamento 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 é seu tipo, seus metadados e 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.

Nossa tarefa se torna mais específica agora: como evitamos copiar UntypedStorage? Podemos gerenciar essa memória de tensor por nós mesmos e construir UntypedStorages apontando para a memória que gerenciamos?

A resposta é sim!

Passando os olhos pelo código C++ 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 o gerenciamento do tempo de vida muito mais simples.

No lado do Python, um ponteiro para uma região de memória é representado por um objeto memoryview, esses 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 em um UntypedStorage sem copiar. Com a capacidade de reconstruir UntypedStorage a partir de memoryview, os dados reais do tensor não precisam estar no fluxo pickle, reduzindo significativamente o tamanho dos dados que temos que 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.

Compartilhando os tensores!

Nota: Já existe um mecanismo de compartilhamento de tensores no PyTorch, mas ele não atende às nossas necessidades. Mais sobre isso depois.

Primeiro, compartilhando memória entre cliente e servidor

Quando vemos 'compartilhar' e 'memória' juntos, todos temos uma vontade de usar shmget e seus amigos. Ele foi "projetado" para ser usado como um mecanismo de compartilhamento de memória, certo? Mas ele tem 2 grandes falhas:

• O shm POSIX é um recurso escasso, o que você pode usar é determinado por como o administrador do sistema configura o sistema. Um exemplo extremo, mas ubíquo, são os contêineres Docker, por padrão você tem apenas 64MiB de shm POSIX utilizável.
• O shm POSIX vive mais do que seu processo, você precisa fazer sua própria gestão. Se o processo de gerenciamento for forçado a encerrar ou não lidar com isso cuidadosamente, o objeto shm pode ser deixado no sistema indefinidamente.

Se você olhar com atenção, o Linux está cheio de chamadas de sistema interessantes. memfd_create é uma que nos interessa: Ela fornece um fd que representa uma alocação de memória anônima. Você pode realizar todos os tipos de operações de arquivo nele: leitura, escrita e, claro, mmap. Se pudermos compartilhar o fd, podemos compartilhar a memória.

Compartilhar um fd tem uma maneira 'padrão', mas arcana, de fazê-lo: sendmsg com SCM_RIGHTS. Podemos aproveitar bibliotecas para nos ajudar a esconder os detalhes assustadores do processo sendmsg, mas ainda precisamos fazer 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 é uma tupla (pid, fd). Funciona perfeitamente no nosso caso.

As palavras acima se resumem a estas 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 serialização

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

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

    # Armazena conteúdo na memória compartilhada
    # 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)  # Obtém um memoryview da memória compartilhada
    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 utilizarão memória compartilhada para acelerar. Você pode parar de ler aqui se já estiver cansado. O restante são detalhes.

Diabos nos Detalhes

Por que não o compartilhamento de tensor interno do PyTorch?

Por 'compartilhamento de tensor interno', quero dizer torch.multiprocessing.reductions.

1. Em alto nível, o método do PyTorch é projetado para 'passar tensor para subprocesso', parece o mesmo, mas com diferença sutil.
2. O PyTorch usa shm POSIX para compartilhar 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 eles são desserializados, tornando-o inadequado para nossas necessidades. Precisamos desserializar o mesmo fluxo de pickle várias vezes.
5. Há muita lógica de compartilhamento relacionada ao CUDA, que são puro ruído e problema para nosso caso de uso.

Por que você diz 'dados do tensor copiados várias vezes'?

Para um típico torch.load em disco:

• O arquivo torch.save em disco é lido na memória.
• Obtenha os dados reais de torch.UntypedStorage como bytes por meio da extração de arquivo Zip (o torch.save gera um arquivo zip).
• O código C++ copiará os dados para sua própria memória gerenciada no construtor de torch.UntypedStorage.

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

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

diffusers têm um módulo dinâmico

Repositórios de modelos podem incluir arquivos Python que são importados em tempo de execução em um namespace diffusers_modules. O cliente não possui esses arquivos em sys.path, o que quebra o processo de deserialização. Felizmente, diffusers escreverá esses arquivos 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 o suporte a bitsandbytes é que o processo de quantização acontece em uma GPU. Uma vez que inicializamos 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, modificamos nosso servidor para gerar um subprocesso, carregá-lo na memória compartilhada e terminar. Isso 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 fazer isso 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)  # Remove 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. Tentamos contornar isso, mas sem sucesso. Tivemos que trocar nosso pickle do fornecido pela biblioteca padrão 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, esse custo será pago apenas uma vez quando estivermos carregando o modelo pela primeira vez (afeta apenas a serialização, não a desserializaçã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 arquivo zip. Isso parece ineficiente, mas é melhor do que nada.

Somente com 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 um dataclass), ele não o serializa, mas apenas mantém uma referência à classe real. Nós 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 arquivo ZIP com ZIP_STORED, para que não precisemos descompactar o arquivo ZIP a cada carregamento subsequente.
2. A interface torch.jit.load também incorre no problema de cópia de memória, então escrevemos um wrapper simples para carregá-lo via o protocolo de buffer do Python, assim 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

Nosso código tem algo assim, ele 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 mencionamos anteriormente, os argumentos das funções são assumidos como objetos simples, facilmente serializáveis, mas esse padrão quebra essa suposição. Para lidar com isso, adicionamos uma lógica especial: cada resultado em cache recebe um ID anexado. Se esse objeto for usado como argumento em outra chamada, o cliente o substitui por seu ID, e o servidor pode então recuperar o objeto real com base no ID.

O modelo pipeline resultante conterá uma referência ao vae. Para simplicidade, apenas o serializamos diretamente aqui. No entanto, ao mover o UntypedStorage real para a memória compartilhada, deduplicamos qualquer dado repetido.

Poderíamos ter usado o mecanismo persistent_id do pickle, mas não tentei esse caminho. Isso é um pouco lamentável.

Benchmarking

E agora para a parte que todos adoram ver.

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

Como você pode ver, o carregamento inicial com overmind leva 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 arquivos de modelo serializados, estima-se que todo o pipeline use 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!)

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

Implantamos várias instâncias de nossa aplicação, uma para cada GPU. Assim, haverá 8 processos por nó. Após implantarmos o overmind, o uso de memória do sistema foi reduzido drasticamente. Não estávamos sofrendo com falta de memória do sistema, mas se estivéssemos, isso teria sido uma grande vitória.

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

Github

Estamos disponibilizando o código aberto no Github, ficaremos felizes se isso ajudar.