fastsafetensors로 모델 가중치 로드하기

fastsafetensors 라이브러리를 사용하면 GPU Direct Storage를 활용해 모델 가중치를 GPU 메모리로 로드할 수 있어요. 기존 safetensors 방식에 비해 로드 속도를 크게 높일 수 있어 대형 모델의 시작 시간(스타트업 시간)을 줄일 때 유용해요. 자세한 내용은 fastsafetensors의 GitHub 저장소를 참고하세요.

출처: 문서

본문

기본적으로 vLLM은 Hugging Face safetensors 형식으로 모델 가중치를 로드해요. fastsafetensors를 쓰면 GPU Direct Storage(GDS)를 통해 디스크에서 GPU 메모리로 직접 가중치를 읽어와 로드 대역폭을 대폭 개선할 수 있어요.

이 기능을 활성화하려면 커맨드라인 인자로 --load-format fastsafetensors를 사용하세요.

vllm serve <model> --load-format fastsafetensors

예를 들어:

vllm serve Qwen/Qwen2.5-0.5B --load-format fastsafetensors

이 옵션은 GPU가 GDS(GPU Direct Storage)를 지원하는 환경에서 특히 효과적이에요. GPU가 GDS를 지원하지 않더라도, fastsafetensors는 사용 가능한 환경에서 로딩 경로를 자동으로 최적화해요.

동작 방식

기존의 safetensors 로딩은 디스크에서 CPU 메모리로, 다시 CPU 메모리에서 GPU 메모리로 데이터를 복사해요. fastsafetensors는 GPU Direct Storage를 사용할 수 있을 때 이 CPU 경유 단계를 생략하고 스토리지에서 GPU 메모리로 직접 데이터를 이동시켜, CPU와 PCIe 대역폭 병목을 줄여 가중치 로딩 시간을 대폭 단축해요. 이는 수 GB~수십 GB에 달하는 대형 모델에서 스타트업 시간을 크게 줄이는 데 특히 유리해요.

사용 시 주의 사항

이 기능을 사용하려면 fastsafetensors 라이브러리가 설치되어 있어야 해요. pip install fastsafetensors로 설치할 수 있어요.
실제 성능 향상은 스토리지 종류(NVMe 등), 파일시스템, GPU의 GDS 지원 여부에 따라 달라져요.
모델을 이미 GPU 메모리에 캐시하거나 I/O가 아닌 다른 병목이 지배적인 경우에는 효과가 제한될 수 있어요.

팁: --load-format의 다른 값으로는 safetensors(기본), InstantTensor, tensorizer 등이 있어요. 각각 가중치 로딩 속도를 높이는 다른 백엔드를 제공하며, 사용 환경과 하드웨어에 맞게 선택하면 돼요.

fastsafetensors로 모델 가중치 로드하기