Entrenamiento equilibrado con AdamW

Entrenamiento eficiente de modelos de IA con PyTorch

Dennis Lee

Data Engineer, Amazon

Entrenamiento eficiente

 

 

Diagrama que muestra los temas del curso con énfasis en optimizadores.

Entrenamiento eficiente de modelos de IA con PyTorch

Optimizadores para un entrenamiento eficiente

 

 

Diagrama que muestra tres optimizadores: AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

Optimizadores para un entrenamiento eficiente

 

 

Diagrama que muestra tres optimizadores: AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

Optimizadores para un entrenamiento eficiente

 

 

Diagrama que muestra tres optimizadores: AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

Compromisos del optimizador

Diagrama que muestra los compromisos entre número de parámetros y precisión para AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

Compromisos del optimizador

Diagrama que muestra los compromisos entre número de parámetros y precisión para AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

Compromisos del optimizador

Diagrama que muestra los compromisos entre número de parámetros y precisión para AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

Compromisos del optimizador

Diagrama que muestra los compromisos entre número de parámetros y precisión para AdamW, Adafactor y Adam de 8 bits.

Entrenamiento eficiente de modelos de IA con PyTorch

¿Cómo funciona AdamW?

Diagrama que ilustra cómo funciona AdamW.

Entrenamiento eficiente de modelos de IA con PyTorch

¿Cómo funciona AdamW?

Diagrama que ilustra cómo funciona AdamW.

Entrenamiento eficiente de modelos de IA con PyTorch

¿Cómo funciona AdamW?

Diagrama que ilustra cómo funciona AdamW.

  • Calcula la media móvil exponencial (EMA) de los gradientes
Entrenamiento eficiente de modelos de IA con PyTorch

¿Cómo funciona AdamW?

Diagrama que ilustra cómo funciona AdamW.

  • Calcula la media móvil exponencial (EMA) de los gradientes
  • Calcula la EMA de los gradientes al cuadrado
Entrenamiento eficiente de modelos de IA con PyTorch

¿Cómo funciona AdamW?

Diagrama que ilustra cómo funciona AdamW.

  • Calcula la media móvil exponencial (EMA) de los gradientes
  • Calcula la EMA de los gradientes al cuadrado
Entrenamiento eficiente de modelos de IA con PyTorch

¿Cómo funciona AdamW?

Diagrama que ilustra cómo funciona AdamW.

  • Calcula la media móvil exponencial (EMA) de los gradientes
  • Calcula la EMA de los gradientes al cuadrado
Entrenamiento eficiente de modelos de IA con PyTorch

Uso de memoria de AdamW

Diagrama que muestra las cantidades usadas en los cálculos de AdamW: gradientes de parámetros, EMA de gradientes y EMA de gradientes al cuadrado.

  • Cada cuadrado es un parámetro y cada color, un estado
Entrenamiento eficiente de modelos de IA con PyTorch

Uso de memoria de AdamW

Diagrama que muestra las cantidades usadas en los cálculos de AdamW: gradientes de parámetros, EMA de gradientes y EMA de gradientes al cuadrado.

  • Cada cuadrado es un parámetro y cada color, un estado
Entrenamiento eficiente de modelos de IA con PyTorch

Uso de memoria de AdamW

Diagrama que muestra las cantidades usadas en los cálculos de AdamW: gradientes de parámetros, EMA de gradientes y EMA de gradientes al cuadrado.

  • Cada cuadrado es un parámetro y cada color, un estado
  • Memoria por parámetro = 8 bytes = 4 bytes por estado * 2 estados
  • Memoria total = Memoria por parámetro (8 bytes) * Número de parámetros
Entrenamiento eficiente de modelos de IA con PyTorch

Estimar el uso de memoria de 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
Entrenamiento eficiente de modelos de IA con PyTorch

Trainer y Accelerator

Diagrama que muestra el compromiso entre capacidad de personalización y facilidad de uso para Accelerator y Trainer.

Entrenamiento eficiente de modelos de IA con PyTorch

Implementar AdamW con 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}
Entrenamiento eficiente de modelos de IA con PyTorch

Implementar AdamW con 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
Entrenamiento eficiente de modelos de IA con PyTorch

Inspeccionar el estado del optimizador

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.], ...])}, ...])
Entrenamiento eficiente de modelos de IA con PyTorch

Calcular el tamaño del optimizador

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
Entrenamiento eficiente de modelos de IA con PyTorch

Calcular el tamaño del optimizador

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
Entrenamiento eficiente de modelos de IA con PyTorch

¡Vamos a practicar!

Entrenamiento eficiente de modelos de IA con PyTorch

Preparing Video For Download...