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

fix deepseek sharding (#1242)

Browse Source
This commit is contained in:
Awni Hannun
2025-02-03 16:59:50 -08:00
committed by GitHub
parent 0989c073b0
commit 21d0ab6e8a
3 changed files with 8 additions and 4 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.