분산 학습의 핵심 개념

분산 학습은 여러 GPU를 활용해 모델 학습 속도를 향상시키는 기술입니다. 여기서 중요한 개념은 다음과 같습니다:

• 병렬 처리 (Parallel): 여러 장치에서 동시에 작업 수행, 프로세스 수에 따라 단일 또는 다중 프로세스 구조 가능
• 분산 환경 (Distributed): 각 GPU마다 별도의 프로세스가 실행되며, 통신을 통해 상태를 동기화하는 방식

본 문서에서는 '병렬'과 '분산'을 혼용하여 사용하며, 다중 GPU 기반의 학습을 의미합니다.

주요 분산 학습 방식

PyTorch에서는 주로 다음 세 가지 방식이 사용됩니다:

• DataParallel (DP): 단일 프로세스 내에서 여러 GPU를 활용하지만, 스케일링에 한계가 있음
• DistributedDataParallel (DDP): 각 GPU에 독립적인 프로세스를 생성하고, 그 사이에서 그래디언트 동기화를 통해 효율적인 분산 학습 지원
• FSDP (Fully Sharded Data Parallel): PyTorch 1.11 이후 도입된 고급 분산 전략으로, 메모리 사용을 최적화하며, DDP와 ZeRO 기법을 포함

또한, Microsoft DeepSpeed은 ZeRO-1/2/3를 통해 더 강력한 메모리 절약 기능을 제공합니다. 이 중 ZeRO-0는 기본적으로 DDP와 동일한 동작을 합니다.

학습 설정 방법

가장 간편한 접근은 Hugging Face Transformers 라이브러리의 Trainer 클래스를 사용하는 것입니다. 이 클래스는 자동으로 데이터 분할 및 분산 환경 초기화를 처리합니다.

예시: 기존 단일 GPU 실행 스크립트

python train_model.py --output_dir ./results

다중 GPU 환경으로 전환 시, 다음과 같이 변경합니다:

deepspeed --include "localhost:1,2,3,4" train_model.py \
  --deepspeed ds_config.json \
  --output_dir ./results

여기서 중요한 점은 CUDA_VISIBLE_DEVICES 환경 변수를 설정하면 deepspeed가 정상 작동하지 않을 수 있다는 것입니다. 반드시 제거해야 합니다.

torchrun을 통한 분산 실행

PyTorch 공식 도구인 torchrun을 사용하면 직접 분산 프로세스를 관리할 수 있습니다.

CUDA_VISIBLE_DEVICES=1,2,3,4 torchrun --nproc_per_node 4 run_clm.py \
  --deepspeed ds_config.json \
  --output_dir ./results

기본적으로 마스터 포트는 29500이며, 충돌 시 다른 포트로 지정 가능합니다:

torchrun --master_port 61234 --nproc_per_node 4 train.py ...

쉘 스크립트 예시:

export CUDA_VISIBLE_DEVICES=$1
gpu_num=$(echo $CUDA_VISIBLE_DEVICES | awk -F ',' '{print NF}')

torchrun --master_port 61234 --nproc_per_node $gpu_num hf_train.py ...

참고 사항

PyTorch 공식 문서의 Parallel and Distributed Training Tutorial와 DDP, FSDP 관련 개발자 노트는 깊이 있는 이해에 매우 유용합니다.