Hugging Face Accelerate 라이브러리가 PyTorch FSDP를 통합하여 GPT-2 Large 모델에서 배치 크기를 DDP 대비 2~3배 증대 및 GPT-2 XL 모델 학습 가능

Context

머신러닝 모델의 파라미터 규모가 증가하면서 단일 GPU 메모리로 대규모 모델을 로드하거나 학습하기 어려워졌다. DDP(Distributed Data Parallel)를 사용하면 GPT-2 Large(762M)에서 배치 크기 7로 제한되고, GPT-2 XL(1.5B)은 배치 크기 1에서도 OOM(Out Of Memory) 오류가 발생한다.

Technical Solution

• 데이터 병렬화 기법 ZeRO를 PyTorch에 네이티브로 구현한 FSDP 도입: 옵티마이저 상태, 그래디언트, 모델 파라미터를 워커 간에 샤딩
• FSDP 샤딩 전략 선택: SHARD_GRAD_OP(Stage 2)와 FULL_SHARD(Stage 3) 구성으로 메모리 사용량 최적화
• min_num_params 파라미터 조정: 1M 또는 2K로 설정하여 샤딩되는 모듈 단위 제어
• CPU 오프로드 활성화: 그래디언트와 옵티마이저 상태를 CPU로 오프로드하여 GPU 메모리 추가 확보
• Accelerate 라이브러리 통합: 코드 변경 없이 config 파일만으로 FSDP 설정 가능

Impact

• GPT-2 Large 모델 배치 크기: DDP 7 → FSDP FULL_SHARD 15 (2.1배 증가), CPU 오프로드 시 22 (3.1배 증가)
• GPT-2 XL 모델: DDP에서 OOM으로 학습 불가 → FSDP로 배치 크기 8 이상 학습 가능
• 학습 시간(GPT-2 Large): DDP 15분 → FSDP SHARD_GRAD_OP 11분 (26% 단축), CPU 오프로드 시 23~24분

Key Takeaway

DDP의 메모리 제약으로 인한 배치 크기 제한을 겪는 대규모 모델 학습 환경에서 FSDP 기반 메모리 샤딩은 코드 변경 없이 배치 크기를 2~3배 확대하고 더 큰 모델 학습을 가능하게 한다.