Hugging Face가 PyTorch의 meta device와 동적 가중치 로딩으로 176억 개 파라미터 모델을 Colab 무료 인스턴스에서 실행 가능하게 구현

Context

기존 PyTorch 모델 로딩 파이프라인은 모델 생성 → 메모리 로드 → 가중치 주입 → 디바이스 이동의 순차 과정을 거친다. 6.7B 파라미터 모델(OPT-6.7B)은 float32 기본 정밀도에서 26.8GB RAM이 필요하고, 176B 파라미터 모델(BLOOM, OPT-176B)은 1.4TB CPU RAM이 필요해 일반 소비자 하드웨어에서 실행 불가능했다.

Technical Solution

• PyTorch 1.9의 meta device를 활용한 빈 모델 생성: 실제 데이터 없이 텐서의 형태(shape)만 가지고 메모리 할당 없이 모델 인스턴스 생성
• init_empty_weights() 컨텍스트 매니저 개발: 기존 Transformers 라이브러리의 150개 모델 코드 수정 없이 빈 모델 자동 생성
• 메타 디바이스의 형태 정보로 device_map 자동 계산: 각 가중치의 형태와 데이터 타입으로 메모리 소비량 사전 계산 후 CPU/GPU/디스크 배치 결정
• offload_folder와 offload_state_dict 파라미터를 통한 디스크 오프로딩: 로드할 수 없는 가중치를 디스크에 저장하고 필요시에만 로드
• dispatch_model 함수의 forward 전/후 훅(hook) 추가: 각 모듈 실행 전에 가중치를 같은 디바이스로 이동, 실행 후 CPU/디스크로 복원

Impact

아티클에서 정량적 성능 지표(속도 개선, 메모리 절감 %)가 명시되지 않음.

Key Takeaway

극도로 큰 모델을 로드할 때 전체 모델을 메모리에 올리는 대신, 계산 그래프의 형태 정보로 미리 배치를 결정하고 단계별로 필요한 부분만 동적으로 로드하는 지연 로딩 패턴이 핵심이다. 이는 단일 대용량 메모리 대신 여러 이기종 스토리지(GPU, CPU RAM, 디스크)를 순차적으로 활용하는 설계 원칙을 보여준다.