Inferência de Modelos Grandes

Um dos maiores avanços que o Accelerate oferece é o Big Model Inference, que permite realizar inferência com modelos que não cabem totalmente na sua placa de vídeo.

Este tutorial mostrará como usar o Big Model Inference no Accelerate e no ecossistema do Hugging Face.

Accelerate

Um fluxo de trabalho típico para carregar um modelo do PyTorch é mostrado abaixo. ModelClass é um modelo que excede a memória da GPU do seu dispositivo (mps ou cuda ou xpu).

import torch

my_model = ModelClass(...)
state_dict = torch.load(checkpoint_file)
my_model.load_state_dict(state_dict)

Com o Big Model Inference, o primeiro passo é inicializar um esqueleto vazio do modelo com o gerenciador de contexto init_empty_weights. Isso não requer memória, pois my_model é “sem parâmetros”.

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

Em seguida, os pesos são carregados no modelo para inferência.

O método load_checkpoint_and_dispatch() carrega um checkpoint dentro do seu modelo vazio e distribui os pesos de cada camada entre todos os dispositivos disponíveis, começando pelos dispositivos mais rápidos (GPU, MPS, XPU, NPU, MLU, SDAA, MUSA) antes de passar para os mais lentos (CPU e disco rígido).

Definir device_map="auto" preenche automaticamente todo o espaço disponível primeiro nas GPUs, depois na CPU e, por fim, no disco rígido (a opção mais lenta possível), caso ainda não haja memória suficiente.

Consulte o guia Designing a device map para mais detalhes sobre como projetar seu próprio device map.

from accelerate import load_checkpoint_and_dispatch

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

Se houver determinados “blocos” de camadas que não devem ser divididos, passe-os para no_split_module_classes (veja aqui para mais detalhes).

Os pesos de um modelo também podem ser divididos em vários checkpoints para economizar memória, como quando o state_dict não cabe na memória (veja aqui para mais detalhes).

Agora que o modelo está totalmente despachado, você pode executar a inferência.

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Cada vez que uma entrada é passada por uma camada, ela é enviada da CPU para a GPU (ou do disco para a CPU e depois para a GPU), a saída é calculada e a camada é removida da GPU, voltando para baixo na fila. Embora isso adicione algum overhead à inferência, isso permite executar modelos de qualquer tamanho no seu sistema, desde que a maior camada caiba na sua GPU.

Várias GPUs, ou “paralelismo de modelo”, podem ser utilizadas, mas apenas uma GPU estará ativa em cada momento. Isso obriga a GPU a aguardar que a GPU anterior envie a ela a saída. Você deve executar seu script normalmente com Python, em vez de outras ferramentas como torchrun e accelerate launch.

Veja um exemplo completo de Inferência de Modelos Grandes abaixo.

import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

with init_empty_weights():
    model = MyModel(...)

model = load_checkpoint_and_dispatch(
    model, checkpoint=checkpoint_file, device_map="auto"
)

input = torch.randn(2,3)
device_type = next(iter(model.parameters())).device.type
input = input.to(device_type)
output = model(input)

Ecossistema Hugging Face

Outras bibliotecas do ecossistema Hugging Face, como Transformers ou Diffusers, suportam Big Model Inference em seus construtores from_pretrained.

Você só precisa adicionar device_map="auto" em from_pretrained para habilitar o Big Model Inference.

Por exemplo, carregue o modelo de 11 bilhões de parâmetros Big Sciences T0pp com o Big Model Inference.

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")

Após carregar o modelo, as etapas de inicialização vazia e despacho inteligente de antes são executadas e o modelo fica totalmente pronto para aproveitar todos os recursos da sua máquina. Por meio desses construtores, você também pode economizar mais memória especificando o parâmetro torch_dtype para carregar um modelo em uma precisão menor.

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto", torch_dtype=torch.float16)

Próximas etapas

Para uma explicação mais detalhada sobre Big Model Inference, confira o guia conceitual!