AdamW를 활용한 균형 잡힌 학습

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

Dennis Lee

Data Engineer, Amazon

효율적인 학습

옵티마이저를 강조 표시한 코스 챕터 주제 다이어그램.

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

AdamW, Adafactor, 8-bit Adam 세 가지 옵티마이저를 나타낸 다이어그램.

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

AdamW, Adafactor, 8-bit Adam 세 가지 옵티마이저를 나타낸 다이어그램.

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

AdamW, Adafactor, 8-bit Adam 세 가지 옵티마이저를 나타낸 다이어그램.

옵티마이저 트레이드오프

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

옵티마이저 트레이드오프

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

옵티마이저 트레이드오프

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

옵티마이저 트레이드오프

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

AdamW의 작동 원리

AdamW의 작동 방식을 설명하는 다이어그램.

AdamW의 작동 원리

AdamW의 작동 방식을 설명하는 다이어그램.

AdamW의 작동 원리

AdamW의 작동 방식을 설명하는 다이어그램.

그래디언트의 지수이동평균(EMA) 계산

AdamW의 작동 원리

AdamW의 작동 방식을 설명하는 다이어그램.

그래디언트의 지수이동평균(EMA) 계산
그래디언트 제곱의 EMA 계산

AdamW의 작동 원리

AdamW의 작동 방식을 설명하는 다이어그램.

그래디언트의 지수이동평균(EMA) 계산
그래디언트 제곱의 EMA 계산

AdamW의 작동 원리

AdamW의 작동 방식을 설명하는 다이어그램.

그래디언트의 지수이동평균(EMA) 계산
그래디언트 제곱의 EMA 계산

AdamW의 메모리 사용량

AdamW 계산에 사용되는 파라미터 그래디언트, 그래디언트 EMA, 그래디언트 제곱 EMA를 나타낸 다이어그램.

각 사각형은 파라미터, 각 색상은 상태를 나타냅니다

AdamW의 메모리 사용량

AdamW 계산에 사용되는 파라미터 그래디언트, 그래디언트 EMA, 그래디언트 제곱 EMA를 나타낸 다이어그램.

각 사각형은 파라미터, 각 색상은 상태를 나타냅니다

AdamW의 메모리 사용량

AdamW 계산에 사용되는 파라미터 그래디언트, 그래디언트 EMA, 그래디언트 제곱 EMA를 나타낸 다이어그램.

각 사각형은 파라미터, 각 색상은 상태를 나타냅니다
파라미터당 메모리 = 8바이트 = 4바이트/상태 × 2개 상태
총 메모리 = 파라미터당 메모리(8바이트) × 파라미터 수

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와 Accelerator

Accelerator와 Trainer의 커스터마이징 가능성 및 사용 편의성 간 트레이드오프를 나타낸 다이어그램.

Trainer로 AdamW 구현

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}

Accelerator로 AdamW 구현

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

옵티마이저 상태 확인

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

옵티마이저 크기 계산

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

옵티마이저 크기 계산

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

연습해 봅시다!

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