Entraînement équilibré avec AdamW

Entraîner efficacement des modèles d’IA avec PyTorch

Dennis Lee

Data Engineer, Amazon

Entraînement efficace

Diagramme montrant les thèmes du chapitre du cours avec un focus sur les optimiseurs.

Optimiseurs pour un entraînement efficace

Diagramme présentant trois optimiseurs : AdamW, Adafactor et Adam 8 bits.

Optimiseurs pour un entraînement efficace

Diagramme présentant trois optimiseurs : AdamW, Adafactor et Adam 8 bits.

Optimiseurs pour un entraînement efficace

Diagramme présentant trois optimiseurs : AdamW, Adafactor et Adam 8 bits.

Compromis des optimiseurs

Diagramme montrant les compromis entre nombre de paramètres et précision pour AdamW, Adafactor et Adam 8 bits.

Compromis des optimiseurs

Diagramme montrant les compromis entre nombre de paramètres et précision pour AdamW, Adafactor et Adam 8 bits.

Compromis des optimiseurs

Diagramme montrant les compromis entre nombre de paramètres et précision pour AdamW, Adafactor et Adam 8 bits.

Compromis des optimiseurs

Diagramme montrant les compromis entre nombre de paramètres et précision pour AdamW, Adafactor et Adam 8 bits.

Comment fonctionne AdamW ?

Diagramme illustrant le fonctionnement d’AdamW.

Comment fonctionne AdamW ?

Diagramme illustrant le fonctionnement d’AdamW.

Comment fonctionne AdamW ?

Diagramme illustrant le fonctionnement d’AdamW.

Calculer la moyenne mobile exponentielle (EMA) des gradients

Comment fonctionne AdamW ?

Diagramme illustrant le fonctionnement d’AdamW.

Calculer la moyenne mobile exponentielle (EMA) des gradients
Calculer l’EMA des gradients au carré

Comment fonctionne AdamW ?

Diagramme illustrant le fonctionnement d’AdamW.

Calculer la moyenne mobile exponentielle (EMA) des gradients
Calculer l’EMA des gradients au carré

Comment fonctionne AdamW ?

Diagramme illustrant le fonctionnement d’AdamW.

Calculer la moyenne mobile exponentielle (EMA) des gradients
Calculer l’EMA des gradients au carré

Utilisation mémoire d’AdamW

Diagramme montrant les quantités utilisées par AdamW : gradients des paramètres, EMA des gradients et EMA des gradients au carré.

Chaque carré est un paramètre et chaque couleur un état

Utilisation mémoire d’AdamW

Diagramme montrant les quantités utilisées par AdamW : gradients des paramètres, EMA des gradients et EMA des gradients au carré.

Chaque carré est un paramètre et chaque couleur un état

Utilisation mémoire d’AdamW

Diagramme montrant les quantités utilisées par AdamW : gradients des paramètres, EMA des gradients et EMA des gradients au carré.

Chaque carré est un paramètre et chaque couleur un état
Mémoire par paramètre = 8 octets = 4 octets par état × 2 états
Mémoire totale = Mémoire par paramètre (8 octets) × Nombre de paramètres

Estimer la mémoire utilisée par 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

Trainer et Accelerator

Diagramme montrant le compromis entre capacité de personnalisation et facilité d’utilisation pour Accelerator et Trainer.

Implémenter AdamW avec 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}

Implémenter AdamW avec 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

Inspecter l’état de l’optimiseur

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.], ...])}, ...])

Calculer la taille de l’optimiseur

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

Calculer la taille de l’optimiseur

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

Passons à la pratique !

Entraîner efficacement des modèles d’IA avec PyTorch

Preparing Video For Download...