Treine modelos com o Accelerator

Treinamento Eficiente de Modelos de IA com PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer e Accelerator

Gráfico comparando facilidade de uso vs. capacidade de personalização para Accelerator e Trainer.

Loops de treino customizados

Trainer: sem loops de treino customizados
Algumas tarefas avançadas em IA generativa exigem duas redes

Imagem de uma generative adversarial network.

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

Trainer e Accelerator

Gráfico comparando facilidade de uso vs. capacidade de personalização para Accelerator e Trainer.

Modificando um loop de treino básico

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

Zere os gradientes
Mova dados para o dispositivo: .to(device)
Faça o forward pass
Calcule a loss de entropia cruzada
Calcule gradientes no backward pass
Atualize parâmetros e taxa de aprendizado

Crie um objeto Accelerator

Accelerator oferece uma interface para treino distribuído

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement (bool, padrão True): gerencia a alocação de dispositivos

Defina o modelo e o otimizador

Carregue um modelo pré-treinado

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

Otimize os parâmetros com Adam

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

Defina o scheduler

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer (obj): otimizador do PyTorch, como Adam
num_warmup_steps (int): passos para aumentar lr linearmente; use int(num_training_steps * 0.1)
num_training_steps (int): passos totais de treino; use len(train_dataloader) * num_epochs

Prepare o modelo para treino eficiente

O método prepare cuida da alocação de dispositivos

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

Construindo um loop de treino com o Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

Zere os gradientes
Antes, movíamos os dados para o dispositivo

Construindo um loop de treino com o Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Zere os gradientes
Antes, movíamos os dados para o dispositivo
Remova linhas que movem dados manualmente

Construindo um loop de treino com o Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Zere os gradientes
Antes, movíamos os dados para o dispositivo
Remova linhas que movem dados manualmente
Faça o forward pass
Calcule a loss de entropia cruzada e os gradientes

Construindo um loop de treino com o Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Zere os gradientes
Antes, movíamos os dados para o dispositivo
Remova linhas que movem dados manualmente
Faça o forward pass
Calcule a loss de entropia cruzada e os gradientes
Troque loss.backward por accelerator
Atualize parâmetros e taxa de aprendizado

Resumo das mudanças

Antes do Accelerator

Precisava mover dados manualmente para os dispositivos
- inputs.to(device)
- targets.to(device)
Gradientes com loss.backward()

Depois do Accelerator

Alocação automática e paralelismo de dados
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Sincroniza gradientes com accelerator.backward(loss)
Loop personalizável
Amigável, independente de hardware, escalável e fácil de manter

Vamos praticar!

Treinamento Eficiente de Modelos de IA com PyTorch