Навчання з низьким споживанням пам’яті з Adafactor

Ефективне тренування моделей ШІ з PyTorch

Dennis Lee

Data Engineer, Amazon

Оптимізатори для ефективного навчання

Піктограми AdamW, Adafactor та 8-бітного Adam.

Компроміси оптимізаторів

Схема компромісів між кількістю параметрів і точністю для AdamW, Adafactor та 8-бітного Adam.

Як працює Adafactor?

Схема кроків Adafactor.

Як працює Adafactor?

Схема кроків Adafactor.

Як працює Adafactor?

Схема кроків Adafactor.

EMA: експоненційне ковзне середнє
Другий момент: EMA квадратів градієнтів

Як працює Adafactor?

Схема кроків Adafactor.

EMA: експоненційне ковзне середнє
Другий момент: EMA квадратів градієнтів

Як Adafactor економить пам’ять?

Схема матриці другого моменту, суми за стовпцями та за рядками.

Заощаджує пам’ять, не зберігаючи матрицю другого моменту

Як Adafactor економить пам’ять?

Схема матриці другого моменту, суми за стовпцями та за рядками.

Заощаджує пам’ять, не зберігаючи матрицю другого моменту
Натомість зберігає суму за стовпцями та за рядками матриці

Як Adafactor економить пам’ять?

Схема матриці другого моменту, суми за стовпцями та за рядками.

Заощаджує пам’ять, не зберігаючи матрицю другого моменту
Натомість зберігає суму за стовпцями та за рядками матриці

Як Adafactor економить пам’ять?

Схема матриці другого моменту, суми за стовпцями та за рядками.

Заощаджує пам’ять, не зберігаючи матрицю другого моменту
Натомість зберігає суму за стовпцями та за рядками матриці
Оцінює повну матрицю як добуток сум за стовпцями та за рядками

Реалізація через Trainer та Accelerator

Діаграма порівняння простоти використання та можливостей налаштування для Accelerator і Trainer.

Реалізація Adafactor у 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}

Реалізація Adafactor у 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

Перегляд стану оптимізатора

Доступ до optimizer через його state

optimizer_state = optimizer.state.values()

Або доступ до optimizer через 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, ...])}, ...])

Обчислення використання пам’яті 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

Порівняно з AdamW: Adafactor споживає значно менше пам’яті!

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

Passons à la pratique !

Ефективне тренування моделей ШІ з PyTorch