Accelerate의 대규모 모델 로딩 원리

기존 PyTorch의 모델 로딩 방식은 모델 생성 후 가중치를 메모리에 올리는 과정을 거치는데, 이 과정에서 모델 크기보다 훨씬 큰 CPU RAM이 필요하다는 한계가 있다. 예를 들어 176B 파라미터 모델을 로드하려면 약 1.4TB의 RAM이 필요하다.

Accelerate는 이를 해결하기 위해 다음과 같은 최적화된 프로세스를 사용한다:

• 가중치가 없는 **빈 모델(empty model)**을 먼저 생성한다.
• 각 레이어가 위치할 디바이스를 결정한다.
• 가중치를 부분적으로 로드하여 모델에 주입하고 디바이스로 이동시키는 과정을 반복한다.

이 기술의 핵심은 PyTorch 1.9의 meta device 활용에 있다. meta device를 사용하면 실제 데이터 없이 텐서의 형태(shape) 정보만 가진 상태로 모델을 생성할 수 있어, 메모리 부족 문제 없이 거대한 모델 구조를 즉시 정의할 수 있다.

Hugging Face는 init_empty_weights 컨텍스트 매니저를 통해 기존 모델 코드를 수정하지 않고도 간편하게 빈 모델을 생성하고 대규모 모델을 효율적으로 로드할 수 있도록 지원한다.