분산 학습: 데이터 · 텐서 · 파이프라인 병렬과 ZeRO

분산 학습: 데이터 · 텐서 · 파이프라인 병렬과 ZeRO

7B 모델의 학습 상태만 해도 112GB로 GPU 한 장(80GB)에 들어가지 않습니다. 그래서 큰 모델은 수백~수천 장의 GPU에 나눠서 학습합니다. 나누는 방법은 크게 넷입니다. 데이터를 나누는 데이터 병렬, 그 데이터 병렬의 메모리 낭비를 없앤 ZeRO, 층 안의 행렬을 쪼개는 텐서 병렬, 층들을 단계별로 나누는 파이프라인 병렬. 이 수업은 각 방법이 무엇을 나누고 어떤 통신 비용을 치르는지를 숫자로 비교합니다.

들어가기 전에

• 스케일링 법칙과 연산 예산 — 파라미터당 16바이트, 활성화 메모리, GPU 시간 어림.
• 트랜스포머 블록 한 층을 손으로 — 어텐션 · FFN이 행렬 곱 몇 개로 이루어지는지.
• 수치 형식 · 혼합 정밀도 · 양자화 — FP16 가중치 · FP32 옵티마이저 상태의 구분.

한눈에 보기

• 데이터 병렬에서 GPU들이 기울기를 '평균 내는' 올리듀스(all-reduce)가 무엇인지 압니다.
• ZeRO가 단계별로 무엇을 나눠 GPU당 메모리를 120GB → 1.9GB로 줄이는지 계산합니다.
• 텐서 병렬이 FFN 행렬을 '열로 자르고 행으로 잘라' 통신을 한 번으로 줄이는 원리를 압니다.
• 파이프라인 버블이 무엇이고, 마이크로배치를 늘리면 왜 줄어드는지 숫자로 봅니다.
• 실제 대형 학습이 세 가지를 어떤 원칙으로 섞는지(3D 병렬) 설명할 수 있습니다.

1. 데이터 병렬: 같은 모델, 다른 데이터

직관(해설용 비유). 같은 시험지 풀이법(모델)을 가진 학생 8명이 문제 묶음을 8등분해 각자 풉니다. 끝나면 각자 '어디를 고치면 좋을지(기울기)'를 모아 평균을 내고, 모두 똑같이 고칩니다.

단계.

1. 모든 GPU가 모델 전체를 복사해 둡니다.
2. 배치를 GPU 수 $N_d$ 로 나눠 각자 순전파 · 역전파를 합니다.
3. 올리듀스(all-reduce): 모든 GPU의 기울기를 더해(평균 내) 모두에게 똑같이 나눠 줍니다.
4. 각자 같은 업데이트를 해서 모델이 계속 똑같이 유지됩니다.

통신량. ZeRO 논문의 설명대로, 표준 올리듀스는 두 단계입니다. ① 리듀스-스캐터(reduce-scatter): 기울기를 $N_d$ 조각으로 나눠, GPU마다 한 조각씩 합계를 맡습니다. ② 올개더(all-gather): 합친 조각을 모두에게 돌립니다. 데이터가 $\Psi$ 원소라면 각 단계가 $\Psi$ 씩 움직여, 한 학습 단계에 약 $2\Psi$ 의 통신이 듭니다.

한계. 계산은 나눠지지만 메모리는 하나도 줄지 않습니다. 모든 GPU가 가중치 · 기울기 · 옵티마이저 상태를 다 들고 있으니까요. 모델이 GPU 한 장에 안 들어가면 데이터 병렬만으로는 학습할 수 없습니다.

2. ZeRO: 데이터 병렬의 중복을 없애기

직관. 학생 8명이 같은 참고서 8권을 각자 들고 다닐 필요는 없습니다. 한 권을 8등분해 각자 한 부분만 들고, 필요할 때 빌려 보면 됩니다.

ZeRO(Zero Redundancy Optimizer)는 데이터 병렬을 유지하면서 모델 상태를 GPU들에 나눠 저장(partition) 합니다. 파라미터 수 $\Psi$, 옵티마이저 상태 배수 $K$ (혼합 정밀도 Adam은 12), GPU 수 $N_d$ 일 때 GPU당 메모리는:

숫자로 해 보기(논문의 그림 1 설정). $\Psi=7.5$B, $K=12$, $N_d=64$:

• 기본: $16\times7.5=120$ GB
• ZeRO-1: $4\times7.5+12\times7.5/64\approx31.4$ GB
• ZeRO-2: $2\times7.5+14\times7.5/64\approx16.6$ GB
• ZeRO-3: $120/64\approx1.9$ GB

논문은 큰 $N_d$ 에서 1단계가 약 4배, 2단계가 약 8배, 3단계가 $N_d$ 배 메모리를 줄인다고 정리합니다. 3단계에서는 순전파 · 역전파 때 필요한 층의 가중치를 그때그때 올개더로 모았다가 버립니다. 그래서 통신이 늘어납니다.

왜 1 · 2단계는 통신이 늘지 않나. 올리듀스가 원래 리듀스-스캐터 + 올개더이기 때문입니다. 기울기를 리듀스-스캐터한 시점에 GPU마다 '자기 조각'의 합계만 있으면 그 조각의 옵티마이저 업데이트를 할 수 있고, 업데이트된 가중치 조각을 올개더하면 끝납니다. 같은 양의 통신으로 중복만 없앤 셈입니다.

3. 텐서 병렬: 행렬 하나를 여러 GPU로

직관. 큰 표 계산을 여러 명이 열을 나눠 맡습니다. 각자 자기 열만 계산하고, 마지막에 합칩니다.

Megatron-LM의 FFN 분할(Shoeybi 외, 2019). FFN은 $Y=\mathrm{GeLU}(XA)$, $Z=YB$ 입니다.

1. 첫 행렬 $A$ 를 열 로 자릅니다: $A=[A_1, A_2]$. 그러면 GPU $i$ 는 $Y_i=\mathrm{GeLU}(XA_i)$ 를 혼자 계산할 수 있습니다. GeLU가 열마다 따로 적용되기 때문입니다(행으로 자르면 GeLU 전에 합쳐야 해서 동기화가 필요합니다).
2. 둘째 행렬 $B$ 를 행 으로 자릅니다: $B=\begin{bmatrix}B_1\\B_2\end{bmatrix}$. GPU $i$ 는 $Y_iB_i$ 를 계산합니다.
3. $Z=Y_1B_1+Y_2B_2$ 이므로, 마지막에 올리듀스 한 번 으로 합칩니다.

어텐션도 헤드 단위로 쿼리 · 키 · 값 행렬을 열로 나눠 같은 방식으로 처리합니다. 논문은 이렇게 하면 트랜스포머 층 하나에 순전파 올리듀스 2번, 역전파 2번이면 된다고 설명합니다.

숫자로 해 보기(해설용). $X=(1,\ 2)$, $A=\begin{pmatrix}1&0&2&1\\0&1&1&-1\end{pmatrix}$ 를 열 두 개씩 나누면 GPU 1은 $XA_1=(1,\ 2)$, GPU 2는 $XA_2=(4,\ -1)$ 을 따로 계산합니다(활성화로 ReLU를 쓰면 $(1,2)$, $(4,0)$). $B$ 를 행 두 개씩 나눠 $B_1=\begin{pmatrix}1\\1\end{pmatrix}$, $B_2=\begin{pmatrix}1\\2\end{pmatrix}$ 라 하면 GPU 1은 $1+2=3$, GPU 2는 $4+0=4$ 를 내고, 올리듀스로 합친 $Z=7$ 은 한 GPU에서 통째로 계산한 결과와 같습니다.

한계. 층마다 통신이 끼어들어 대역폭이 아주 높은 연결이 필요합니다. Megatron 논문들은 텐서 병렬을 한 서버 안(NVLink로 연결된 GPU들)에서 쓰는 것을 권합니다. 첫 논문은 이 방식으로 83억(8.3B) 파라미터 모델을 GPU 512장에서 학습하며 15.1 PFLOPS, 76% 확장 효율을 보고했습니다.

4. 파이프라인 병렬: 층을 단계로 나누기

직관. 공장 조립 라인입니다. GPU 1이 1~8층, GPU 2가 9~16층, … 을 맡고 활성화를 다음 GPU로 넘깁니다. 통신은 단계 경계에서만 일어나 텐서 병렬보다 훨씬 적습니다.

문제: 버블. 배치를 통째로 흘리면 GPU 1이 일할 때 나머지는 놀고, 마지막 GPU가 일할 때 앞쪽은 놉니다. 이 노는 시간이 파이프라인 버블 입니다.

GPipe의 해법(Huang 외, 2018). 배치를 $M$ 개 마이크로배치 로 나눠 연달아 흘립니다. GPU 1이 첫 마이크로배치를 넘기자마자 둘째를 시작하므로 여러 GPU가 동시에 일합니다. 단계 수가 $K$ 일 때 버블 비율은 $O\!\left(\frac{K-1}{M+K-1}\right)$ 이고, 논문은 $M\ge4K$ 이면 버블이 무시할 만했다고 보고합니다.

숫자로 해 보기(해설용, $K=4$ 단계).

Megatron 3D 병렬 논문(Narayanan 외, 2021)은 같은 내용을 '이상적인 계산 시간 대비 버블' $(p-1)/m$ ($p$ 단계, $m$ 마이크로배치)로 씁니다. 마이크로배치가 많을수록 좋지만, 그만큼 활성화를 보관해야 해서 메모리가 늘어납니다. 그래서 GPipe는 재계산(활성화 체크포인팅)과 함께 쓰고, 이후 연구는 순전파 · 역전파 순서를 바꾼 스케줄로 메모리를 줄였습니다.

5. 3D 병렬: 셋을 섞는 원칙

12345flowchart TB
  C[GPU 클러스터] --> DP[데이터 병렬 · ZeRO<br/>복제본끼리 기울기 평균]
  DP --> PP[파이프라인 병렬<br/>서버 사이에 층을 단계로]
  PP --> TP[텐서 병렬<br/>서버 안 GPU끼리 행렬 분할]
  TP --> K[각 GPU: 층 일부의 행렬 조각]

Narayanan 외(2021)는 세 방법을 조합(PTD-P)해 1조 파라미터 GPT 모델을 A100 3,072장에서 502 PFLOP/s(GPU당 이론 최대의 52%)로 학습했다고 보고합니다. 논문의 조합 원칙을 요약하면:

왜 이렇게 나누나. 통신이 잦은 것일수록 빠른 연결 안에 가둡니다. 텐서 병렬을 느린 서버 간 연결로 하면 대부분 시간을 기다리며 보냅니다.

6. 언제 무엇을 쓰나 (실무 감각)

• 모델이 GPU 한 장에 들어가고 데이터만 많다 → 데이터 병렬.
• 모델 상태가 안 들어간다(수십억 파라미터) → ZeRO-2 · 3 또는 이를 구현한 FSDP류 도구를 먼저 시도. 코드 변경이 적습니다.
• 수백억 이상, 수백 장 이상 → 텐서 + 파이프라인 + 데이터 병렬을 조합.
• 미세조정만 한다면 → LoRA · QLoRA 로 학습할 파라미터와 메모리를 먼저 줄이는 것이 보통 더 쌉니다.

연습 문제

1. 7B 모델, $K=12$, GPU 8장일 때 ZeRO 각 단계의 GPU당 모델 상태 메모리는?
   • 풀이: 기본 112GB, ZeRO-1 $28+84/8=38.5$ GB, ZeRO-2 $14+98/8=26.25$ GB, ZeRO-3 $112/8=14$ GB입니다.
2. 파이프라인 단계 8개, 마이크로배치 32개면 GPipe식 버블 비율은?
   • 풀이: $7/(32+7)\approx0.18$ 입니다.
3. 텐서 병렬에서 FFN 첫 행렬을 행으로 자르면 어떤 문제가 생기나요?
   • 풀이: 각 GPU가 부분합만 갖게 되어, 비선형 함수(GeLU)를 적용하기 전에 합치는 통신(동기화)이 한 번 더 필요해집니다.
4. 데이터 병렬만으로 GPU를 2배 늘리면 GPU당 메모리는?
   • 풀이: 그대로입니다. 모델 상태 전체를 모든 GPU가 복제하기 때문입니다(배치를 나눠 활성화만 줄일 수 있습니다).

스스로 점검

1. 올리듀스를 두 단계로 나누면 무엇이 좋은가요?
   • 답: 리듀스-스캐터 후 각 GPU가 자기 조각의 합계만 갖게 되어, ZeRO처럼 조각별로 옵티마이저를 나눠 맡을 수 있고 통신량은 그대로입니다.
2. ZeRO-3가 통신을 더 쓰는 이유는?
   • 답: 가중치도 나눠 저장하므로, 층을 계산할 때마다 필요한 가중치를 다른 GPU에서 모아 와야(올개더) 하기 때문입니다.
3. 텐서 병렬을 서버 안에, 파이프라인 병렬을 서버 사이에 두는 이유는?
   • 답: 텐서 병렬은 층마다 통신해 대역폭이 높은 연결이 필요하고, 파이프라인은 단계 경계에서만 통신해 느린 연결에서도 견딜 수 있기 때문입니다.

용어 정리

더 깊이 보려면

• The Ultra-Scale Playbook: huggingface.co/spaces/nanotron/ultrascale… — 각 병렬화를 실험 수치와 함께 설명한 공개 안내서
• ZeRO 원문 arxiv.org/abs/1910.02054 · Megatron-LM 원문 arxiv.org/abs/1909.08053

다음에 볼 것

• 추론 최적화: KV 캐시 · MQA/GQA · 배칭 · 추측 디코딩 · 서빙 지표 — 학습이 끝난 모델을 싸고 빠르게 서비스하는 법
• 플래시어텐션: 메모리를 아끼는 정확한 어텐션 — GPU 한 장 안에서 메모리 이동을 줄이는 법

라이선스

본문은 KAITABASE 자체 해설입니다. 메모리 · 통신 · 성능 수치는 ZeRO · Megatron-LM · GPipe · Megatron 3D 병렬 원문에서 확인한 것만 옮겼고, 행렬 분할 예와 버블 표는 해설용 계산이며 파이썬으로 검산했습니다.

출처

• Rajbhandari 외, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2019): arxiv.org/abs/1910.02054
• Shoeybi 외, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019): arxiv.org/abs/1909.08053
• Huang 외, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (2018): arxiv.org/abs/1811.06965
• Narayanan 외, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021): arxiv.org/abs/2104.04473
• Hugging Face, The Ultra-Scale Playbook: huggingface.co/spaces/nanotron/ultrascale…