Adafactor를 활용한 메모리 효율적 학습

PyTorch로 AI 모델 효율적으로 학습시키기

Dennis Lee

Data Engineer, Amazon

학습 효율을 위한 옵티마이저

 

 

AdamW, Adafactor, 8-bit Adam을 나타내는 아이콘.

PyTorch로 AI 모델 효율적으로 학습시키기

옵티마이저 트레이드오프

AdamW, Adafactor, 8-bit Adam의 파라미터 수와 정밀도 간 트레이드오프를 나타내는 다이어그램.

PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 작동 원리

 

Adafactor 단계를 나타내는 다이어그램.

PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 작동 원리

 

Adafactor 단계를 나타내는 다이어그램.

PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 작동 원리

 

Adafactor 단계를 나타내는 다이어그램.

 

  • EMA: 지수 이동 평균
  • 2차 모멘트: 기울기 제곱의 EMA
PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 작동 원리

 

Adafactor 단계를 나타내는 다이어그램.

 

  • EMA: 지수 이동 평균
  • 2차 모멘트: 기울기 제곱의 EMA
PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 메모리 절약 방법

 

2차 모멘트 행렬, 열 합계, 행 합계를 나타내는 다이어그램.

  • 2차 모멘트 행렬을 저장하지 않아 메모리 절약
PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 메모리 절약 방법

 

2차 모멘트 행렬, 열 합계, 행 합계를 나타내는 다이어그램.

  • 2차 모멘트 행렬을 저장하지 않아 메모리 절약
  • 대신 행렬의 열 합계와 행 합계를 저장
PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 메모리 절약 방법

 

2차 모멘트 행렬, 열 합계, 행 합계를 나타내는 다이어그램.

  • 2차 모멘트 행렬을 저장하지 않아 메모리 절약
  • 대신 행렬의 열 합계와 행 합계를 저장
PyTorch로 AI 모델 효율적으로 학습시키기

Adafactor의 메모리 절약 방법

 

2차 모멘트 행렬, 열 합계, 행 합계를 나타내는 다이어그램.

  • 2차 모멘트 행렬을 저장하지 않아 메모리 절약
  • 대신 행렬의 열 합계와 행 합계를 저장
  • 열 합계와 행 합계를 곱해 전체 행렬 추정
PyTorch로 AI 모델 효율적으로 학습시키기

Trainer 및 Accelerator 구현

Accelerator와 Trainer의 사용 편의성 대 커스터마이징 가능성을 비교한 차트.

PyTorch로 AI 모델 효율적으로 학습시키기

Trainer로 Adafactor 구현하기

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}
PyTorch로 AI 모델 효율적으로 학습시키기

Accelerator로 Adafactor 구현하기

# 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
PyTorch로 AI 모델 효율적으로 학습시키기

옵티마이저 상태 확인하기

  • state를 통해 optimizer에 접근
optimizer_state = optimizer.state.values()
  • 또는 trainer를 통해 optimizer에 접근
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, ...])}, ...])
PyTorch로 AI 모델 효율적으로 학습시키기

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
PyTorch로 AI 모델 효율적으로 학습시키기

연습해 봅시다!

PyTorch로 AI 모델 효율적으로 학습시키기

Preparing Video For Download...