Run:ai Model Streamer - ускорение загрузки LLM

Run:ai Model Streamer – Python SDK, разработанный для оптимизации загрузки моделей машинного обучения. Он поддерживает загрузку моделей в различных форматах (.pt, .h5, .safetensors и др.) из сетевых файловых систем, хранилищ S3 и локальных дисков.

Особенность Streamer - использование многопоточности для параллельной загрузки тензоров из файла в выделенный буфер оперативной памяти.

Каждый тензор идентифицируется уникальным ключом, который впоследствии используется приложением для загрузки тензора в память GPU. Это дает возможность загружать тензоры в память GPU одновременно с чтением других тензоров из хранилища в оперативную память, минимизируя время простоя GPU.

Streamer использует высокопроизводительный слой на C++, а Python API обеспечивает удобную интеграцию Streamer в существующие проекты, например, для автомасштабируемых серверов инференса, где минимизация времени простоя GPU критически важна.

Тест производительности Run:ai Model Streamer выполнялся на NVIDIA A10G с моделью Llama-3-8B (15 GB) и сравнивался с загрузчиками SafeTensors от Hugging Face и Tensorizer от CoreWeave.

При использовании локальных SSD, Run:ai Model Streamer достигал максимальной пропускной способности SSD (1 ГБ/с для GP3 и 2 ГБ/с для IO2), сокращая время загрузки модели в 6 раз по сравнению с SafeTensors Loader.

На Amazon S3 Run:ai Model Streamer загружал модель за 4.88 секунды, значительно превосходя Tensorizer (37.36 секунд).

Streamer поддерживает только приложения PyTorch.

Размер буфера оперативной памяти регулируется параметром RUNAI_STREAMER_MEMORY_LIMIT

Пример запуска с локального диска:

# Install streamer from pip  
 pip install runai-model-streamer 
 
 # Load the tensors to the buffer and stream to the GPU 
 
 from runai_model_streamer import SafetensorsStreamer 
 
 file_path = "/path/to/file.safetensors" 
 
 with SafetensorsStreamer() as streamer: 
 
     streamer.stream_file(file_path) 
 
     for name, tensor in streamer.get_tensors(): 
 
         tensor.to('CUDA:0')

Лицензирование: Apache 2.0 License.

Бенчмарки в блоге RunAI (https://www.run.ai/blog/run-ai-model-streamer-performance-benchmarks)

Документация (https://github.com/run-ai/runai-model-streamer/blob/master/docs/README.md)

GitHub (https://github.com/run-ai/runai-model-streamer)

Источник: github.com

Комментарии: