zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-12-16 02:08:55 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Lazy loading models for faster convert and merge (#462)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-02-20 13:36:55 -08:00
committed by GitHub
parent 8eee4399f4
commit dc4f2e0a6b
3 changed files with 41 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
llms/mlx_lm/convert.py
View File

@@ -96,7 +96,7 @@ def convert(
):
print("[INFO] Loading")
model_path = get_model_path(hf_path)
model, config, tokenizer = fetch_from_hub(model_path)
model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
weights = dict(tree_flatten(model.parameters()))
dtype = mx.float16 if quantize else getattr(mx, dtype)
@@ -110,7 +110,8 @@ def convert(
if isinstance(mlx_path, str):
mlx_path = Path(mlx_path)
save_weights(mlx_path, weights)
del model
save_weights(mlx_path, weights, donate_weights=True)
py_files = glob.glob(str(model_path / "*.py"))
for file in py_files: