DeepSpeed

DeepSpeed (LLM 학습·추론 가속)

DeepSpeed는 Microsoft가 만든 딥러닝 시스템 최적화 라이브러리예요. 모델이 점점 커지면서 메모리에 통째로 올리지 못하는 상황이 오잖아요? 그 문제를 ZeRO 같은 기법으로 풀어서, 수백억 파라미터 모델도 학습하고 추론할 수 있게 도와줘요. PyTorch 모델을 그대로 감싸서 쓰기 때문에, 기존 학습 코드를 크게 바꾸지 않고도 분산 학습과 메모리 절약을 누릴 수 있어요.

DeepSpeed의 정체성은 크게 두 가지예요. 하나는 학습을 위한 ZeRO 계열 최적화이고, 다른 하나는 추론을 위한 DeepSpeed-Inference예요. 둘 다 하나의 엔진이 담당해서, 학습했던 체크포인트를 그대로 추론에 쓸 수 있다는 점이 큰 장점이에요.

출처: DeepSpeed 공식 사이트

DeepSpeed는 왜 등장했을까요

큰 모델을 학습하려면 몇 가지 벽이 있어요. GPU 메모리에 모델이 안 들어가는 문제, 여러 GPU에 나눠 담았을 때 통신 비용이 커지는 문제, 그리고 배치 크기를 키울 때 학습 수렴이 불안정해지는 문제까지요. DeepSpeed는 이 벽들을 시스템 레벨에서 부수는 데 집중해요.

대표적인 혁신으로는 ZeRO, 3D-Parallelism, DeepSpeed-MoE, ZeRO-Infinity 같은 것들이 있어요. 특히 ZeRO는 옵티마이저 상태·그래디언트·파라미터를 GPU들에 나눠 담아서 중복을 없애는 방식이라, 같은 하드웨어로 훨씬 큰 모델을 학습하게 해줘요.

엔진으로 모델 감싸기

DeepSpeed 학습은 deepspeed 엔진이 중심이에요. 엔진은 torch.nn.Module 타입이면 어떤 모델이든 감쌀 수 있고, 학습과 체크포인트 저장에 필요한 API를 아주 작게 유지해요.

import deepspeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=params
)

내부적으로는 분산 데이터 병렬 학습, 혼합 정밀도, 사전 정의된 학습률 스케줄러를 자동으로 처리해 줘요. 그래서 직접 DataParallel이나 DistributedDataParallel을 붙일 필요가 줄어들죠.

설정 파일로 기능 켜고 끄기

DeepSpeed의 강점 중 하나는 기능 대부분을 코드가 아니라 설정 JSON으로 다루는 점이에요. 런처에 ds_config.json을 넘기면, ZeRO 단계·오프로딩·통신 방식 등을 코드 수정 없이 바꿀 수 있어요.

deepspeed --hostfile=myhostfile client_entry.py client_args \
  --deepspeed --deepspeed_config ds_config.json

--hostfile은 노드와 GPU 정보가 담긴 파일을 가리키고, ds_config.json이 실제 최적화 옵션을 담아요. 멀티 노드 환경이라면 .deepspeed_env 파일에 VAR=VAL 형식으로 환경변수를 적어 두면 런처가 알아서 전파해 줘요.

NCCL_IB_DISABLE=1
NCCL_SOCKET_IFNAME=eth0

NCCL 관련 특수 환경변수가 필요한 클러스터라면 이 파일이 특히 유용해요. 기본적으로 NCCL·PYTHON 관련 환경변수는 자동으로 전파되니까, 그 외 추가 변수만 명시하면 돼요.

3D 병렬화로 스케일 확장

모델이 수백억 파라미터 수준으로 커지면 데이터 병렬만으로는 부족해요. 그때 텐서 병렬·파이프라인 병렬·데이터 병렬을 함께 쓰는 3D-Parallelism이 등장해요. DeepSpeed는 이 세 축을 한 프레임워크 안에서 조율하기 때문에, 학습 규모를 키우면서도 통신 오버헤드를 관리할 수 있어요.

더 알아보기