Balanserad träning med AdamW

Effektiv AI-modellträning med PyTorch

Dennis Lee

Data Engineer, Amazon

Effektiv träning

Diagram som visar kursens kapitelämnen med fokus på optimerare.

Optimerare för träningseffektivitet

Diagram som visar tre optimerare: AdamW, Adafactor och 8-bit Adam.

Optimerare för träningseffektivitet

Diagram som visar tre optimerare: AdamW, Adafactor och 8-bit Adam.

Optimerare för träningseffektivitet

Diagram som visar tre optimerare: AdamW, Adafactor och 8-bit Adam.

Optimeringsavvägningar

Diagram som visar avvägningarna mellan antal parametrar och precision för AdamW, Adafactor och 8-bit Adam.

Optimeringsavvägningar

Diagram som visar avvägningarna mellan antal parametrar och precision för AdamW, Adafactor och 8-bit Adam.

Optimeringsavvägningar

Diagram som visar avvägningarna mellan antal parametrar och precision för AdamW, Adafactor och 8-bit Adam.

Optimeringsavvägningar

Diagram som visar avvägningarna mellan antal parametrar och precision för AdamW, Adafactor och 8-bit Adam.

Hur fungerar AdamW?

Diagram som illustrerar hur AdamW fungerar.

Hur fungerar AdamW?

Diagram som illustrerar hur AdamW fungerar.

Hur fungerar AdamW?

Diagram som illustrerar hur AdamW fungerar.

Beräkna exponentiellt glidande medelvärde (EMA) av gradienterna

Hur fungerar AdamW?

Diagram som illustrerar hur AdamW fungerar.

Beräkna exponentiellt glidande medelvärde (EMA) av gradienterna
Beräkna EMA av kvadrerade gradienter

Hur fungerar AdamW?

Diagram som illustrerar hur AdamW fungerar.

Beräkna exponentiellt glidande medelvärde (EMA) av gradienterna
Beräkna EMA av kvadrerade gradienter

Hur fungerar AdamW?

Diagram som illustrerar hur AdamW fungerar.

Beräkna exponentiellt glidande medelvärde (EMA) av gradienterna
Beräkna EMA av kvadrerade gradienter

Minnesanvändning för AdamW

Diagram som visar de kvantiteter som ingår i beräkningarna för AdamW: parametergradienter, EMA av gradienter och EMA av kvadrerade gradienter.

Varje ruta är en parameter och varje färg ett tillstånd

Minnesanvändning för AdamW

Diagram som visar de kvantiteter som ingår i beräkningarna för AdamW: parametergradienter, EMA av gradienter och EMA av kvadrerade gradienter.

Varje ruta är en parameter och varje färg ett tillstånd

Minnesanvändning för AdamW

Diagram som visar de kvantiteter som ingår i beräkningarna för AdamW: parametergradienter, EMA av gradienter och EMA av kvadrerade gradienter.

Varje ruta är en parameter och varje färg ett tillstånd
Minne per parameter = 8 byte = 4 byte per tillstånd * 2 tillstånd
Totalt minne = Minne per parameter (8 byte) * Antal parametrar

Uppskatta minnesanvändning för 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 och Accelerator

Diagram som visar avvägningen mellan anpassningsmöjligheter och användarvänlighet för Accelerator och Trainer.

Implementera AdamW med 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}

Implementera AdamW med 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

Granska optimeringstillståndet

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

Beräkna optimerarens storlek

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

Beräkna optimerarens storlek

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

Laten we oefenen!

Effektiv AI-modellträning med PyTorch

Preparing Video For Download...