[Pytorch] 분산 학습 4: DeepSpeed

Date: 2025.10.16 Updated: 2025.10.16

카테고리: Pytorch

DeepSpeed

DeepSpeed 란?

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

주요 기능
- 모델 분할: 모델 분할을 통해 단일 GPU 메모리 한계를 넘어서는 매우 큰 모델을 학습할 수 있도록 지원
- 효율적인 데이터 병렬 처리: GPU에서의 데이터 병렬 처리를 최적화하여, 학습 속도를 개선
- ZeRO 최적화: ZeRO(Zero Redundancy Optimizer) 기술을 사용해 메모리 사용량을 줄이면서 대규모 모델의 학습을 가능하게 함.
- 모델 학습 가속: 커널 최적화와 효율적인 커뮤니케이션을 통해 모델 학습 속도를 높임.

ZeRO (Zero Redundancy Optimizer)

DeepSpeed의 핵심 기능 중 하나는 ZeRO이다. 파라미터, 그래디언트, 옵티마이저 상태를 GPU끼리 분산해 GPU 메모리 사용량을 크게 절감시킨다. 특히 ZeRO-3는 DDP와 달리 파라미터 자체까지 shard하여 GPU 하나가 모델 전체를 항상 올리지 않고도 학습을 가능하게 한다.

ZeRO-1: Optimizer states 분산
ZeRO-2: Optimizer states + Gradients 분산
ZeRO-3: Optimizer states + Gradients + Parameters 분산

Mixed Precision Training

DeepSpeed를 사용하면서 학습 정밀도도 마찬가지로 조절 가능하다. 예를 들어, FB16, BF16등을 활용해 연산 속도와 GPU 메모리 사용량을 줄일 수 있다. 참고로 사용하는 GPU 종류에 따라 최적화된 정밀도는 다르다.

A6000, A100: BF16
H100: FP8 > BF16 (가능하면 Transformer Engine 기반 FP8이 효율적)

Gradient Accumulation & CPU Offloading

대규모 배치를 GPU 메모리 한계 내에서 구현하기 위해, DeepSpeed는 그래디언트를 여러 스텝에 걸쳐 누적(accumulation)하는 방식을 사용한다. 이렇게 하면 실제 GPU에 올리는 마이크로 배치 크기는 작지만, 여러 스텝을 모아 하나의 대형 배치로 학습한 것과 동일한 효과를 얻을 수 있다.

또한 옵티마이저 상태(optimizer states)를 GPU 대신 CPU나 NVMe에 저장(offload)함으로써 VRAM 사용량을 크게 줄인다. 이는 ZeRO 최적화와 결합되어 GPU에 필요한 텐서만 로드하고, 나머지는 CPU에서 비동기 전송으로 처리하므로 DDP·FSDP보다 오프로딩 구성이 간단하다. 다만 CPU 전송 대기 시간(PCIe latency)이 생길 수 있어, BF16 연산이나 통신-연산 오버랩(overlap)으로 이를 완화한다.

결과적으로 이 접근은 메모리 효율성을 극대화하면서도 안정적인 대규모 학습을 가능하게 하며, 제한된 GPU 환경에서도 긴 시퀀스나 대형 모델을 다룰 수 있게 한다. 단, 그래디언트 누적 단계가 많을수록 한 스텝의 학습 시간이 길어질 수 있으므로, 포스터에서는 “메모리 절감 ↔ 연산 지연의 트레이드오프”를 강조하면 된다.

Pipeline & Tensor Parallelism

DeepSpeed는 모델을 여러 GPU로 세밀하게 쪼개 병렬 처리할 수 있도록 파이프라인(PP)과 텐서 병렬(TP) 방식을 지원한다. 파이프라인 병렬은 트랜스포머처럼 깊은 모델을 여러 스테이지로 나눠 마이크로배치를 순차적으로 흘려보내며, 1F1B 스케줄로 각 GPU가 동시에 다른 배치를 처리하도록 해 GPU 유휴 시간을 최소화한다.

텐서 병렬은 각 레이어의 연산(예: Linear, Attention)을 행렬 단위로 분할해 여러 GPU가 동시에 연산한 뒤 all-reduce로 결과를 합치는 방식이다. 이는 모델의 “너비(width)” 방향으로 병렬화된 형태이며, 특히 거대한 피드포워드나 어텐션 블록에서 큰 효과를 발휘한다.

두 병렬화를 함께 사용하면 깊이(Depth) + 너비(Width) 양축을 동시에 분산시킬 수 있어, 단일 GPU로는 불가능한 규모의 모델 학습이 가능해진다. 다만 통신 오버헤드와 파이프라인 버블이 성능을 저해할 수 있으므로, 마이크로배치 수와 스테이지 분할 균형을 조정해 통신–연산 오버랩을 극대화하는 것이 핵심이다.

DDP vs FSDP vs DeepSpeed

DeepSpeed는 DDP나 FSDP보다 대규모 LLM 학습에 최적화된 통합 분산 엔진이다. DDP가 단순한 데이터 병렬화에 머무르고, FSDP가 파라미터 샤딩 중심의 메모리 절감에 초점을 맞추는 반면, DeepSpeed는 ZeRO 최적화, CPU/NVMe Offloading, Gradient Accumulation, Mixed Precision, Pipeline·Tensor Parallelism을 하나의 프레임워크 안에서 유기적으로 결합한다. 이로써 GPU 메모리 용량에 구애받지 않고 거대 모델을 효율적으로 학습할 수 있으며, 통신–연산 오버랩과 비동기 I/O를 통해 확장성과 처리 효율을 동시에 확보한다. 따라서 DeepSpeed는 단순 병렬화 라이브러리가 아니라, 모델 규모·메모리·통신·학습속도 간의 균형을 자동으로 조정하는 LLM 전용 학습 인프라로서 DDP·FSDP보다 훨씬 유연하고 실용적인 선택이다.

DeepSpeed 실행

설치와 환경

# Install
pip install deepspeed

# Quick env report (CUDA/NCCL/BF16/TPU 여부 등)
ds_report

실행 방법

# ds_config.json
{
  "bf16": { "enabled": true },                // use BF16 if supported
  "train_batch_size": 0,                      // computed by DeepSpeed
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 2,                               // 3 for full param sharding
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "reduce_scatter": true,
    "overlap_comm": true
  }
}

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
model_id = "gpt2"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
args = TrainingArguments(
    output_dir="out",
    per_device_train_batch_size=1,            # micro-batch
    deepspeed="ds_config.json",               # hand over to DeepSpeed
    num_train_epochs=1
)
trainer = Trainer(model=model, args=args, train_dataset=[{"input_ids":[0]}])
trainer.train()

# 4 GPUs example
deepspeed --num_gpus 4 train.py --deepspeed ds_config.json

실행 방법 B. Python API로 직접 초기화

import torch
import deepspeed as ds
model = torch.nn.Linear(768, 768)
optim = torch.optim.AdamW(model.parameters(), lr=1e-4)
engine, optim, _, _ = ds.initialize(model=model, optimizer=optim, config="ds_config.json")

for _ in range(10):
    x = torch.randn(8, 768, device=engine.device)
    y = model(x).sum()
    engine.backward(y)
    engine.step()

Reference

Blog: DeepSpeed 완벽 이해하기!
BLog: 모델 학습 with Huggingface (DeepSpeed)
Huggingface: ”DeepSpeed”

Meaningful