DeepSpeed 완벽 이해하기!

1. DeepSpeed란?

DeepSpeed는 Microsoft에서 개발한 대규모 분산 학습 최적화 라이브러리이다. 기존의 PyTorch DDP만으로는 학습하기 어려운 수십억~수백억 파라미터 모델을 효율적으로 학습시키기 위해 등장했다. GPU 메모리를 절약하고 학습 속도를 높이기 위한 ZeRO 최적화, Mixed Precision, Offloading, Pipeline/Tensor 병렬화 등 다양한 기법을 통합적으로 제공하는 것이 특징이다. PyTorch 기반으로 작동하며 사용법과 통합이 용이하다.

 

2. 주요 특징

2.1 ZeRO (Zero Redundancy Optimizer)

DeepSpeed의 핵심 기능이다. 파라미터, 그라디언트, 옵티마이저 상태를 GPU끼리 분산(shard)해 GPU 메모리 사용량을 대폭 절감한다. 예를 들어 ZeRO-3는 파라미터 자체까지 shard하여 GPU 하나가 모델 전체를 항상 올리지 않고도 학습이 가능하다. PyTorch DDP는 모델 전체를 각 GPU에 복제하는 방식이라 메모리 사용량이 높지만, ZeRO는 이를 최소화한다.

• ZeRO-1: optimizer states 분산
• ZeRO-2: optimizer states + gradients 분산
• ZeRO-3: optimizer states + gradients + parameters 까지 모두 shard

2.2 Mixed Precision Training

FP16, BF16 등을 활용해 연산 속도와 GPU 메모리 사용량을 줄인다. 예를 들어 32GB GPU로 16GB 모델을 학습할 때 FP16을 적용하면 같은 GPU에서 2배 가까운 배치 사이즈를 사용할 수 있다.

2.3 Gradient Accumulation & CPU Offloading

batch를 여러 step으로 나눠 그래디언트를 누적(accumulate)하고, optimizer 상태를 CPU에 올려 GPU 메모리 부담을 줄인다. DDP나 FSDP보다 CPU offload가 쉽게 적용된다.

2.4 Pipeline & Tensor Parallelism

모델을 여러 stage나 텐서 단위로 쪼개 다수의 GPU에서 병렬로 학습한다. Transformer처럼 깊은 네트워크를 파이프라인 스테이지로 나눠 처리하거나 텐서 연산 자체를 쪼개어 분산 연산하는 방식으로, PyTorch의 기본 DDP/FSDP가 지원하지 않는 수준의 fine-grained 병렬 학습이 가능하다.

 

2.5 분산학습 프레임워크 비교

방법	무엇을 병렬화?	특징
DDP	데이터 병렬화	모든 GPU가 같은 모델 전체를 복제, 서로 다른 데이터 batch 학습
FSDP	데이터 + 파라미터 shard	파라미터 shard로 GPU 메모리 줄임, 그래도 데이터 병렬
DeepSpeed Pipeline	모델 자체를 stage별로 분할	GPU에 모델을 나눠서 올려 메모리 분산, forward/backward도 순차 수행

 

DDP, FSDP가 데이터 병렬을 위한 프레임워크라면, DeepSpeed는 ZeRO(데이터+파라미터 병렬) 뿐 아니라 파이프라인 병렬도 지원해 초대형 모델 학습에 훨씬 적합하다.

 

3. DeepSpeed Pipeline Parallelism 실제 동작 예시

예시: 3-stage 파이프라인

• 모델을 stage 3개로 나누었다고 가정 (ex. Transformer block을 나눔)
• GPU도 3대 (GPU0, GPU1, GPU2)

미니배치1
GPU0: stage1 --> output1
GPU1:           stage2 --> output2
GPU2:                      stage3 --> final_output

미니배치2
GPU0: stage1 --> output1
GPU1:           stage2 --> output2
GPU2:                      stage3 --> final_output

...

즉 forward는

• GPU0이 stage1을 계산 → output을 GPU1로 넘김
• GPU1이 stage2 계산 → output을 GPU2로 넘김
• GPU2가 stage3 계산 → 최종 output

동시에 pipeline bubble을 줄이기 위해

time step1: GPU0(미니배치1 stage1)
time step2: GPU0(미니배치2 stage1), GPU1(미니배치1 stage2)
time step3: GPU0(미니배치3 stage1), GPU1(미니배치2 stage2), GPU2(미니배치1 stage3)
...

이렇게 미니배치를 나누어 서로 다른 stage가 동시에 돌아가도록 해 idle time(bubble)을 최소화한다.

 

Backward 흐름

GPU2: backward(stage3)
GPU1: backward(stage2)
GPU0: backward(stage1)

backward는 forward의 반대 방향으로 흘러감.

 

✅ Pipeline parallelism
모델을 여러 stage로 쪼개 GPU에 분산 → 순차적으로 forward/backward 수행 → bubble 최소화 위해 미니배치 나누어 pipeline 채우기

 

4. 사용 방법

4.1 설치

pip install deepspeed

PyTorch, CUDA, NCCL이 설치되어 있어야 하며, 대규모 분산을 위해 Slurm 같은 스케줄러 환경이 준비되어 있으면 좋다.

4.1 기본 학습 스크립트 구조

DeepSpeed는 deepspeed 런처를 사용해 분산 학습을 시작한다.

deepspeed --num_gpus=4 train.py --deepspeed --deepspeed_config ds_config.json

여기서 ds_config.json은 학습 및 ZeRO, FP16 설정 등을 담은 설정 파일이다.

4.2 예시 ds_config.json

{
  "train_batch_size": 64,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu"
    }
  }
}

• train_batch_size: 전체 학습에서 목표로 하는 effective batch size. gradient_accumulation_steps와 num_gpus를 곱해 GPU당 처리량을 결정한다.
• gradient_accumulation_steps: gradient를 여러 step 동안 누적 후 한꺼번에 optimizer step을 수행해 더 큰 batch size를 흉내낸다.
• fp16.enabled: float32 대신 float16(반정밀도) 연산을 수행해 연산 속도와 GPU 메모리 사용을 줄인다.
• zero_optimization.stage: ZeRO 최적화 단계. stage 2는 optimizer states와 gradients를 shard하여 메모리 footprint를 크게 줄인다.
• offload_optimizer.device: optimizer 상태를 CPU로 옮겨 GPU 메모리 부담을 더 줄인다.

GPU 클러스터 구성 → Slurm으로 자원 관리 → PyTorch 모델 → DeepSpeed로 래핑 + ds_config → Slurm에서 srun/sbatch로 실행

---

DeepSpeed는 초거대 모델을 학습하기 위해 등장한 프레임워크로, ZeRO를 통해 GPU 메모리 부담을 극적으로 줄이고, FP16, Offloading, Parallelism 등을 조합해 대규모 분산 학습을 효율적으로 수행할 수 있게 한다.