DeepSpeed 및 FairScale를 통한 Hugging Face Transformers ZeRO 통합
Hugging Face Transformers v4.2.0은 DeepSpeed 및 FairScale에 대한 실험적 지원을 도입하여, Zero Redundancy Optimizer (ZeRO)를 통합함으로써 사용자가 GPU 메모리 사용량을 최적화하여 더 큰 모델을 학습시키고 배치 크기를 늘릴 수 있도록 합니다. 이 통합을 통해 --sharded_ddp (FairScale) 및 --deepspeed (DeepSpeed) 명령줄 인수를 사용하여 Trainer API를 통해 이러한 최적화를 직접 사용할 수 있습니다.
멀티 GPU 학습을 위한 ZeRO 메모리 최적화
DeepSpeed 또는 FairScale를 통해 ZeRO를 통합하면 표준 Distributed Data Parallel (DDP) 베이스라인과 비교하여 최대 가능한 배치 크기(BS)를 늘리는 동시에 학습 및 평가 시간을 크게 단축할 수 있습니다. 두 개의 24GB Titan RTX GPU에서 t5-large 모델을 사용한 벤치마크에서 다음과 같은 성능 향상이 관찰되었습니다:
| Method | Max Batch Size | Train Time | Eval Time |
|---|---|---|---|
| Baseline (DDP) | 16 | 30.9458 | 56.3310 |
| fp16 | 20 | 21.4943 | 53.4675 |
| sharded_ddp (FairScale) | 30 | 25.9085 | 47.5589 |
| sharded_ddp + fp16 | 30 | 17.3838 | 45.6593 |
| DeepSpeed (no CPU offload) | 40 | 10.4007 | 34.9289 |
| DeepSpeed (with CPU offload) | 50 | 20.9706 | 32.1409 |
DeepSpeed는 배치 크기와 학습 속도에서 가장 높은 이득을 보여주었으며, FairScale는 설정 파일 없이 단일 명령줄 인수만 필요하므로 배포가 더 쉬운 것으로 언급되었습니다.
대규모 모델을 위한 단일 GPU 학습
DeepSpeed는 CPU offloading을 통해 단일 GPU의 메모리 용량을 초과할 수 있는 모델의 학습을 가능하게 합니다. 단일 24GB RTX-3090 카드에서 t5-3b 모델을 사용한 테스트에서, 표준 단일 GPU 설정은 배치 크기가 1임에도 불구하고 Out of Memory (OOM) 오류로 실패했습니다. DeepSpeed를 사용하면 배치 크기 20으로 모델을 성공적으로 학습할 수 있었으며, 시스템은 배치 크기가 30일 때만 OOM에 도달했습니다.
ZeRO의 기술적 메커니즘
ZeRO (Zero Redundancy Optimizer)는 데이터 병렬 학습에 분산 데이터 스토리지를 추가하여 메모리를 최적화합니다. 전체 모델 상태를 모든 GPU에 복제하는 대신, ZeRO는 파라미터, 그래디언트, 그리고 옵티마이저 상태를 사용 가능한 GPU에 분할하여 저장합니다.
분산 파티셔닝
각 GPU는 파라미터, 그래디언트, 그리고 옵티마이저 상태의 단일 샤드(shard)만을 저장합니다. 런타임 시, 각 GPU는 특정 레이어에 필요한 데이터를 다른 참여 GPU로부터 즉시 가져오므로 데이터 저장의 중복이 발생하지 않습니다.
ZeRO-Offload
ZeRO-Offload는 특정 프로세싱 및 메모리 요구 사항을 GPU에서 호스트 CPU로 이동시키며, 이는 t5-3b와 같은 거대한 모델을 제한된 GPU 하드웨어에 맞추는 데 매우 중요합니다.
메모리 단편화 관리
DeepSpeed는 GPU 메모리 단편화(OOM 오류가 전체 가용 메모리에도 불구하고 연속적인 블록이 충분히 크지 않아 발생하는 현상)를 해결하기 위해 GPU 메모리를 독립적으로 관리하여 장기 및 단기 할당을 분리합니다.
배포 및 통합
ZeRO 최적화는 모델 아키텍처를 수정할 필요가 없으며, 학습 코드에서만 변경이 필요합니다. Hugging Face Trainer 사용자는 다음 플래그를 사용하여 이러한 기능을 구현할 수 있습니다:
--sharded_ddp: FairScale 통합을 활성화합니다.--deepspeed: DeepSpeed 통합을 활성화합니다 (JSON 설정 파일이 필요합니다).
언급된 향후 개선 사항에는 DeepSpeed Sparse Attention, 1-bit Adam, 그리고 FairScale와 DeepSpeed 모두에서 모델 파라미터 샤딩을 지원할 것에 대한 기대가 포함됩니다.