Hugging Face Accelerate: PyTorch로 대형 모델 실행

Hugging Face Accelerate는 사용자가 소비자 하드웨어의 가용 RAM이나 GPU 메모리를 초과하는 매우 큰 언어 모델(LLM)을 로드하고 실행할 수 있게 합니다. 특정 PyTorch 기능과 수정된 로딩 파이프라인을 활용하여 Accelerate는 모델 가중치를 여러 GPU, CPU RAM 및 디스크 저장소 등 사용 가능한 하드웨어 자원에 분산시킵니다.

모델 로딩 시 메모리 제약 극복

전통적인 PyTorch 모델 로딩은 선형적인 과정을 따릅니다: 모델을 생성하고, 메모리 내 state_dict에 가중치를 로드한 뒤, 그 가중치를 모델에 적용하고, 모델을 디바이스로 이동합니다. 매우 큰 모델의 경우, 이 과정은 RAM 측면에서 감당하기 어려울 정도로 비용이 많이 듭니다. 예를 들어, float32 정밀도로 67억 개의 파라미터를 가진 모델은 초기 모델 생성만으로 약 26.8GB의 RAM이 필요하고, state_dict를 로드하는 데 또 다른 26.8GB가 필요해, 모델이 GPU에 도달하기 전까지 총 53GB가 넘습니다.

이를 해결하기 위해 Accelerate는 보다 메모리 효율적인 파이프라인을 구현합니다:

가중치가 없는 빈 모델을 생성합니다.
각 레이어가 위치할 디바이스를 결정하기 위해 디바이스 맵을 정합니다.
가중치를 작은 조각(샤드)으로 로드합니다.
그 가중치를 빈 모델에 로드합니다.
추론을 위한 지정된 디바이스로 가중치를 이동합니다.
남은 모든 가중치에 대해 이 과정을 반복합니다.

메타 디바이스를 통한 효율적인 모델 초기화

Accelerate는 PyTorch 1.9에서 도입된 "meta" 디바이스를 활용해 실제 데이터 메모리를 할당하지 않고 모델을 인스턴스화합니다. 메타 디바이스상의 텐서는 형태와 데이터 타입 정보만 저장하므로, CPU나 GPU RAM을 소비하지 않고도 임의로 큰 모델을 생성할 수 있습니다.

device 키워드를 지원하도록 Transformers 라이브러리의 모든 모델을 재작성하는 것은 비현실적이기 때문에, Hugging Face는 init_empty_weights() 컨텍스트 매니저를 개발했습니다. 이를 통해 어떤 모델이든 메타 디바이스 상의 "쉘"로 인스턴스화할 수 있어, 실제 가중치를 로드하지 않고도 메모리 요구량을 계산할 수 있는 구조적 정보를 제공합니다.

자동 디바이스 매핑 및 자원 할당

Accelerate는 infer_auto_device_map 함수를 사용해 모델 가중치를 사용 가능한 하드웨어에 자동으로 분산시킵니다. 시스템은 다음 순서로 자원을 우선순위에 두고 할당합니다: GPU → CPU RAM → 디스크 오프로드.

디바이스 맵 구성

사용 사례에 따라 사용자는 다양한 매핑 전략을 선택할 수 있습니다:

"auto" 또는 "balanced": 사용 가능한 모든 GPU에 가중치를 균등하게 분배합니다.
"balanced_low_0": GPU에 가중치를 균등하게 분배하지만 첫 번째 GPU(GPU 0)의 부하를 최소화합니다. 이는 첫 번째 GPU가 모델 출력(예: 텍스트 생성)에 필요할 때 유용합니다.
"sequential": GPU를 순서대로 채우며, 이후 GPU가 사용되지 않을 수 있습니다.

단일 레이어가 여러 디바이스에 걸쳐 분할되는 것을 방지하기 위해(잔차 연결이 깨지는 것을 방지) Accelerate는 no_split_module_classes(예: ["OPTDecoderLayer"])를 지정하도록 허용합니다.

RAM 사용량 감소를 위한 샤드 체크포인트

단일 거대한 state_dict 파일을 로드하는 것은 대부분의 하드웨어에서 비현실적입니다; 예를 들어 BLOOM 모델(176B 파라미터)은 bfloat16 형식의 가중치를 로드하는 데만 352GB의 RAM이 필요합니다. 이를 완화하기 위해 Hugging Face는 sharded checkpoints를 사용합니다.

샤드 체크포인트에서는:

pytorch_model.bin.index.json 파일은 각 파라미터 이름을 특정 샤드 파일에 매핑합니다.
가중치는 여러 표준 PyTorch state dict 파일에 분할됩니다(예: BLOOM의 경우 72개 파일).

이 접근 방식은 전체 모델이 아니라 가장 큰 단일 샤드(예: BLOOM의 경우 7.19GB)만을 메모리에 올려두면 되도록 보장합니다. GPU와 CPU RAM이 충분하지 않을 경우, Accelerate는 지정된 offload_folder에 가중치를 디스크에 오프로드할 수 있습니다. offload_state_dict=True 옵션을 사용하면 다른 샤드가 처리되는 동안 CPU에 상주하는 모델 부분을 일시적으로 오프로드하여 RAM 사용량을 더욱 줄일 수 있습니다.

동적 훅을 통한 실행

여러 디바이스에 걸쳐 분할된 모델을 실행하기 위해 Accelerate는 PyTorch 훅에서 영감을 받은 메커니즘을 사용합니다. dispatch_model 함수는 모든 모듈 및 서브모듈에 훅을 연결해 각 forward 패스 전후에 실행됩니다. 이러한 훅은 다음 작업을 수행합니다:

모든 모듈 입력이 가중치와 동일한 디바이스에 있는지 확인합니다.
forward 패스 직전에 가중치를 CPU에서 GPU 0으로 이동하고, forward 패스 직후에 다시 CPU로 반환합니다.
forward 패스 전 디스크에서 RAM으로, 그 다음 GPU 0으로 가중치를 로드하고, forward 패스 후 즉시 메모리를 해제합니다.

이 접근 방식은 복잡한 파이프라인 병렬성을 사용하지 않고 GPU를 순차적으로 활용하지만, 훨씬 작은 하드웨어 환경에서도 거대한 모델을 실행할 수 있게 합니다.