기존 PyTorch 코드에 Accelerate 추가하기

분산 학습 프레임워크마다 자기만의 방식이 있어서, PyTorch 코드를 분산 환경에 맞추려면 커스텀 코드를 많이 써야 했어요. Accelerate는 그 세부를 대신 처리해 주는 친절한 인터페이스예요. 이 페이지에서는 기본 PyTorch 학습 루프를 점진적으로 Accelerate로 바꿔 보면서, 어떤 부분이 사라지고 어떤 부분이 추가되는지 보여드릴게요.

출처: Add Accelerate to your code (공식)

원래 PyTorch 학습 루프

시작점으로 model과 optimizer가 이미 준비된 기본 PyTorch 학습 루프를 봅시다.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Accelerator 생성하기

Accelerator는 코드를 Accelerate에 맞추는 핵심 클래스예요. 사용 중인 분산 설정(프로세스 수, 하드웨어 유형)을 알고 있고, 여러 디바이스에 걸친 프로세스 관리 메서드를 제공해요. 스크립트에서 가장 먼저 import하고 인스턴스를 만드는 게 규칙이에요.

from accelerate import Accelerator

accelerator = Accelerator()

Accelerator는 PyTorch 객체를 어느 디바이스로 옮길지도 알기 때문에, 디바이스 관련 처리를 Accelerate에 맡기세요.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

PyTorch 객체 준비하기

prepare() 메서드가 모델·옵티마이저·스케줄러 등을 분산 학습에 맞게 준비해요. 단일 GPU든 멀티 GPU든 적절한 컨테이너에 모델을 배치하고, 옵티마이저와 스케줄러를 Accelerate의 AcceleratedOptimizer·AcceleratedScheduler로 감싸며, 프로세스 간에 샤딩되는 데이터로더를 만들어요.

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

객체는 전달한 순서 그대로 반환돼요. Accelerate는 torch.optim.Optimizer 같은 각 PyTorch 클래스에서 상속한 객체만 준비해요.

학습 루프 변경하기

Accelerate의 DataLoader가 자동으로 올바른 디바이스에 배치해 주므로, 입력과 타깃에 대한 to(device) 호출을 제거해요. 그리고 역전파는 accelerator.backward()로 바꾸는데, 이 메서드가 그래디언트를 스케일링하고 분산 설정(예: DeepSpeed, Megatron)에 맞는 올바른 backward()를 골라 줘요.

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

모두 적용한 새 학습 루프는 이렇게 됩니다.

from accelerate import Accelerator
accelerator = Accelerator()

device = accelerator.device
model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()

모델 저장과 로드

모든 프로세스가 끝난 뒤에 unwrap_model()로 래퍼를 풀어서 저장해야 해요. prepare()가 모델을 분산 학습을 위한 인터페이스로 감쌌기 때문이에요. 풀지 않으면 큰 모델의 추가 레이어까지 저장돼서, 원래 모델로 가중치를 로드할 수 없어요.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

save_model()은 모델 state dict를 풀어서 저장하고, 샤딩된 체크포인트나 safetensors 형식으로도 저장할 수 있어요.

가중치를 로드하려면 역시 unwrap_model()로 먼저 푼 뒤 로드해요.

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory, "pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

학습 중간 상태(모델, 옵티마이저, 랜덤 제너레이터, 스케줄러)를 저장·복원하려면 save_state()와 load_state()를 써요.

기존 PyTorch 코드에 Accelerate 추가하기