Treinamento com pouca memória usando Adafactor

Treinamento Eficiente de Modelos de IA com PyTorch

Dennis Lee

Data Engineer, Amazon

Otimizadores para treinar com eficiência

 

 

Ícones representando AdamW, Adafactor e Adam de 8 bits.

Treinamento Eficiente de Modelos de IA com PyTorch

Trade-offs dos otimizadores

Diagrama mostrando os trade-offs entre número de parâmetros e precisão para AdamW, Adafactor e Adam de 8 bits.

Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor funciona?

 

Diagrama mostrando as etapas do Adafactor.

Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor funciona?

 

Diagrama mostrando as etapas do Adafactor.

Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor funciona?

 

Diagrama mostrando as etapas do Adafactor.

 

  • EMA: média móvel exponencial
  • Segundo momento: EMA dos gradientes ao quadrado
Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor funciona?

 

Diagrama mostrando as etapas do Adafactor.

 

  • EMA: média móvel exponencial
  • Segundo momento: EMA dos gradientes ao quadrado
Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor economiza memória?

 

Diagrama mostrando a matriz do segundo momento, soma das colunas e soma das linhas.

  • Economize memória sem armazenar a matriz do segundo momento
Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor economiza memória?

 

Diagrama mostrando a matriz do segundo momento, soma das colunas e soma das linhas.

  • Economize memória sem armazenar a matriz do segundo momento
  • Em vez disso, armazene a soma das colunas e a soma das linhas
Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor economiza memória?

 

Diagrama mostrando a matriz do segundo momento, soma das colunas e soma das linhas.

  • Economize memória sem armazenar a matriz do segundo momento
  • Em vez disso, armazene a soma das colunas e a soma das linhas
Treinamento Eficiente de Modelos de IA com PyTorch

Como o Adafactor economiza memória?

 

Diagrama mostrando a matriz do segundo momento, soma das colunas e soma das linhas.

  • Economize memória sem armazenar a matriz do segundo momento
  • Em vez disso, armazene a soma das colunas e a soma das linhas
  • Estime a matriz completa multiplicando as duas somas
Treinamento Eficiente de Modelos de IA com PyTorch

Implementação com Trainer e Accelerator

Gráfico comparando facilidade de uso vs. capacidade de customização para Accelerator e Trainer.

Treinamento Eficiente de Modelos de IA com PyTorch

Implemente Adafactor com o 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}
Treinamento Eficiente de Modelos de IA com PyTorch

Implemente Adafactor com o 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
Treinamento Eficiente de Modelos de IA com PyTorch

Inspecione o estado do otimizador

  • Acesse o optimizer pelo state
optimizer_state = optimizer.state.values()
  • Ou acesse o 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, ...])}, ...])
Treinamento Eficiente de Modelos de IA com PyTorch

Calcule o uso de memória do 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
  • Compare com AdamW: Adafactor usa bem menos memória!
Number of AdamW parameters: 131,566,188
AdamW size: 502 MB
Treinamento Eficiente de Modelos de IA com PyTorch

Vamos praticar!

Treinamento Eficiente de Modelos de IA com PyTorch

Preparing Video For Download...