[AI System] Training With Multiple GPUs

"Training With Multiple GPUs"

[Data Parallelism Training Process]

▣ Data Parallelism

· Dataset을 여러 GPU에 나누어 각자 배정된 데이터를 동시에 연산 수행

· 모든 GPU는 동일한 Model을 공유

→ 각 GPU는 자신에게 할당된 Data Sample을 사용해 Gradient 계산

▣ Components

· PS(Parameter Server) : 여러 Worker가 보낸 Gradient를 모아 처리

· Workers : 실제 Gradient 계산 수행

▣ Training Process

> Step 0 : Dataset Splitting + Model Initialization + Download

· 전체 Dataset D를 여러 부분으로 나누어{D₁, D₂, D₃, D₄} 여러 Worker Node에 무작위로 균등하게 분배

· Parameter Server가 초기 Model(w₀)을 만들고, 해당 Model을 모든 Worker에게 전달

    → 각 Worker는 같은 초기 Parameter로 시작 (∵ 데이터는 다르지만 시작 Model은 같아야 이후 Gradient 평균이 의미를 가짐)

> Step 1 : Gradient Computation

· 각 Worker는 현재 Model과 자신에게 할당된 데이터(D_i)를 사용해서 Gradient(g_i)를 계산 = Backpropagation

> Step 2 : Gradient Upload

· 각 Worker는 계산한 Gradient를 Parameter Server로 보내고, Server는 이를 모아 평균 Gradient를 생성

> Step 3 : Model Update

· Server는 모아진 Gradient를 이용해 Parameter를 Update

> Step 4 : Model Broadcasting

· Server가 방금 갱신한 Model(w₁)을 다시 모든 Worker에게 전송

→ 다음 Iteration에서 모든 Worker가 동일한 최신 Model로 다시 Gradient 계산

> Repeat Step 1 ~ 4 : Training Rounds (= Communication Rounds)

w₀ → w₁ → w₂ → … → w_T-1 → w_T

· Conventional Centralized Training의 Multiple Epoch 학습 과정과 유사

[Comparison With Centralized Training]

▣ Full-Batch Gradient Descent

· 데이터가 균등하게 분배되어 있다면 각 Worker가 가진 데이터 수가 같으므로, Worker 평균 = 전체 데이터 평균

· Worker가 Gradient 계산 → Gradient를 Server로 전송 → Server가 Model Update

· Worker가 Gradient 계산 → Worker가 자기 Local Model을 한 번 Update → Update된 Model을 Server로 전송 → 동일한 결과

∴ Model Averaging = Gradient Averaging

▣ Mini-Batch Gradient Descent

> Centralized Training

① 각 Iteration마다 Epoch의 첫 Mini-Batch이면 Dataset을 Shuffle

② 크기가 B인 Mini-Batch(D_b)를 뽑고 그 Mini-Batch에 대해 Gradient를 계산

③ Model Update

→ 한 번의 Update마다 전체 데이터가 아니라 일부 Sample만 사용

> Distributed Mini-Batch Training

① Worker가 N개 있을 때, 각 Worker가 자기 데이터 안에서 크기 B/N의 Mini-Batch(D_b)를 선택

② 각 Worker는 자기 Mini-Batch에 대한 Gradient(g_tⁱ)를 계산해서 Server에 전송

③ Server는 그것을 평균내어 Global Update 수행

→ 전체 관점에서는 한 Iteration에서 총 B개 Sample을 본 것과 유사

· Mini-Batch를 균등 무작위로 Sampling하면 두 방식은 통계적으로 같은 값으로 수렴

∵ Full-Batch는 모든 데이터를 다 쓰므로 Gradient가 정확히 동일

→ But, Mini-Batch는 매 Iteration마다 일부 데이터만 사용하므로, 실제 Gradient는 랜덤하게 변화

∴ 각 Iteration의 Update는 Sampling된 데이터에 따라 달라질 가능성 존재

** Distributed Training과 Centralized Training은 본질적으로 같음

[Memory & Delay]

▣ Memory

> Parameter Memory per GPU

· Centralized Training : 2m (SGD) / Data Parallelism : 2m (SGD)

→ 각 GPU가 동일한 Model 전체를 복제해서 가지고 있기 때문에 데이터 병렬화만으로는 SGD 기준 Parameter Memory 감소 X

∵ 데이터의 일부만 받았어도, 각 GPU가 해당 데이터를 처리하려면 완전한 형태의 신경망 Model이 필요하여 Model 전체를 복제

> Activation Memory per GPU

· Centralized Training : α · B / Data Parallelism : α · B/N

→ 목표 Mini-Batch Size B를 고정하면 각 GPU는 그 중 B/N개 Sample만 처리하므로 Activation Memory가 1/N으로 감소

Ex. ResNet-18 has roughly 11.7M Parameters. If we store parameters in FP32(4Bytes) : 11.7M x 4B ≒ 46.8MB

> Centralized Training

· Parameter Memory : 46.8MB x 2 = 93.6MB (SGD)

· Activation Memory : 40.9MB x 512 = 20.94GB (Mini-Batch Size 512)

→ Total Memory : 21.03GB

> Distributed Training (Data Parallelism)

· Parameter Memory : 46.8MB x 2 = 93.6MB (SGD)

· Activation Memory : 40.9MB x 512/10 = 2.094GB (Mini-Batch Size 512)

→ Total Memory : 2.19GB

∴ 12GB GPU에서 Centralized Training은 불가능하지만 Distributed Training은 가능

→ Data Parallelism은 메모리 제한 때문에 돌리지 못하는 큰 Batch 학습을 가능하게 함

▣ Delay

· 통신 비용 : 2m per GPU (Gradient Upload, Model Download)

· Delay (per Training Round) = GPU의 Gradient 계산 시간 (∝ B/N)

                                                    + Gradient Upload를 위한 GPU의 통신 시간(∝ m)

                                                    + Server의 Gradient Descent 계산 시간

                                                    + Model Transmission Time(Broadcasting the Updated Model to all Workers)

△ 각 GPU가 적은 데이터만 처리하므로 계산 시간 1/N로 감소

▼ 통신 지연 추가

→ 일반적으로 병렬 계산 덕분에 전체 학습 시간 감소

▣ Data Parallelism vs Federated Learning

> Communication

· Data Parallelism : 주로 GPU와 PS 사이의 유선 통신

· Federated Learning : Edge Device와 Server/Base Station(기지국) 사이의 무선 통신

> Data Distribution

· Data Parallelism : 데이터를 의도적으로 나누기 때문에 대체로 비슷한 분포 가능

· Federated Learning : 사용자별로 자연스럽게 수집된 데이터라서 Node마다 분포가 다를 수 있음 (Non-IID)

** Data Parallelism에서도 Gradient Compression with Error Feedback 적용 가능

→ Gradient를 그대로 보내지 않고 압축해서 보내 통신 효율 증가

∴ Data Parallelism

· Centralized Training과 동등한 성능

· Activation Memory 감소

· Optimizer에 따라 Parameter Memory도 감소 가능

· 통신 지연 조건에 따라 전체 Training Time 감소 가능

[Reference]

· Training With Multiple GPUs (AI System : Week 7) - Dongjun Han