Entraînement économe en mémoire avec Adafactor

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

Dennis Lee

Data Engineer, Amazon

Optimiseurs pour un entraînement efficace

Icônes représentant AdamW, Adafactor et Adam 8 bits.

Compromis des optimiseurs

Schéma des compromis entre nombre de paramètres et précision pour AdamW, Adafactor et Adam 8 bits.

Comment fonctionne Adafactor ?

Schéma des étapes d’Adafactor.

Comment fonctionne Adafactor ?

Schéma des étapes d’Adafactor.

Comment fonctionne Adafactor ?

Schéma des étapes d’Adafactor.

EMA : moyenne mobile exponentielle
Second moment : EMA des gradients au carré

Comment fonctionne Adafactor ?

Schéma des étapes d’Adafactor.

EMA : moyenne mobile exponentielle
Second moment : EMA des gradients au carré

Comment Adafactor économise-t-il la mémoire ?

Diagramme montrant la matrice du second moment, la somme des colonnes et des lignes.

Économise la mémoire en ne stockant pas la matrice du second moment

Comment Adafactor économise-t-il la mémoire ?

Diagramme montrant la matrice du second moment, la somme des colonnes et des lignes.

Économise la mémoire en ne stockant pas la matrice du second moment
Stocke à la place la somme des colonnes et des lignes de la matrice

Comment Adafactor économise-t-il la mémoire ?

Diagramme montrant la matrice du second moment, la somme des colonnes et des lignes.

Économise la mémoire en ne stockant pas la matrice du second moment
Stocke à la place la somme des colonnes et des lignes de la matrice

Comment Adafactor économise-t-il la mémoire ?

Diagramme montrant la matrice du second moment, la somme des colonnes et des lignes.

Économise la mémoire en ne stockant pas la matrice du second moment
Stocke à la place la somme des colonnes et des lignes de la matrice
Estime la matrice complète en multipliant somme des colonnes et des lignes

Implémentation avec Trainer et Accelerator

Graphique comparant facilité d’usage vs. possibilité de personnalisation pour Accelerator et Trainer.

Implémenter Adafactor avec Trainer

training_args = TrainingArguments(output_dir="./results",
                                  evaluation_strategy="epoch",

                                  optim="adafactor")

trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics)
trainer.train()

{'epoch': 1.0, 'eval_accuracy': 0.6, 'eval_f1': 0.5}

Implémenter Adafactor avec Accelerator

# Suppose PyTorch 2.5 ou supérieur
from torch.optim import Adafactor

optimizer = Adafactor(params=model.parameters(), lr=lr)


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.71

Inspecter l’état de l’optimiseur

Accéder à optimizer via son state

optimizer_state = optimizer.state.values()

Ou accéder à optimizer via trainer

optimizer_state = trainer.optimizer.state.values()

print(optimizer_state)

dict_values([{'step': tensor(3.),
              'exp_avg_sq_row': tensor([1.0000e-30, 1.0000e-30, 1.0000e-30,  ...]), 
              'exp_avg_sq_col': tensor([4.6147e-11, 5.5115e-11, 1.6338e-10, ...])}, ...])

Calculer l’usage mémoire d’Adafactor

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

Nombre de paramètres Adafactor : 178 712
Taille d’Adafactor : 1 Mo

Comparé à AdamW : Adafactor utilise bien moins de mémoire !

Nombre de paramètres AdamW : 131 566 188
Taille d’AdamW : 502 Mo

Passons à la pratique !

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