Inferência distribuída

A inferência distribuída pode ser dividida em duas categorias:

1. Carregar um modelo inteiro em cada GPU e enviar pedaços de um lote pela cópia do modelo de cada GPU de cada vez
2. Carregar partes de um modelo em cada GPU e processar uma única entrada por vez

Este guia aborda o primeiro parêntese. Para o segundo, consulte Big Model Inference.

Enviando automaticamente blocos de um lote para cada modelo carregado

Esta é a solução que mais consome memória, pois exige que cada GPU mantenha uma cópia completa do modelo na memória em um determinado momento.

Normalmente, ao fazer isso, os usuários enviam o modelo para um dispositivo específico para carregá-lo da CPU e, em seguida, movem cada prompt para um dispositivo diferente.

Um pipeline básico usando a biblioteca diffusers pode ser algo como:

import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)

Em seguida, é realizada a inferência com base no prompt específico:

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    pipe.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    result = pipe(prompt).images[0]
    result.save(f"result_{rank}.png")

Notará-se que é necessário verificar o rank para saber qual prompt enviar, o que pode ser um pouco trabalhoso.

Um usuário também pode pensar que, com o Accelerate, usar o Accelerator para preparar um dataloader para esse tipo de tarefa seria uma maneira simples de gerenciar isso. (Para saber mais, confira a seção relevante no Quick Tour)

Ele consegue gerenciar? Sim. No entanto, ele adiciona código extra desnecessário? Também sim.

Com o Accelerate, podemos simplificar esse processo usando o gerenciador de contexto Accelerator.split_between_processes() (que também existe em PartialState e AcceleratorState). Essa função dividirá automaticamente qualquer dado que você passar para ela (seja um prompt, um conjunto de tensores, um dicionário dos dados anteriores, etc.) entre todos os processos (com a possibilidade de ser preenchido com padding) para que você possa usá-lo imediatamente.

Vamos reescrever o exemplo acima usando este gerenciador de contexto:

import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipe(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")

E, em seguida, para executar o código, podemos usar o Accelerate:

Se você gerou um arquivo de configuração para ser usado com accelerate config:

accelerate launch distributed_inference.py

Se você tem um arquivo de configuração específico que deseja usar:

accelerate launch --config_file my_config.json distributed_inference.py

Ou, se você não quiser criar nenhum arquivo de configuração e iniciar em duas GPUs:

Nota: Você receberá alguns avisos sobre valores que estão sendo estimados com base no seu sistema. Para removê-los, você pode executar accelerate config default ou passar por accelerate config para criar um arquivo de configuração.

accelerate launch --num_processes 2 distributed_inference.py

Agora reduzimos facilmente o código de boilerplate necessário para dividir esses dados para apenas algumas linhas de código.

Mas e se tivermos uma distribuição ímpar de prompts para as GPUs? Por exemplo, e se tivermos 3 prompts, mas apenas 2 GPUs?

Sob o gerenciador de contexto, a primeira GPU receberia os dois primeiros prompts e a segunda GPU o terceiro, garantindo que todos os prompts sejam divididos e que nenhum overhead seja necessário.

Porém, e se quiséssemos fazer algo com os resultados de todos os GPUs? (Por exemplo, reunir todos e realizar algum tipo de pós-processamento) Você pode passar apply_padding=True para garantir que as listas de prompts sejam preenchidas até o mesmo comprimento, com dados extras sendo obtidos da última amostra. Dessa forma, todos os GPUs terão o mesmo número de prompts, e você poderá reunir os resultados.

Isso é necessário apenas ao tentar executar uma ação, como reunir os resultados, em que os dados em cada dispositivo precisam ter o mesmo comprimento. A inferência básica não requer isso.

Por exemplo:

import torch
from accelerate import PartialState  # Can also be Accelerator or AcceleratorState
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
distributed_state = PartialState()
pipe.to(distributed_state.device)

# Assume two processes
with distributed_state.split_between_processes(["a dog", "a cat", "a chicken"], apply_padding=True) as prompt:
    result = pipe(prompt).images

Na primeira GPU, os prompts serão ["a dog", "a cat"], e na segunda GPU serão ["a chicken", "a chicken"]. Certifique-se de descartar a amostra final, pois ela será um duplicado da anterior.

Você pode encontrar exemplos mais complexos aqui, como o uso com LLMs.