Hugging Face가 PyTorch FSDP와 meta device를 활용한 단계별 모델 로딩으로 Llama 2 70B 파인튜닝 시 CPU RAM 사용량을 2TB에서 1.5GB 수준으로 감소

Context

Llama 70B 모델을 FSDP로 파인튜닝할 때 각 프로세스가 전체 모델을 로드하면 CPU RAM이 2TB를 초과하여 프로세스가 강제 종료된다. 전체 체크포인트 저장 시 NCCL Timeout 에러가 발생하고 GPU 메모리 사용량과 학습 속도가 최적화되지 않은 상태다.

Technical Solution

• Meta device를 사용한 선택적 로딩: rank 0에서만 사전학습 모델의 state dict를 CPU에 로드하고, 나머지 rank는 torch.empty()로 빈 파라미터 생성 후 sync_module_states=True로 FSDP가 학습 시작 전에 가중치 브로드캐스트 처리
• Hugging Face Transformers PR#25107과 Accelerate PR#1777을 통해 대규모 모델 초기화 최적화: rank 0의 CPU peak memory 32744MB에서 rank 1의 1506MB 수준까지 분산
• Flash Attention V2 설치 및 적용으로 VRAM 사용량과 연산 시간 감소
• Gradient Checkpointing 활성화로 중간 활성화값 저장 메모리 절감
• Accelerate와 SLURM 통합으로 2개 노드 × 8개 GPU(A100 80GB)의 다중 노드 학습 환경 구성

Key Takeaway

대규모 언어 모델 파인튜닝에서 meta device를 활용한 선택적 로딩과 동기화는 CPU 메모리 병목을 제거하는 실용적 패턴이며, Hugging Face 프레임워크의 FSDP 통합은 분산 학습의 인프라 복잡성을 크게 단순화한다.