Accelerate documentation
Accelerate를 사용하여 Local SGD 사용
Accelerate를 사용하여 Local SGD 사용
Local SGD는 분산 학습 기법으로, 모든 단계마다 그라디언트를 동기화하지 않습니다. 따라서 각 프로세스는 자체 모델 가중치 버전을 업데이트하고, 지정된 단계 수 이후에 모든 프로세스에 걸쳐 가중치를 평균화하여 동기화합니다. 이를 통해 통신 효율이 향상되며, NVLink와 같은 더 빠른 인터커넥트가 없는 컴퓨터의 경우 특히 상당한 학습 속도 향상을 가져올 수 있습니다. 그라디언트 축적(통신 효율 향상을 위해 유효 배치 크기를 늘려야 함)과 달리, Local SGD는 배치 크기나 학습률/스케줄을 변경할 필요가 없습니다. 그러나 필요에 따라 Local SGD를 그라디언트 축적과 함께 사용할 수도 있습니다.
이 튜토리얼에서는 Local SGD Accelerate를 빠르게 설정하는 방법을 확인합니다. 표준 Accelerate 설정과 비교할 때, 이 설정에는 코드 2줄만 추가로 필요합니다.
이 예제는 두 개의 배치마다 그라디언트 축적을 수행하는 매우 단순한 PyTorch 학습 루프를 사용합니다:
device = "cuda"
model.to(device)
gradient_accumulation_steps = 2
for index, batch in enumerate(training_dataloader):
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss = loss / gradient_accumulation_steps
loss.backward()
if (index + 1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()Accelerate로 변환
앞서 보여준 코드를 LocalSGD나 그라디언트 축적 헬퍼 없이 Accelerate를 사용하도록 변환합니다:
+ from accelerate import Accelerator
+ accelerator = Accelerator()
+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+ model, optimizer, training_dataloader, scheduler
+ )
for index, batch in enumerate(training_dataloader):
inputs, targets = batch
- inputs = inputs.to(device)
- targets = targets.to(device)
outputs = model(inputs)
loss = loss_function(outputs, targets)
loss = loss / gradient_accumulation_steps
+ accelerator.backward(loss)
if (index+1) % gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step()Accelerate가 모델 동기화를 처리하도록 설정
이제 남은 작업은 Accelerate가 모델 파라미터 동기화 및 그라디언트 축적을 처리하도록 하는 것입니다. 단순화를 위해 8단계마다 동기화가 필요하다고 가정해 보겠습니다. 이를 위해 with LocalSGD 문장을 하나 추가하고, 모든 옵티마이저 스텝 이후에 local_sgd.step()를 한 번 호출하면 됩니다:
+local_sgd_steps=8
+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
for batch in training_dataloader:
with accelerator.accumulate(model):
inputs, targets = batch
outputs = model(inputs)
loss = loss_function(outputs, targets)
accelerator.backward(loss)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
+ local_sgd.step()내부적으로 Local SGD 코드는 자동 그래디언트 동기화를 비활성화합니다(단, 그래디언트 축적은 여전히 정상적으로 작동합니다). 대신 학습 루프의 끝에서 local_sgd_steps 스텝마다 모델 파라미터를 평균화합니다.
한계
현재 구현은 DeepSpeed와 같은 기능을 사용하지 않는 기본 다중 GPU(또는 다중 CPU) 학습에서만 작동합니다.
참고
이 단순한 접근법의 진정한 기원은 정확히 알 수 없으나, 로컬 SGD의 개념은 꽤 오래되었으며 적어도 다음으로 거슬러 올라갑니다:
Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). Parallel SGD: When does averaging help?. arXiv preprint arXiv:1606.07365.
Local SGD이라는 용어는 다음 논문에 기인합니다(우리가 인지하지 못하는 더 이전의 참고 문헌이 있을 수 있습니다).
Stich, Sebastian Urban. “Local SGD Converges Fast and Communicates Little.” ICLR 2019-International Conference on Learning Representations. No. CONF. 2019.
Update on GitHub