Accelerator로 모델 훈련하기

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

Dennis Lee

Data Engineer, Amazon

Trainer와 Accelerator

Accelerator와 Trainer의 사용 편의성 및 커스터마이즈 능력 비교 차트.

커스텀 훈련 루프

Trainer: 커스텀 훈련 루프 미지원
생성형 AI의 일부 고급 작업에는 두 개의 네트워크 필요

생성적 적대 신경망 이미지.

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

Trainer와 Accelerator

Accelerator와 Trainer의 사용 편의성 및 커스터마이즈 능력 비교 차트.

기본 훈련 루프 수정

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

그래디언트 초기화
지정된 디바이스로 데이터 이동: .to(device)
순전파 수행
교차 엔트로피 손실 계산
역전파로 그래디언트 계산
모델 파라미터 및 학습률 업데이트

Accelerator 객체 생성

Accelerator는 분산 훈련 인터페이스를 제공합니다

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement (bool, 기본값 True): 기본적으로 디바이스 배치를 처리

모델 및 옵티마이저 정의

사전 훈련된 모델 로드

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

Adam으로 모델 파라미터 최적화

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

스케줄러 정의

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer (obj): Adam 등 PyTorch 옵티마이저
num_warmup_steps (int): lr을 선형 증가시킬 스텝 수, int(num_training_steps * 0.1)로 설정
num_training_steps (int): 총 훈련 스텝 수, len(train_dataloader) * num_epochs로 설정

효율적인 훈련을 위한 모델 준비

prepare 메서드가 디바이스 배치를 처리합니다

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

Accelerator로 훈련 루프 구성

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

그래디언트 초기화
데이터를 디바이스로 이동

Accelerator로 훈련 루프 구성

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

그래디언트 초기화
데이터를 디바이스로 이동
수동으로 데이터 이동하는 코드 제거

Accelerator로 훈련 루프 구성

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

그래디언트 초기화
데이터를 디바이스로 이동
수동으로 데이터 이동하는 코드 제거
순전파 수행
교차 엔트로피 손실 및 그래디언트 계산

Accelerator로 훈련 루프 구성

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

그래디언트 초기화
데이터를 디바이스로 이동
수동으로 데이터 이동하는 코드 제거
순전파 수행
교차 엔트로피 손실 및 그래디언트 계산
loss.backward를 accelerator로 대체
모델 파라미터 및 학습률 업데이트

변경 사항 요약

Accelerator 적용 전

데이터를 수동으로 디바이스에 이동해야 함
- inputs.to(device)
- targets.to(device)
loss.backward()로 그래디언트 계산

Accelerator 적용 후

자동 디바이스 배치 및 데이터 병렬 처리
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
accelerator.backward(loss)로 그래디언트 동기화 처리
커스터마이즈 가능한 루프
사용자 친화적, 하드웨어 독립적, 확장 가능, 유지보수 용이

연습해 봅시다!

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