Hugging Face가 ZeRO 메모리 최적화를 Transformers Trainer에 통합해 배치 크기 16에서 50으로 3배 증가 및 학습 시간 30.9초에서 20.9초로 32% 단축

Context

GPU 메모리 증가 속도보다 머신러닝 모델이 더 빠르게 커지면서 대규모 모델을 학습하거나 로드할 수 없는 상황이 발생하고 있다. 단일 GPU로는 수십억 개 매개변수 모델을 학습할 수 없으며, 다중 GPU 환경에서도 메모리 효율이 낮아 배치 크기와 학습 속도에 제약이 있다.

Technical Solution

• ZeRO 메모리 최적화를 Transformers Trainer에 통합: --sharded_ddp (FairScale), --deepspeed 커맨드라인 인자로 모델 매개변수 및 최적화 상태를 GPU들 간에 샤딩
• FairScale sharded_ddp 옵션 추가: 기존 DistributedDataParallel(DDP) 대신 모델 가중치와 그래디언트를 여러 GPU에 분산 저장
• DeepSpeed CPU 오프로딩 옵션 제공: GPU 메모리 부족 시 CPU 메모리로 중간 계산 결과와 최적화 상태를 임시 저장하고 필요시 로드
• 혼합 정밀도(FP16) 지원: FP16을 추가로 적용하면 메모리 사용량 추가 감소 가능
• 구성 파일 기반 배포: DeepSpeed는 JSON 설정 파일로 최적화 정책 제어, FairScale은 커맨드라인 인자만으로 활성화

Impact

T5-Large 모델 기준 2개 24GB GPU 환경에서:

• 기본 설정: 배치 크기 16, 학습 시간 30.95초, 평가 시간 56.33초
• FP16 추가: 배치 크기 20, 학습 시간 21.49초 (30% 단축)
• FairScale Sharded DDP: 배치 크기 30, 학습 시간 25.91초, 평가 시간 47.56초
• FairScale + FP16: 배치 크기 30, 학습 시간 17.38초 (44% 단축), 평가 시간 45.66초 (19% 단축)
• DeepSpeed CPU 오프로딩 없음: 배치 크기 40, 학습 시간 10.40초 (66% 단축), 평가 시간 34.93초 (38% 단축)
• DeepSpeed CPU 오프로딩 포함: 배치 크기 50, 학습 시간 20.97초, 평가 시간 32.14초 (43% 단축)

Key Takeaway

ZeRO 메모리 최적화를 Trainer 수준에서 통합하면 모델 코드 수정 없이 커맨드라인 인자만으로 GPU 메모리 효율을 2~3배 향상시킬 수 있으며, DeepSpeed가 더 강력하지만 FairScale은 배포 복잡도가 낮은 대안을 제공한다.