Treino equilibrado com AdamW

Treinamento Eficiente de Modelos de IA com PyTorch

Dennis Lee

Data Engineer, Amazon

Treino eficiente

 

 

Diagrama mostrando os tópicos do capítulo no curso, com destaque em otimizadores.

Treinamento Eficiente de Modelos de IA com PyTorch

Otimizadores para treino eficiente

 

 

Diagrama mostrando três otimizadores: AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Otimizadores para treino eficiente

 

 

Diagrama mostrando três otimizadores: AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Otimizadores para treino eficiente

 

 

Diagrama mostrando três otimizadores: AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Trade-offs dos otimizadores

Diagrama mostrando os trade-offs entre número de parâmetros e precisão para AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Trade-offs dos otimizadores

Diagrama mostrando os trade-offs entre número de parâmetros e precisão para AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Trade-offs dos otimizadores

Diagrama mostrando os trade-offs entre número de parâmetros e precisão para AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Trade-offs dos otimizadores

Diagrama mostrando os trade-offs entre número de parâmetros e precisão para AdamW, Adafactor e Adam 8-bit.

Treinamento Eficiente de Modelos de IA com PyTorch

Como o AdamW funciona?

Diagrama ilustrando como o AdamW funciona.

Treinamento Eficiente de Modelos de IA com PyTorch

Como o AdamW funciona?

Diagrama ilustrando como o AdamW funciona.

Treinamento Eficiente de Modelos de IA com PyTorch

Como o AdamW funciona?

Diagrama ilustrando como o AdamW funciona.

  • Calcule a média móvel exponencial (EMA) dos gradientes
Treinamento Eficiente de Modelos de IA com PyTorch

Como o AdamW funciona?

Diagrama ilustrando como o AdamW funciona.

  • Calcule a média móvel exponencial (EMA) dos gradientes
  • Calcule a EMA dos gradientes ao quadrado
Treinamento Eficiente de Modelos de IA com PyTorch

Como o AdamW funciona?

Diagrama ilustrando como o AdamW funciona.

  • Calcule a média móvel exponencial (EMA) dos gradientes
  • Calcule a EMA dos gradientes ao quadrado
Treinamento Eficiente de Modelos de IA com PyTorch

Como o AdamW funciona?

Diagrama ilustrando como o AdamW funciona.

  • Calcule a média móvel exponencial (EMA) dos gradientes
  • Calcule a EMA dos gradientes ao quadrado
Treinamento Eficiente de Modelos de IA com PyTorch

Uso de memória do AdamW

Diagrama mostrando as quantidades usadas nos cálculos do AdamW: gradientes dos parâmetros, EMA dos gradientes e EMA dos gradientes ao quadrado.

  • Cada quadrado é um parâmetro e cada cor é um estado
Treinamento Eficiente de Modelos de IA com PyTorch

Uso de memória do AdamW

Diagrama mostrando as quantidades usadas nos cálculos do AdamW: gradientes dos parâmetros, EMA dos gradientes e EMA dos gradientes ao quadrado.

  • Cada quadrado é um parâmetro e cada cor é um estado
Treinamento Eficiente de Modelos de IA com PyTorch

Uso de memória do AdamW

Diagrama mostrando as quantidades usadas nos cálculos do AdamW: gradientes dos parâmetros, EMA dos gradientes e EMA dos gradientes ao quadrado.

  • Cada quadrado é um parâmetro e cada cor é um estado
  • Memória por parâmetro = 8 bytes = 4 bytes por estado * 2 estados
  • Memória total = Memória por parâmetro (8 bytes) * Nº de parâmetros
Treinamento Eficiente de Modelos de IA com PyTorch

Estimar uso de memória do AdamW

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

num_parameters = sum(p.numel() for p in model.parameters())
print(f"Number of model parameters: {num_parameters:,}")

Number of model parameters: 65,783,042

estimated_memory = num_parameters * 8 / (1024 ** 2)
print(f"Estimated memory usage of AdamW: {estimated_memory:.0f} MB")

Estimated memory usage of AdamW: 502 MB
Treinamento Eficiente de Modelos de IA com PyTorch

Trainer e Accelerator

Diagrama mostrando o trade-off entre possibilidade de customização e facilidade de uso para Accelerator e Trainer.

Treinamento Eficiente de Modelos de IA com PyTorch

Implementar AdamW com Trainer

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics,
                  optimizers=(optimizer, lr_scheduler))

trainer.train()

{'epoch': 1.0, 'eval_accuracy': 0.7, 'eval_f1': 0.8}
Treinamento Eficiente de Modelos de IA com PyTorch

Implementar AdamW com Accelerator

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    lr_scheduler.step()
    optimizer.zero_grad()
    print(f"Loss = {loss}")

Loss = 0.7
Treinamento Eficiente de Modelos de IA com PyTorch

Inspecionar o estado do otimizador

optimizer_state = optimizer.state.values()
print(optimizer_state)

dict_values([{'step': tensor(3.),

              'exp_avg': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...]),

              'exp_avg_sq': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...])}, ...])
Treinamento Eficiente de Modelos de IA com PyTorch

Calcular o tamanho do otimizador

def compute_optimizer_size(optimizer_state):
    total_size_megabytes, total_num_elements = 0, 0

    for params in optimizer_state:

        for name, tensor in params.items():
            tensor = torch.tensor(tensor)

            num_elements = tensor.numel()

            element_size = tensor.element_size()

            total_num_elements += num_elements

            total_size_megabytes += num_elements * element_size / (1024 ** 2)

    return total_size_megabytes, total_num_elements
Treinamento Eficiente de Modelos de IA com PyTorch

Calcular o tamanho do otimizador

total_size_megabytes, total_num_elements = \
    compute_optimizer_size(trainer.optimizer.state.values())
print(f"Number of optimizer parameters: {total_num_elements:,}")

Number of optimizer parameters: 131,566,188

print(f"Optimizer size: {total_size_megabytes:.0f} MB")

Optimizer size: 502 MB
Treinamento Eficiente de Modelos de IA com PyTorch

Vamos praticar!

Treinamento Eficiente de Modelos de IA com PyTorch

Preparing Video For Download...