Antrenament eficient cu Adafactor

Antrenament eficient al modelelor AI cu PyTorch

Dennis Lee

Data Engineer, Amazon

Optimizatori pentru eficiența antrenamentului

Pictograme reprezentând AdamW, Adafactor și 8-bit Adam.

Compromisuri ale optimizatorilor

Diagramă care prezintă compromisurile dintre numărul de parametri și precizie pentru AdamW, Adafactor și 8-bit Adam.

Cum funcționează Adafactor?

Diagramă care prezintă pașii Adafactor.

Cum funcționează Adafactor?

Diagramă care prezintă pașii Adafactor.

Cum funcționează Adafactor?

Diagramă care prezintă pașii Adafactor.

EMA: medie mobilă exponențială
Al doilea moment: EMA a gradienților la pătrat

Cum funcționează Adafactor?

Diagramă care prezintă pașii Adafactor.

EMA: medie mobilă exponențială
Al doilea moment: EMA a gradienților la pătrat

Cum economisește Adafactor memorie?

Diagramă ce ilustrează matricea celui de-al doilea moment, suma pe coloane și suma pe rânduri.

Economisire memorie prin nestocarea matricei celui de-al doilea moment

Cum economisește Adafactor memorie?

Diagramă ce ilustrează matricea celui de-al doilea moment, suma pe coloane și suma pe rânduri.

Economisire memorie prin nestocarea matricei celui de-al doilea moment
În schimb, se stochează suma pe coloane și suma pe rânduri

Cum economisește Adafactor memorie?

Diagramă ce ilustrează matricea celui de-al doilea moment, suma pe coloane și suma pe rânduri.

Economisire memorie prin nestocarea matricei celui de-al doilea moment
În schimb, se stochează suma pe coloane și suma pe rânduri

Cum economisește Adafactor memorie?

Diagramă ce ilustrează matricea celui de-al doilea moment, suma pe coloane și suma pe rânduri.

Economisire memorie prin nestocarea matricei celui de-al doilea moment
În schimb, se stochează suma pe coloane și suma pe rânduri
Matricea completă se estimează prin înmulțirea celor două sume

Implementare cu Trainer și Accelerator

Grafic care compară ușurința de utilizare față de posibilitatea de personalizare pentru Accelerator și Trainer.

Implementarea Adafactor cu 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}

Implementarea Adafactor cu 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

Inspectarea stării optimizatorului

Accesați optimizer prin state

optimizer_state = optimizer.state.values()

Sau accesați optimizer prin 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, ...])}, ...])

Calculul utilizării memoriei de către 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

Comparație cu AdamW: Adafactor folosește mult mai puțină memorie!

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

Lass uns üben!

Antrenament eficient al modelelor AI cu PyTorch