DeepSpeed 시작하기

DeepSpeed 시작하기 (Getting Started)

DeepSpeed 모델 학습은 DeepSpeed 엔진 위에서 이루어져요. 엔진이 torch.nn.Module 타입이라면 어떤 모델이든 감쌀 수 있고, 학습과 체크포인트에 필요한 API가 아주 적어요. 설치부터 시작해서 실제로 분산 학습을 띄우는 흐름을 보여드릴게요.

출처: DeepSpeed Getting Started (공식)

설치하기

DeepSpeed는 일반적으로 pip으로 설치해요. GPU 가속이 필요하면 CUDA 버전에 맞춰 설치해야 하고, AMD 환경에서는 ROCm 이미지를 쓸 수 있어요.

pip install deepspeed

AMD 환경이라면 제공되는 ROCm 이미지로 바로 시작할 수 있어요.

docker pull deepspeed/rocm501:ds060_pytorch110

엔진 초기화하기

학습 코드는 deepspeed.initialize를 호출해 모델을 DeepSpeed 엔진으로 감싸는 것으로 시작해요. 파라미터를 넘겨주면 엔진이 분산 데이터 병렬 학습과 혼합 정밀도, 학습률 스케줄러를 알아서 구성해 줘요.

import deepspeed

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

args에는 --deepspeed_config로 지정한 설정 파일 경로가 포함돼요. 설정 파일에 ZeRO 단계나 통신 전략을 적어 두면, 엔진이 그에 맞춰 동작해요.

설정 파일 구성하기

DeepSpeed의 기능은 대부분 설정 JSON으로 켜고 끌 수 있어요. 예를 들어 ZeRO 단계를 고르거나, 오프로딩을 켜거나, FP16/BF16을 활성화하는 것 모두 설정 파일에서 처리돼요.

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

--hostfile은 참여할 노드와 GPU를 나열한 파일이에요. 여러 노드에 걸쳐 학습할 때 이 파일로 자원 구성을 결정해요.

no-SSH 런칭 모드

클라우드나 Kubernetes처럼 컨테이너 오케스트레이션이 자유로운 환경에서는 패스워드 없는 SSH 설정이 번거로울 수 있어요. 그럴 때 --no_ssh 옵션을 쓰면 SSH 없이도 학습을 띄울 수 있어요.

deepspeed --hostfile=myhostfile --no_ssh --node_rank=<n> \
  --master_addr=<addr> --master_port=<port> \
  client_entry.py client_args \
  --deepspeed --deepspeed_config ds_config.json

--hostfile=myhostfile: 노드와 GPU 정보가 든 hostfile을 지정해요.
--no_ssh: SSH 없는 모드를 켜요.
--node_rank=<n>: 노드의 rank를 지정해요. 0부터 n-1 사이의 유일한 정수여야 해요.
--master_addr=<addr>: 리더 노드(rank 0)의 주소예요.
--master_port=<port>: 리더 노드의 포트예요.

이 모드는 PyTorch의 torchrun 런처와 비슷하게 동작해요.

멀티 노드 환경변수 전파

여러 노드에서 학습할 때 사용자 정의 환경변수를 전달하고 싶을 수 있어요. DeepSpeed는 기본적으로 모든 NCCL·PYTHON 관련 환경변수를 전파해요. 추가로 전파할 변수는 현재 실행 경로나 홈 디렉토리의 .deepspeed_env 파일에 한 줄에 VAR=VAL 형식으로 적으면 돼요.

NCCL_IB_DISABLE=1
NCCL_SOCKET_IFNAME=eth0

예를 들어 일부 클러스터는 학습 전에 특별한 NCCL 변수를 설정해야 하는데, 그런 변수를 홈 디렉토리의 .deepspeed_env에 적어 두면 매번 명령에 포함할 필요가 없어요.

DeepSpeed 시작하기