Entrena modelos con Accelerator

Entrenamiento eficiente de modelos de IA con PyTorch

Dennis Lee

Data Engineer, Amazon

Trainer y Accelerator

Gráfico que compara facilidad de uso vs. capacidad de personalización para Accelerator y Trainer.

Bucles de entrenamiento personalizados

Trainer: sin bucles de entrenamiento personalizados
Algunas tareas avanzadas en IA generativa requieren dos redes

Imagen de una red generativa antagónica.

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

Trainer y Accelerator

Gráfico que compara facilidad de uso vs. capacidad de personalización para Accelerator y Trainer.

Modificar un bucle de entrenamiento 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()

Pone a cero los gradientes
Mueve los datos a un dispositivo: .to(device)
Haz un forward pass
Calcula la pérdida de entropía cruzada
Calcula gradientes en el backward pass
Actualiza parámetros y tasa de aprendizaje

Crea un objeto Accelerator

Accelerator ofrece una interfaz para entrenamiento distribuido

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

device_placement (bool, por defecto True): gestiona la asignación de dispositivo

Define el modelo y el optimizador

Carga un modelo preentrenado

from transformers import AutoModelForSequenceClassification

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

Optimiza los parámetros con Adam

from torch.optim import Adam

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

Define el 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): optimizador de PyTorch, como Adam
num_warmup_steps (int): pasos para aumentar linealmente el lr, usa int(num_training_steps * 0.1)
num_training_steps (int): pasos totales, usa len(train_dataloader) * num_epochs

Prepara el modelo para entrenar eficientemente

El método prepare gestiona la asignación de dispositivos

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

                        optimizer,

                        dataloader,

                        lr_scheduler)

Crear un bucle de entrenamiento con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

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

Pone a cero los gradientes
Antes movíamos los datos al dispositivo

Crear un bucle de entrenamiento con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

Pone a cero los gradientes
Antes movíamos los datos al dispositivo
Elimina las líneas que mueven datos manualmente

Crear un bucle de entrenamiento con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

Pone a cero los gradientes
Antes movíamos los datos al dispositivo
Elimina las líneas que mueven datos manualmente
Haz un forward pass
Calcula la pérdida de entropía cruzada y los gradientes

Crear un bucle de entrenamiento con Accelerator

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

Pone a cero los gradientes
Antes movíamos los datos al dispositivo
Elimina las líneas que mueven datos manualmente
Haz un forward pass
Calcula la pérdida de entropía cruzada y los gradientes
Sustituye loss.backward por accelerator
Actualiza parámetros y tasa de aprendizaje

Resumen de cambios

Antes de Accelerator

Había que mover los datos manualmente
- inputs.to(device)
- targets.to(device)
Gradientes con loss.backward()

Con Accelerator

Colocación automática y paralelismo de datos
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
Sincroniza gradientes con accelerator.backward(loss)
Bucle personalizable
Fácil de usar, agnóstico del hardware, escalable y mantenible

¡Vamos a practicar!

Entrenamiento eficiente de modelos de IA con PyTorch