Hugging Face Accelerate

Accelerate는 같은 PyTorch 코드를 어떤 분산 설정에서든 실행할 수 있게 해주는 라이브러리예요. 코드를 네 줄만 바꾸면 분산 학습·추론을 할 수 있다는 점이 핵심이에요. DeepSpeed, FSDP, 혼합 정밀도까지 자동 지원이라 "스케일에서의 학습과 추론을 쉽고 효율적으로" 만들어 줘요.

torch.distributed와 torch_xla 위에 만들어져 있어서, 플랫폼에 맞춘 커스텀 코드를 직접 작성할 필요가 없어요. 그래서 멀티 GPU, TPU, CPU 어디서든 같은 학습 코드가 그대로 동작해요.

출처: Hugging Face Accelerate 문서 (공식)

네 줄만 바꾸면 되는 이유

기존 PyTorch 학습 코드에 Accelerate를 더하면, 모델·옵티마이저·데이터로더·스케줄러를 accelerator.prepare()로 감싸고 loss.backward()를 accelerator.backward()로 바꾸는 정도면 충분해요.

+ from accelerate import Accelerator
+ accelerator = Accelerator()

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

  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)
+     accelerator.backward(loss)
      optimizer.step()
      scheduler.step()

그러면 이 코드를 Accelerate의 CLI로 어떤 환경에서든 띄울 수 있어요.

accelerate launch {my_script.py}

Accelerator가 해주는 일

Accelerator는 분산 설정(프로세스 수, 하드웨어 유형)을 알고 있어요. 어떤 디바이스로 PyTorch 객체를 옮길지도 알기 때문에, 디바이스 관리를 Accelerate에 맡기는 걸 권장해요.

from accelerate import Accelerator

accelerator = Accelerator()
device = accelerator.device

prepare()는 모델을 단일 GPU·멀티 GPU 컨테이너에 맞게 배치하고, 옵티마이저·스케줄러를 Accelerate가 관리하는 래퍼로 감싸며, 프로세스 간에 샤딩되는 데이터로더를 만들어 줘요. 객체들은 전달한 순서 그대로 반환돼요.

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

학습 루프 바꾸기

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

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

학습 기능 추가하기

Accelerate는 그라디언트 누적, 그라디언트 클리핑, 혼합 정밀도 같은 기능을 손쉽게 추가할 수 있게 해줘요.

그라디언트 누적은 여러 배치에 걸쳐 그라디언트를 쌓아서 더 큰 배치 크기로 학습하는 효과를 내요. gradient_accumulation_steps을 지정하고 accumulate() 컨텍스트를 쓰면 돼요.

accelerator = Accelerator(gradient_accumulation_steps=2)

for input, label in training_dataloader:
    with accelerator.accumulate(model):
        predictions = model(input)
        loss = loss_function(predictions, label)
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

혼합 정밀도는 fp16 같은 낮은 정밀도 데이터 타입으로 계산해서 학습을 가속해요. mixed_precision을 지정하고 필요할 때 autocast() 컨텍스트를 쓰면 돼요.

accelerator = Accelerator(mixed_precision="fp16")
with accelerator.autocast():
    loss = complex_loss_function(outputs, target)

모델 안에서 loss를 계산하는 방식(Transformers 모델처럼)이 성능에 가장 좋아요. 모델 밖의 계산은 전체 정밀도로 이뤄지거든요.

저장과 로드

학습이 끝나면 unwrap_model()로 래퍼를 풀어서 저장해야 해요. 래퍼를 풀지 않으면 분산 학습 레이어까지 함께 저장돼서, 원래 모델로 가중치를 로드하지 못하게 돼요.

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

학습 중 상태를 저장·복원하려면 save_state()와 load_state()를 쓰세요. 옵티마이저, 랜덤 제너레이터, 학습률 스케줄러 상태까지 같은 스크립트 안에서 복원할 수 있어요.

Hugging Face Accelerate