Model parallel generation (#676)

Author: angeloskath

mlx_lm/benchmark.py

@@ -6,7 +6,7 @@
 
 from mlx_lm import batch_generate, load, stream_generate
 from mlx_lm.generate import DEFAULT_MODEL
-from mlx_lm.utils import pipeline_load
+from mlx_lm.utils import pipeline_load, sharded_load
 
 
 def setup_arg_parser():
@@ -49,6 +49,11 @@ def setup_arg_parser():
         help="Number of timing trials",
         type=int,
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     return parser
 
 
@@ -59,6 +64,8 @@ def main():
 
     group = mx.distributed.init()
     rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
 
     def rprint(*args, **kwargs):
         if rank == 0:
@@ -67,7 +74,9 @@ def rprint(*args, **kwargs):
     model_path = args.model or DEFAULT_MODEL
 
     if group.size() > 1:
-        model, tokenizer, config = pipeline_load(args.model, return_config=True)
+        model, tokenizer, config = sharded_load(
+            args.model, pipeline_group, tensor_group, return_config=True
+        )
     else:
         model, tokenizer, config = load(
             args.model, return_config=True, tokenizer_config={"trust_remote_code": True}

mlx_lm/chat.py

@@ -7,7 +7,7 @@
 from .generate import stream_generate
 from .models.cache import make_prompt_cache
 from .sample_utils import make_sampler
-from .utils import load
+from .utils import load, sharded_load
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
@@ -79,35 +79,54 @@ def setup_arg_parser():
         default=None,
         help="System prompt to be used for the chat template",
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     return parser
 
 
 def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
+    group = mx.distributed.init()
+    rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
     if args.seed is not None:
         mx.random.seed(args.seed)
 
-    model, tokenizer = load(
-        args.model,
-        adapter_path=args.adapter_path,
-        tokenizer_config={
-            "trust_remote_code": True if args.trust_remote_code else None
-        },
-    )
+    if group.size() > 1:
+        if args.adapter_path:
+            parser.error("Adapters not supported in distributed mode")
+        model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)
+    else:
+        model, tokenizer = load(
+            args.model,
+            adapter_path=args.adapter_path,
+            tokenizer_config={
+                "trust_remote_code": True if args.trust_remote_code else None
+            },
+        )
 
     def print_help():
-        print("The command list:")
-        print("- 'q' to exit")
-        print("- 'r' to reset the chat")
-        print("- 'h' to display these commands")
+        rprint("The command list:")
+        rprint("- 'q' to exit")
+        rprint("- 'r' to reset the chat")
+        rprint("- 'h' to display these commands")
 
-    print(f"[INFO] Starting chat session with {args.model}.")
+    rprint(f"[INFO] Starting chat session with {args.model}.")
     print_help()
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
-        query = input(">> ")
+        query = input(">> " if rank == 0 else "")
         if query == "q":
             break
         if query == "r":
@@ -139,8 +158,8 @@ def print_help():
             ),
             prompt_cache=prompt_cache,
         ):
-            print(response.text, flush=True, end="")
-        print()
+            rprint(response.text, flush=True, end="")
+        rprint()
 
 
 if __name__ == "__main__":

mlx_lm/examples/pipeline_generate.py mlx_lm/examples/sharded_generate.pymlx_lm/examples/pipeline_generate.py renamed to mlx_lm/examples/sharded_generate.py

@@ -1,33 +1,34 @@
-# Copyright © 2024 Apple Inc.
+# Copyright © 2025 Apple Inc.
 
 """
 Run with:
 
 ```
 mlx.launch \
- --hostfile /path/to/hosts.json \
- /path/to/pipeline_generate.py \
- --prompt "hello world"
+    --backend jaccl \
+    --env MLX_METAL_FAST_SYNCH=1 \
+    --hostfile /path/to/hosts.json \
+    /path/to/sharded_generate.py \
+    --prompt 'Hello world'
 ```
 
-Make sure you can run MLX over MPI on two hosts. For more information see the
-documentation:
+For more information on running distributed programs with MLX see the documentation:
 
-https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
+https://ml-explore.github.io/mlx/build/html/usage/distributed.html .
 """
 
 import argparse
 
 import mlx.core as mx
 
 from mlx_lm import stream_generate
-from mlx_lm.utils import pipeline_load
+from mlx_lm.utils import sharded_load
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
+    parser = argparse.ArgumentParser(description="LLM distributed inference example")
     parser.add_argument(
         "--model",
-        default="mlx-community/DeepSeek-R1-3bit",
+        default="mlx-community/Llama-3.3-70B-Instruct-4bit",
         help="HF repo or path to local model.",
     )
     parser.add_argument(
@@ -43,16 +44,23 @@
         default=256,
         help="Maximum number of tokens to generate",
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     args = parser.parse_args()
 
     group = mx.distributed.init()
     rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
 
     def rprint(*args, **kwargs):
         if rank == 0:
             print(*args, **kwargs)
 
-    model, tokenizer = pipeline_load(args.model)
+    model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)
 
     messages = [{"role": "user", "content": args.prompt}]
     prompt = tokenizer.apply_chat_template(

mlx_lm/models/deepseek_v2.py

@@ -6,6 +6,7 @@
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .pipeline import PipelineMixin
@@ -315,13 +316,21 @@ def __init__(self, config: ModelArgs):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         inds, scores = self.gate(x)
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(x)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -395,7 +404,8 @@ def __call__(
                 cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -429,6 +439,62 @@ def sanitize(self, weights):
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.self_attn.q_b_proj = shard_linear(
+                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                )
+            layer.self_attn.kv_b_proj = shard_linear(
+                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV2MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
     @property
     def layers(self):
         return self.model.pipeline_layers