Hugging Face Accelerate가 FSDP의 정밀도 처리를 DeepSpeed와 일치하도록 수정해 두 프레임워크 간 학습 결과 편차 제거

Context

Hugging Face Accelerate를 통해 DeepSpeed와 PyTorch FSDP로 Mistral-7B 모델을 학습할 때 동일한 설정에서도 손실 함수 수렴 패턴이 달랐다. DeepSpeed는 학습률 스케일링 없이 수렴했지만 FSDP는 학습률을 4배로 올려야만 수렴했으며, 이는 두 프레임워크의 정밀도 처리 방식 차이에서 기인했다.

Technical Solution

• DeepSpeed의 내부 동작 분석: DeepSpeedZeroOptimizer_Stage3에서 trainable_param_groups을 _create_fp32_partitions 함수로 처리하면서 모든 마스터 가중치를 float32로 업캐스팅
• FSDP의 정밀도 처리 현황 파악: 모델 로드 시 torch_dtype을 따르지만 "평탄화된 파라미터" 생성 시 float32 강제 적용 불가, 옵티마이저 단계에서 torch_dtype으로 연산 수행
• Hugging Face Accelerate 0.30.0에서 FSDP용 자동 업캐스팅 기능 추가: mixed_precision 활성화 시 FSDP도 float32로 옵티마이저 연산 수행하도록 변경
• 두 가지 FSDP 모드 제공: "메모리 제약" 모드(bf16 전체 연산)와 "혼합정밀도" 모드(float32로 옵티마이저 연산)
• 마이그레이션 가이드 업스트림: FSDP와 DeepSpeed 간 동등한 샤딩 전략, 모델 로드 방식, 체크포인트 처리 방법 문서화

Impact

• IBM Granite 7B 모델 기준 4개 A100 GPU에서 FSDP(정렬 모드) 3158.7 tokens/sec/device, DeepSpeed 3094.5 tokens/sec/device로 동등한 처리량 확인
• FSDP(정렬 모드) Model Flops Utilization 0.41, DeepSpeed 0.40으로 유사한 효율성 달성
• 두 프레임워크 간 학습 손실 곡선과 그래디언트 노름 특성이 동일한 학습률 설정에서 일치

Key Takeaway

혼합정밀도 학습 시 trainable parameters는 float32로 유지하는 것이 수렴 안정성의 핵심이며, GPU 수가 많을 때는 업캐스팅 오버헤드가 분산되어 무시할 수 있지만 소수 GPU에서는 메모리 비용(2배)이 유의미하므로 프레임워크별로 유연한 선택지를 제공하는 것이 중요하다.