Speeding up Model Loading with fastsafetensors
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2025
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866912401140482048 |
|---|---|
| author | Yoshimura, Takeshi Chiba, Tatsuhiro Sethi, Manish Waddington, Daniel Sundararaman, Swaminathan |
| author_facet | Yoshimura, Takeshi Chiba, Tatsuhiro Sethi, Manish Waddington, Daniel Sundararaman, Swaminathan |
| contents | The rapid increases in model parameter sizes introduces new challenges in pre-trained model loading. Currently, machine learning code often deserializes each parameter as a tensor object in host memory before copying it to device memory. We found that this approach underutilized storage throughput and significantly slowed down loading large models with a widely-used model file formats, safetensors. In this work, we present fastsafetensors, a Python library designed to optimize the deserialization of tensors in safetensors files. Our approach first copies groups of on-disk parameters to device memory, where they are directly instantiated as tensor objects. This design enables further optimization in low-level I/O and high-level tensor preprocessing, including parallelized copying, peer-to-peer DMA, and GPU offloading. Experimental results show performance improvements of 4.8x to 7.5x in loading models such as Llama (7, 13, and 70 billion parameters), Falcon (40 billion parameters), and the Bloom (176 billion parameters). |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_23072 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Speeding up Model Loading with fastsafetensors Yoshimura, Takeshi Chiba, Tatsuhiro Sethi, Manish Waddington, Daniel Sundararaman, Swaminathan Distributed, Parallel, and Cluster Computing The rapid increases in model parameter sizes introduces new challenges in pre-trained model loading. Currently, machine learning code often deserializes each parameter as a tensor object in host memory before copying it to device memory. We found that this approach underutilized storage throughput and significantly slowed down loading large models with a widely-used model file formats, safetensors. In this work, we present fastsafetensors, a Python library designed to optimize the deserialization of tensors in safetensors files. Our approach first copies groups of on-disk parameters to device memory, where they are directly instantiated as tensor objects. This design enables further optimization in low-level I/O and high-level tensor preprocessing, including parallelized copying, peer-to-peer DMA, and GPU offloading. Experimental results show performance improvements of 4.8x to 7.5x in loading models such as Llama (7, 13, and 70 billion parameters), Falcon (40 billion parameters), and the Bloom (176 billion parameters). |
| title | Speeding up Model Loading with fastsafetensors |
| topic | Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2505.23072 |