Hugging Face Accelerate: DeepSpeed와 FSDP 정밀도 조화
Hugging Face Accelerate는 PyTorch Fully Sharded Data Parallel (FSDP)과 Microsoft DeepSpeed의 정밀도 처리를 맞추기 위한 업데이트를 도입했습니다. 이 변경 사항은 Accelerate 0.30.0 릴리스에 통합되어, 사용자가 내부 정밀도 기본값 차이로 인해 이전에 발생하던 수렴성 차이 없이 두 ZeRO Redundancy Optimizer 구현 간을 전환할 수 있게 합니다.
정밀도 차이와 수렴성
DeepSpeed와 FSDP가 매개변수 정밀도를 처리하는 방식의 차이는 훈련 결과가 달라지는 원인이 될 수 있습니다. bfloat16으로 로드된 Mistral-7B 기본 모델을 사용한 테스트에서, DeepSpeed는 안정적인 수렴을 보였지만 FSDP는 학습률을 GPU 수에 맞게 수동으로 스케일링하거나 1e-5 이하로 낮추지 않으면 손실을 감소시키지 못했습니다.
이 현상은 내부 업캐스팅 때문입니다. DeepSpeed의 DeepSpeedZeroOptimizer_Stage3는 _create_fp32_partitions를 통해 학습 가능한 매개변수 그룹을 자동으로 float32로 업캐스팅하여, 설계상 마스터 가중치를 전체 정밀도로 유지합니다. 이를 통해 낮은 정밀도에서는 불안정할 수 있는 학습률에서도 옵티마이저가 수렴할 수 있습니다.
정밀도 워크플로우 비교
| Process | FSDP | DeepSpeed |
|---|---|---|
| Preparation | torch_dtype 사용 |
torch_dtype 무시; float32로 생성 |
| Optimizer Initialization | 매개변수를 torch_dtype로 생성 |
매개변수를 float32로 생성 |
| Optimizer (Pre-step) | (있는 경우) torch_dtype로 업캐스팅 |
모든 것을 float32로 업캐스팅 |
| Optimizer (Actual step) | torch_dtype에서 수행 |
float32에서 수행 |
DeepSpeed의 기본 업캐스팅은 수렴을 보장하지만, 메모리 사용량을 2배로 증가시킬 수 있어 GPU 수가 적은 경우에 큰 영향을 미칩니다. 반면, torch 기반 FSDP 구현은 업캐스팅을 강제하지 않아, 옵티마이저가 저정밀도에서 동작하도록 함으로써 메모리 제한 상황에 더 큰 유연성을 제공합니다.
Accelerate 0.30.0에서 새로운 FSDP 모드
이러한 프레임워크를 조화시키기 위해, Hugging Face Accelerate는 수렴 안정성과 메모리 효율성 사이에서 사용자의 우선순위에 맞추어 두 가지 별도 FSDP 모드를 지원합니다:
- 혼합 정밀도 모드: 준비 단계와 옵티마이저 단계에서 매개변수를
fp32로 업캐스팅하여 DeepSpeed와 일치시키고, 훈련은bf16을 유지합니다. - 메모리 제한 모드: 준비, 훈련, 옵티마이저 단계 모두에서 저정밀도(
bf16)로 완전히 동작합니다.
프레임워크 비교 표
| 프레임워크 | 모델 로딩 | 혼합 정밀도 | 준비 | 훈련 | 옵티마이저 |
|---|---|---|---|---|---|
| FSDP (메모리 제한) | bf16 |
None | bf16 |
bf16 |
bf16 |
| FSDP (혼합 정밀도) | bf16 |
bf16 |
fp32 |
bf16 |
fp32 |
| DeepSpeed | bf16 |
bf16 |
fp32 |
bf16 |
fp32 |
처리량 성능
네 개의 A100 GPU에서 IBM Granite 7B 모델을 사용한 벤치마크 결과, 정렬된 모드의 FSDP와 DeepSpeed (Zero3)가 거의 동일한 처리량 성능을 제공함을 보여줍니다:
| 프레임워크 | 토큰 / 초 / 디바이스 | 스텝 시간 (초) | 모델 FLOPs 활용도 (MFU) |
|---|---|---|---|
| FSDP (정렬) | 3158.7 | 10.4 | 0.41 |
| DeepSpeed | 3094.5 | 10.6 | 0.40 |
마이그레이션 및 구성
Hugging Face는 사용자가 FSDP와 DeepSpeed 사이를 마이그레이션할 수 있도록 돕는 concept guide를 공개했습니다. 전환은 주로 Accelerate 설정 파일이나 DeepSpeed와 FSDP 플러그인 클래스를 통해 처리됩니다. 마이그레이션 시 주요 고려 사항은 다음과 같습니다:
- 샤딩 전략: 프레임워크 간에 동등한 샤딩을 달성하기.
- 모델 로딩: 효율적인 로딩 패턴 구현.
- 가중치 프리패칭: 가중치를 GPU 간에 이동하는 방법 관리.
- 체크포인팅: 각 프레임워크가 모델 상태를 저장하고 로드하는 방식의 차이 처리.