대규모 모델 학습·추론 최적화 시리즈 4편: FSDP · ZeRO · DeepSpeed

"단일 GPU로는 불가능한 대형 모델, 어떻게 학습할까?"
FSDP, ZeRO, DeepSpeed는 모델 파라미터와 optimizer state를 여러 GPU에 분산해
거대한 모델을 효율적으로 학습할 수 있도록 해줍니다.

 

---

FSDP · ZeRO · DeepSpeed

 

문제 정의: 대규모 모델 학습의 한계

• 모델 파라미터 수 → 수십억 이상 (GPT-3, LLaMA 등)
• GPU 메모리 한계로 모델 전체 weight조차 올릴 수 없음
• Optimizer state (Adam: m, v 벡터)까지 고려하면 메모리 3~4배 필요
• 데이터 병렬만으로는 한계 → 파라미터 분산 병렬(Sharding) 필요

---

ZeRO (Zero Redundancy Optimizer)

ZeRO는 모델 학습 시 중복 저장되는 상태를 GPU마다 나누어 저장합니다.

• Stage 1: Optimizer states 분산 (m, v shard)
• Stage 2: Gradient 분산 (각 GPU가 gradient 일부만 보유)
• Stage 3: 파라미터 shard (모델 weight 자체도 쪼개서 GPU에 분산)

ZeRO-3까지 가면 단일 GPU가 모델 전체 weight를 가질 필요 없음 → 거대한 모델 학습 가능

---

FSDP (Fully Sharded Data Parallel)

FSDP는 PyTorch 공식 API로 제공되는 완전 샤딩 데이터 병렬 방식.

• 모델 weight, gradient, optimizer state 모두 shard
• Forward/Backward 시 필요한 shard만 gather → 연산
• Optimizer step 후 다시 shard 분산
• ZeRO-3와 동등하지만 PyTorch native 구현 → 안정적

 

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

fsdp_model = FSDP(model, auto_wrap_policy=...)

---

DeepSpeed

DeepSpeed는 Microsoft가 개발한 분산 학습 라이브러리로,
ZeRO 구현 + 다양한 분산 학습 기능을 제공.

• ZeRO Stage 1/2/3 지원
• CPU Offloading → 일부 optimizer state CPU 메모리에 저장
• Activation Checkpointing + Mixed Precision 통합
• Pipeline Parallel + Tensor Parallel 조합 가능

 

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

---

장점

1. 거대한 모델 학습 가능 (GPU 메모리 효율 극대화)
2. Throughput 향상 (GPU 간 load balancing)
3. 다양한 최적화 조합 (AMP, Gradient Checkpointing, Offloading 등)
4. 클러스터 환경 확장 용이 (수십~수백 GPU까지 확장)

---

단점 & 고려사항

1. 통신 오버헤드 증가 → 고속 네트워크(NVLink, Infiniband) 필수
2. 초기 세팅 복잡 → config, wrap policy, bucket size 등 튜닝 필요
3. 작은 모델에는 오히려 비효율 → 중대형 모델에서만 이점 큼

---

실제 효과 (예시)

모델 일반 Data Paralle lZeRO-3 메모리 절감

 

모델	일반 Data Paralle	lZeRO-3	메모리 절감
GPT-3 13B	불가 (OOM)	가능	-80%
BERT-Large	64 GPU	16 GPU	동일 batch 유지

ZeRO-3/FSDP 덕분에 한정된 GPU 수로 대형 모델 학습 가능

---

실무 적용 팁

• 작은 모델 → DDP로 충분
• 1B+ 모델 → ZeRO-2 이상 고려
• 10B+ 모델 → ZeRO-3/FSDP 필수
• GPU 메모리 + 통신 속도 모니터링 필수 (병목 방지)

---

결론

FSDP / ZeRO / DeepSpeed는 대형 모델 학습의 필수 도구입니다.
모델 파라미터를 GPU 메모리 안에 다 못 올리는 상황에서도 학습이 가능하고,
더 적은 GPU로도 동일한 batch size를 유지할 수 있습니다.

 

 

 

다음 시리즈는 아래에 있습니다!

https://machineindeep.tistory.com/76

추론 속도 최적화 시리즈 1편: Operator Fusion (연산자 융합)