Entraînement économe en mémoire avec Adafactor

Entraînement efficace de 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 montrant les 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?

Schéma de la matrice du second moment, somme des colonnes et somme des rangées.

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

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

Schéma de la matrice du second moment, somme des colonnes et somme des rangées.

Économise la mémoire en ne stockant pas la matrice du second moment
À la place, stocke la somme des colonnes et des rangées de la matrice

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

Schéma de la matrice du second moment, somme des colonnes et somme des rangées.

Économise la mémoire en ne stockant pas la matrice du second moment
À la place, stocke la somme des colonnes et des rangées de la matrice

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

Schéma de la matrice du second moment, somme des colonnes et somme des rangées.

Économise la mémoire en ne stockant pas la matrice du second moment
À la place, stocke la somme des colonnes et des rangées de la matrice
Estime la matrice complète en multipliant ces deux sommes

Implémentation avec Trainer et Accelerator

Graphique comparant la simplicité d’usage et la 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

# Assumes PyTorch 2.5 or higher
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")

Number of Adafactor parameters: 178,712
Adafactor size: 1 MB

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

Number of AdamW parameters: 131,566,188
AdamW size: 502 MB

Passons à la pratique !

Entraînement efficace de modèles d’IA avec PyTorch