DeepSpeed를 사용한 대규모 모델 학습 가속

Hugging Face는 DeepSpeed의 Zero Redundancy Optimizer (ZeRO)를 accelerate 라이브러리에 통합하여, GPU 간 메모리 중복을 줄임으로써 대규모 모델을 보다 효율적으로 학습할 수 있게 했습니다. 이 통합을 통해 개발자는 메모리 부족(OOM) 오류를 피하고, 큰 코드 변경 없이 배치 크기를 늘릴 수 있습니다.

ZeRO 데이터 병렬 처리 이해

ZeRO(Zero Redundancy Optimizer)는 데이터 병렬 처리 중 메모리 사용량을 최적화하기 위해 학습 상태를 사용 가능한 GPU에 걸쳐 샤딩합니다. 프레임워크는 모델 학습의 메모리 발자국을 줄이기 위해 여러 단계로 동작합니다:

• Stage 1: 옵티마이저 상태를 데이터 병렬 워커/GPU에 걸쳐 샤딩합니다.
• Stage 2: 옵티마이저 상태와 그래디언트를 워커/GPU에 걸쳐 샤딩합니다.
• Stage 3: 옵티마이저 상태, 그래디언트, 모델 파라미터를 워커/GPU에 걸쳐 샤딩합니다.
• Optimizer Offload: Stage 2를 확장하여 그래디언트와 옵티마이저 상태를 CPU 또는 디스크로 오프로드합니다.
• Param Offload: Stage 3을 확장하여 모델 파라미터를 CPU 또는 디스크로 오프로드합니다.

Accelerate를 통한 코드 없는 통합

사용자는 accelerate config 명령과 Accelerate DeepSpeed 플러그인을 사용하여 훈련 코드를 수정하지 않고 DeepSpeed ZeRO Stage-2를 활용할 수 있습니다.

두 개의 24GB NVIDIA Titan RTX GPU를 사용해 DeBERTa-v2-xlarge-mnli 모델(9억 파라미터)을 파인튜닝한 벤치마크에서 DeepSpeed ZeRO Stage-2는 Distributed Data Parallel(DDP)보다 크게 우수한 성능을 보였습니다:

방법	최대 배치 크기	에포크당 학습 시간 (초)	에포크당 평가 시간 (초)	F1 점수	정확도
DDP	8	103.57	2.04	0.931	0.904
DeepSpeed ZeRO Stage 2	40	28.98	1.79	0.936	0.912

DeepSpeed는 DDP에 비해 최대 배치 크기 5배 증가와 전체 훈련 시간 약 3.5배 가속을 가능하게 했으며, 성능 지표는 감소하지 않았습니다.

DeepSpeed 설정 파일을 활용한 고급 구성

보다 세밀한 제어를 위해 사용자는 accelerate config를 통해 DeepSpeed 설정 JSON 파일을 제공할 수 있습니다. 설정 파일이 옵티마이저와 스케줄러를 정의할 경우, 사용자는 표준 PyTorch 옵티마이저와 스케줄러를 accelerate.utils.DummyOptim 및 accelerate.utils.DummyScheduler로 교체하여 최소한의 코드 수정만 하면 됩니다.

MuDoConv 데이터셋에서 BlenderBot-400M-distill 모델에 이 방식을 적용한 테스트 결과, DeepSpeed ZeRO Stage-2는 DDP의 100에 비해 최대 배치 크기 200을 허용했습니다. 이는 훈련 속도 1.44배 가속 및 평가 속도 1.23배 가속을 가져왔습니다.

CPU/디스크 오프로드를 활용한 초대형 모델 학습

CPU 오프로드를 적용한 DeepSpeed ZeRO Stage-3은 배치 크기 1조차도 GPU 메모리에 들어가지 않는 모델의 학습을 가능하게 합니다.

단일 24GB NVIDIA Titan RTX GPU로 GPT-XL 모델(15억 파라미터)을 학습하는 테스트에서 DDP는 즉시 OOM 오류를 발생시켰습니다. 반면, 옵티마이저 상태, 그래디언트 및 파라미터를 CPU에 오프로드한 DeepSpeed ZeRO Stage-3은 배치 크기 16으로 학습을 성공적으로 수행했으며, 에포크당 6608.35초가 소요되었습니다:

방법	최대 배치 크기	에포크당 학습 시간 (초)	비고
DDP	-	-	OOM 오류
DeepSpeed ZeRO Stage 3	16	6608.35	-