Training bilanciato con AdamW

Efficient AI Model Training with PyTorch

Dennis Lee

Data Engineer, Amazon

Training efficiente

Diagramma che mostra gli argomenti del corso con evidenza sugli optimizer.

Optimizer per un training efficiente

Diagramma che rappresenta tre optimizer: AdamW, Adafactor e Adam a 8 bit.

Optimizer per un training efficiente

Diagramma che rappresenta tre optimizer: AdamW, Adafactor e Adam a 8 bit.

Optimizer per un training efficiente

Diagramma che rappresenta tre optimizer: AdamW, Adafactor e Adam a 8 bit.

Compromessi degli optimizer

Diagramma che mostra i compromessi tra numero di parametri e precisione per AdamW, Adafactor e Adam a 8 bit.

Compromessi degli optimizer

Diagramma che mostra i compromessi tra numero di parametri e precisione per AdamW, Adafactor e Adam a 8 bit.

Compromessi degli optimizer

Diagramma che mostra i compromessi tra numero di parametri e precisione per AdamW, Adafactor e Adam a 8 bit.

Compromessi degli optimizer

Diagramma che mostra i compromessi tra numero di parametri e precisione per AdamW, Adafactor e Adam a 8 bit.

Come funziona AdamW?

Diagramma che illustra come funziona AdamW.

Come funziona AdamW?

Diagramma che illustra come funziona AdamW.

Come funziona AdamW?

Diagramma che illustra come funziona AdamW.

Calcola la media mobile esponenziale (EMA) dei gradienti

Come funziona AdamW?

Diagramma che illustra come funziona AdamW.

Calcola la media mobile esponenziale (EMA) dei gradienti
Calcola l’EMA dei gradienti al quadrato

Come funziona AdamW?

Diagramma che illustra come funziona AdamW.

Calcola la media mobile esponenziale (EMA) dei gradienti
Calcola l’EMA dei gradienti al quadrato

Come funziona AdamW?

Diagramma che illustra come funziona AdamW.

Calcola la media mobile esponenziale (EMA) dei gradienti
Calcola l’EMA dei gradienti al quadrato

Uso di memoria di AdamW

Diagramma che mostra le quantità usate nei calcoli per AdamW: gradienti dei parametri, EMA dei gradienti ed EMA dei gradienti al quadrato.

Ogni quadrato è un parametro e ogni colore è uno stato

Uso di memoria di AdamW

Diagramma che mostra le quantità usate nei calcoli per AdamW: gradienti dei parametri, EMA dei gradienti ed EMA dei gradienti al quadrato.

Ogni quadrato è un parametro e ogni colore è uno stato

Uso di memoria di AdamW

Diagramma che mostra le quantità usate nei calcoli per AdamW: gradienti dei parametri, EMA dei gradienti ed EMA dei gradienti al quadrato.

Ogni quadrato è un parametro e ogni colore è uno stato
Memoria per parametro = 8 byte = 4 byte per stato * 2 stati
Memoria totale = Memoria per parametro (8 byte) * Numero di parametri

Stima dell’uso di memoria di 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 e Accelerator

Diagramma che mostra il compromesso tra possibilità di personalizzazione e facilità d’uso per Accelerator e Trainer.

Implementare 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}

Implementare 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

Ispezionare lo stato dell’optimizer

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

Calcolare la dimensione dell’optimizer

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

Calcolare la dimensione dell’optimizer

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

Passiamo alla pratica!

Efficient AI Model Training with PyTorch

Preparing Video For Download...