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

llms/mlx_lm/examples/pipeline_generate.py

@@ -4,10 +4,11 @@
 Run with:
 
 ```
-/path/to/mpirun \
- -np 2 \
+mlx.launch \
  --hostfile /path/to/hosts.txt \
- python /path/to/pipeline_generate.py --prompt "hello world"
+ --backend mpi \
+ /path/to/pipeline_generate.py \
+ --prompt "hello world"
 ```
 
 Make sure you can run MLX over MPI on two hosts. For more information see the
@@ -17,62 +18,110 @@ https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
 """
 
 import argparse
+import json
+from pathlib import Path
 
 import mlx.core as mx
+from huggingface_hub import snapshot_download
+from mlx.utils import tree_flatten
 from mlx_lm import load, stream_generate
-
-parser = argparse.ArgumentParser(description="LLM pipelined inference example")
-parser.add_argument(
-    "--model",
-    default="mlx-community/DeepSeek-R1-3bit",
-    help="HF repo or path to local model.",
-)
-parser.add_argument(
-    "--prompt",
-    "-p",
-    default="Write a quicksort in C++.",
-    help="Message to be processed by the model ('-' reads from stdin)",
-)
-parser.add_argument(
-    "--max-tokens",
-    "-m",
-    type=int,
-    default=256,
-    help="Maximum number of tokens to generate",
-)
-args = parser.parse_args()
-
-model, tokenizer = load(args.model, lazy=True)
-
-messages = [{"role": "user", "content": args.prompt}]
-prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-
-group = mx.distributed.init()
-rank = group.rank()
-model.model.pipeline(group)
-mx.eval(model.parameters())
-
-# Synchronize processes before generation to avoid timeout if downloading
-# model for the first time.
-mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
+from mlx_lm.utils import load_model, load_tokenizer
 
 
-def rprint(*args, **kwargs):
-    if rank == 0:
-        print(*args, **kwargs)
+def download(repo: str, allow_patterns: list[str]) -> Path:
+    return Path(
+        snapshot_download(
+            repo,
+            allow_patterns=allow_patterns,
+        )
+    )
 
 
-for response in stream_generate(model, tokenizer, prompt, max_tokens=args.max_tokens):
-    rprint(response.text, end="", flush=True)
+def shard_and_load(repo):
+    # Get model path with everything but weight safetensors
+    model_path = download(
+        args.model,
+        allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
+    )
 
-rprint()
-rprint("=" * 10)
-rprint(
-    f"Prompt: {response.prompt_tokens} tokens, "
-    f"{response.prompt_tps:.3f} tokens-per-sec"
-)
-rprint(
-    f"Generation: {response.generation_tokens} tokens, "
-    f"{response.generation_tps:.3f} tokens-per-sec"
-)
-rprint(f"Peak memory: {response.peak_memory:.3f} GB")
+    # 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")
+    rank = group.rank()
+    model.model.pipeline(group)
+
+    # Figure out which files we need for the local shard
+    with open(model_path / "model.safetensors.index.json", "r") as fid:
+        weight_index = json.load(fid)["weight_map"]
+
+    local_files = set()
+    for k, _ in tree_flatten(model.parameters()):
+        local_files.add(weight_index[k])
+
+    # 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, 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.
+    mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
+    return model, tokenizer
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
+    parser.add_argument(
+        "--model",
+        default="mlx-community/DeepSeek-R1-3bit",
+        help="HF repo or path to local model.",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        default="Write a quicksort in C++.",
+        help="Message to be processed by the model ('-' reads from stdin)",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        "-m",
+        type=int,
+        default=256,
+        help="Maximum number of tokens to generate",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init(backend="mpi")
+    rank = group.rank()
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
+    model, tokenizer = shard_and_load(args.model)
+
+    messages = [{"role": "user", "content": args.prompt}]
+    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+
+    for response in stream_generate(
+        model, tokenizer, prompt, max_tokens=args.max_tokens
+    ):
+        rprint(response.text, end="", flush=True)
+
+    rprint()
+    rprint("=" * 10)
+    rprint(
+        f"Prompt: {response.prompt_tokens} tokens, "
+        f"{response.prompt_tps:.3f} tokens-per-sec"
+    )
+    rprint(
+        f"Generation: {response.generation_tokens} tokens, "
+        f"{response.generation_tps:.3f} tokens-per-sec"
+    )
+    rprint(f"Peak memory: {response.peak_memory:.3f} GB")