zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-19 11:28:07 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge branch 'ml-explore:main' into adding-orpo-training

Browse Source
This commit is contained in:
Gökdeniz Gülmez
2025-02-04 11:04:40 +01:00
committed by GitHub
parent 2c96da5155 21d0ab6e8a
commit c33c245c11
7 changed files with 52 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
llms/mlx_lm/examples/pipeline_generate.py
View File

@@ -44,7 +44,8 @@ def shard_and_load(repo):
allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
)
# Lazy load and shard model
# Lazy load and shard model to figure out
# which weights we need
model, _ = load_model(model_path, lazy=True, strict=False)
group = mx.distributed.init(backend="mpi")
@@ -62,8 +63,11 @@ def shard_and_load(repo):
# Download weights for local shard
download(args.model, allow_patterns=local_files)
# Load and shard the model, and load the weights
tokenizer = load_tokenizer(model_path)
model, _ = load_model(model_path)
model, _ = load_model(model_path, lazy=True, strict=False)
model.model.pipeline(group)
mx.eval(model.parameters())
# Synchronize processes before generation to avoid timeout if downloading
# model for the first time.