Load model faster for checkpoint conversion (#5)

* add accelerate lib for faster model loading * formatting * remove torch * fix load option w/o init_empty_weight context * formatting * fix accelerate lib to mllib * add einops for mllib
friendliai · Jul 26, 2023 · e6b2eec · e6b2eec
1 parent aab2250
commit e6b2eec
Show file tree

Hide file tree

Showing 3 changed files with 236 additions and 16 deletions.
diff --git a/periflow/sdk/resource/checkpoint.py b/periflow/sdk/resource/checkpoint.py
@@ -905,6 +905,7 @@ def convert(
                 torch_dtype=torch.float32,
                 cache_dir=cache_dir,
                 trust_remote_code=True,
+                low_cpu_mem_usage=True,  # For model loading faster and using ~1x model size CPU memory. https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.from_pretrained.example
             ).state_dict()
             logger.info(
                 "Hugging Face checkpoint(%s) is successfully loaded!",