Use depends in pipeline parallel (#483)

Author: awni

mlx_lm/models/deepseek_v2.py

@@ -414,6 +414,8 @@ def __call__(
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                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]]

mlx_lm/models/deepseek_v3.py

@@ -446,6 +446,8 @@ def __call__(
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                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]]

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.29.1
+mlx>=0.29.2
 numpy
 transformers>=4.39.3
 protobuf