Deepspeed 배치 사이즈 관련 파라미터 정리

Deepspeed를 활용하여 multi-gpu로 분산 훈련을 진행할 때 배치 사이즈의 개념과 파라미터로 들어가는 값들에 대해서 정리해보았다.

 

https://www.deepspeed.ai/docs/config-json/#batch-size-related-parameters

DeepSpeed Configuration JSON

 

train_batch_size = (train_micro_batch_size_per_gpu) * (gradient_accumulation) * (number of GPUs)

 

train_batch_size

• 실질적인 훈련 배치 사이즈
• 1 스텝의 모델 업데이트가 이루어질 때까지의 데이터 샘플 개수
• 한 개의 GPU가 처리하는 배치 사이즈 * Gradient 누적 횟수 * 총 GPU 개수 

 

train_micro_batch_size_per_gpu

• 한 개의 GPU가 gradient 누적 없이 한 번에 처리하는 배치 사이즈

 

gradient_accumulation_steps

• Gradient를 평균내고 합치기 전에 얼만큼 누적을 해둘 것인지
• 값이 커지면, 그래디언트 통신을 줄일 수 있음
• 값이 작아지면, 더 큰 train_batch_size를 사용할 수 있어서 작은 GPU에서 더 큰 모델을 훈련할 수 있음

 

 

 

예를 들어, 다음과 같은 조건으로 훈련을 진행한다면,

• 총 데이터 샘플 수: 320,000
• train_micro_batch_size_per_gpu: 1
• gradient_accumulation_steps: 64
• 총 GPU 개수: 16

 

1. 하나의 GPU는 1개의 데이터 샘플을 처리하고, gradient를 64번 누적한다. 즉, 하나의 GPU는 64개의 데이터 샘플을 보고 그에 대한 gradient를 누적한 후에 모델을 업데이트한다. (1 * 64) 
2. 16개의 GPU가 있으므로, 16 x 64 = 1024개의 데이터 샘플이 실질적으로 한 번의 모델 업데이트에 사용된다. (1 * 64 *16)
3. 총 샘플 수가 320,000개이므로 1 epoch을 끝내기 위해서는 320000/1024 = 312.5 번의 업데이트가 필요하다.