최근 거대 언어 모델(LLM)의 폭발적인 성장과 함께, 단일 GPU의 메모리 한계를 극복하기 위한 분산 학습 기술은 선택이 아닌 필수가 되었습니다. Microsoft에서 개발한 DeepSpeed는 그 중심에 있으며, 특히 ZeRO (Zero Redundancy Optimizer) 알고리즘은 메모리 효율성을 극대화하여 기존 데이터 병렬 처리의 한계를 뛰어넘습니다. 본 가이드에서는 현업 엔지니어가 직면하는 메모리 부족(OOM) 문제를 해결하고, 효율적인 분산 학습 환경을 구축하기 위한 ZeRO의 단계별 설정 방법과 실무 최적화 전략을 심도 있게 다룹니다.
1. ZeRO(Zero Redundancy Optimizer)의 핵심 개념과 필요성
일반적인 데이터 병렬 처리(Data Parallelism) 방식에서는 모든 GPU가 동일한 모델 파라미터, 그래디언트, 옵티마이저 상태를 복제하여 가집니다. 이는 모델 규모가 커질수록 GPU 메모리 점유율을 기하급수적으로 높이는 원인이 됩니다. ZeRO는 이러한 중복을 제거하여 메모리를 절약하면서도 계산 효율을 유지합니다.
ZeRO의 3단계 최적화 비교
| 단계 (Stage) | 최적화 대상 | 메모리 절감 효과 | 특징 및 권장 상황 |
|---|---|---|---|
| ZeRO-1 | Optimizer States (옵티마이저 상태) | 약 4배 절감 | 안정적이며 통신 오버헤드가 적음. 기본 분산 학습에 권장. |
| ZeRO-2 | Optimizer States + Gradients (그래디언트) | 약 8배 절감 | 대부분의 중대형 모델 학습에 가장 범용적으로 사용됨. |
| ZeRO-3 | Optimizer States + Gradients + Parameters (파라미터) | GPU 개수에 비례하여 절감 | 단일 GPU에 담을 수 없는 초거대 모델(10B+) 학습 시 필수. |
2. 실무 적용을 위한 DeepSpeed ZeRO 구성 핵심 전략
DeepSpeed를 설정할 때 가장 중요한 것은 ds_config.json 파일의 세밀한 튜닝입니다. 특히 하드웨어 자원(VRAM, 대역폭)에 따라 각 스테이지의 세부 옵션을 조정해야 학습 속도 저하를 최소화할 수 있습니다.
Offload 기능을 활용한 메모리 확장
ZeRO-2와 ZeRO-3에서는 GPU 메모리가 부족할 경우 CPU 메모리나 NVMe 스토리지로 데이터를 넘기는 Offload 기능을 제공합니다. 이는 연산 속도는 다소 느려지지만, 물리적인 GPU 한계를 넘어서는 모델을 학습할 수 있게 합니다.
3. 개발자를 위한 실전 DeepSpeed 설정 및 구현 Example (7가지)
다음은 실제 연구 및 서비스 개발 현장에서 즉시 적용 가능한 설정 예시입니다. 각 예시는 특정 목적에 최적화되어 있습니다.
Example 1: ZeRO-1 기본 설정 (안정적인 분산 학습)
{
"train_batch_size": 32,
"steps_per_print": 10,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.0001,
"betas": [0.9, 0.999],
"eps": 1e-8,
"weight_decay": 0.01
}
},
"zero_optimization": {
"stage": 1,
"allgather_partitions": true,
"reduce_scatter": true,
"allgather_bucket_size": 2e8,
"overlap_comm": true
},
"fp16": {
"enabled": true
}
}
Example 2: ZeRO-2 고성능 설정 (학습 속도 최적화)
그래디언트 축적과 버킷 크기를 최적화하여 연산 효율을 높인 설정입니다.
{
"train_batch_size": 64,
"gradient_accumulation_steps": 2,
"zero_optimization": {
"stage": 2,
"allgather_partitions": true,
"reduce_scatter": true,
"allgather_bucket_size": 5e8,
"overlap_comm": true,
"reduce_bucket_size": 5e8
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16
}
}
Example 3: ZeRO-2 CPU Offload (VRAM 부족 해결)
GPU 메모리가 24GB 이하인 환경에서 중형 모델을 학습할 때 유용합니다.
{
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
}
},
"fp16": { "enabled": true }
}
Example 4: ZeRO-3 초거대 모델 설정 (Parameter Partitioning)
수십억 개의 파라미터를 가진 모델을 여러 GPU에 나누어 로드합니다.
{
"zero_optimization": {
"stage": 3,
"offload_optimizer": { "device": "cpu" },
"offload_param": { "device": "cpu" },
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto"
}
}
Example 5: 파이썬 코드 내 모델 초기화 루틴
DeepSpeed 엔진을 초기화하는 표준 방식입니다.
import deepspeed
import torch
# 모델 및 옵티마이저 정의
model = MyLargeModel()
engine, optimizer, _, _ = deepspeed.initialize(
args=args,
model=model,
model_parameters=model.parameters(),
config="ds_config.json"
)
for data, target in dataloader:
data = data.to(engine.device)
target = target.to(engine.device)
outputs = engine(data)
loss = criterion(outputs, target)
engine.backward(loss)
engine.step()
Example 6: Hugging Face Trainer 통합 설정
Transformers 라이브러리와 함께 사용할 때의 설정입니다.
# training_args 설정 시 deepspeed 경로 지정
training_args = TrainingArguments(
output_dir="./output",
per_device_train_batch_size=4,
deepspeed="./ds_config_stage3.json", # 설정 파일 경로
bf16=True # Ampere 아키텍처 이상 권장
)
Example 7: 다중 노드(Multi-node) 분산 학습 실행 스크립트
여러 대의 서버를 연결하여 학습할 때 사용하는 쉘 명령어 예시입니다.
# hostfile에는 각 서버의 IP와 GPU 개수가 기재되어 있어야 함
deepspeed --hostfile=./my_hostfile \
--num_nodes=2 \
--num_gpus=8 \
train.py --deepspeed_config ds_config.json
4. 성능 향상을 위한 해결 방법 및 팁
- 통신 대역폭 확인: ZeRO-3는 노드 간 통신량이 많으므로 NVLink나 100G 이상의 InfiniBand 환경이 구축되어 있는지 확인하십시오.
- Gradient Clipping: 분산 학습 시 그래디언트 폭주를 막기 위해
"gradient_clipping": 1.0설정을 추가하는 것이 좋습니다. - Mixed Precision: FP16 보다는 하드웨어가 지원한다면(A100 이상) BF16을 사용하십시오. 수치적 안정성이 훨씬 뛰어납니다.
5. 결론 및 향후 전망
DeepSpeed ZeRO는 복잡한 병렬 처리 로직을 추상화하여 개발자가 모델 구조에 집중할 수 있게 해줍니다. 1단계에서 3단계로 갈수록 메모리 절감량은 커지지만 통신 비용이 발생하므로, 프로젝트의 모델 크기와 하드웨어 사양에 맞는 최적의 단계를 선택하는 것이 성공적인 AI 모델 개발의 핵심입니다.
참고 문헌 및 출처
- Microsoft Research: "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" (2020)
- DeepSpeed Official Documentation
- Hugging Face Documentation: "Efficient Training on Multiple GPUs"
- GitHub Repository: microsoft/DeepSpeed
