Run DeepSeek R1 or V3 with MLX Distributed

mlx_distributed_deepseek.md

Setup

On every machine in the cluster install openmpi and mlx-lm:

conda install conda-forge::openmpi
pip install -U mlx-lm

Next download the pipeline parallel run script. Download it to the same path on every machine:

curl -O https://raw.githubusercontent.com/ml-explore/mlx-examples/refs/heads/main/llms/mlx_lm/examples/pipeline_generate.py

Make a hosts.json file on the machine you plan to launch the generation. For two machines it should look like this:

[
  {"ssh": "hostname1"},
  {"ssh": "hostname2"}
]

Also make sure you can ssh hostname from every machine to every other machine. Check-out the MLX documentation for more information on setting up and testing MPI.

Set the wired limit on the machines to use more memory. For example on a 192GB M2 Ultra set this:

sudo sysctl iogpu.wired_limit_mb=180000

Run

Run the generation with a command like the following:

mlx.launch \
  --hostfile path/to/hosts.json \
  --backend mpi \
  path/to/pipeline_generate.py \ 
  --prompt "What number is larger 6.9 or 6.11?" \
  --max-tokens 128 \
  --model mlx-community/DeepSeek-R1-4bit

For DeepSeek R1 quantized in 3-bit you need in aggregate 350GB of RAM accross the cluster of machines, e.g. two 192 GB M2 Ultras. To run the model quantized to 4-bit you need 450GB in aggregate RAM or three 192 GB M2 Ultras.

Yes, I have set up sysctl for both devices, which allows me to control remote device downloads of models. However, downloads usually encounter errors and display network issues. I am in China, is it due to regional restrictions?
`mpirun -np 3 --hostfile hosts.txt python3 pipeline_generate.py --prompt "What number is larger 6.9 or 6.11?" --model mlx-community/DeepSeek-R1-3bit
/Users/zhangchi/Library/Python/3.9/lib/python/site-packages/urllib3/init.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+, currently the 'ssl' module is compiled with 'LibreSSL 2.8.3'. See: urllib3/urllib3#3020
warnings.warn(
/Users/zhangchi/Library/Python/3.9/lib/python/site-packages/urllib3/init.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+, currently the 'ssl' module is compiled with 'LibreSSL 2.8.3'. See: urllib3/urllib3#3020
warnings.warn(
Fetching 5 files: 100%|██████████| 5/5 [00:00<00:00, 22525.80it/s]
Fetching 5 files: 100%|██████████| 5/5 [00:00<00:00, 46707.17it/s]
Fetching 70 files: 17%|█▋ | 12/70 [06:54<37:14, 38.52s/it]--------------------------------------------------------------------------
PRTE has lost communication with a remote daemon.

HNP daemon : [prterun-Mac-Studio-6-41295@0,0] on node Mac-Studio-6
Remote daemon: [prterun-Mac-Studio-6-41295@0,1] on node Mac-Studio-8

This is usually due to either a failure of the TCP network
connection to the node, or possibly an internal failure of
the daemon itself. We cannot recover from this failure, and
therefore will terminate the job.
--------------------------------------------------------------------------`

Does the new Mac Studios with M3 Ultra and max 512Gb of unified memory can run native FP8 deepseek-r1 model? Does the new Mac Studios with M3 Ultra and max 512Gb of unified memory support FP8?

You can run a 4-bit quantized model on the 512GB machine. 8-bit is too big.

You can run a 4-bit quantized model on the 512GB machine. 8-bit is too big.

If i hvae 4 such Mac Studios, each with 512GB, how can i run the 8-bit model distributed on the 4 machines?

• If each machine has 512 you only need 2 of them (and I wouldn't recommend using more because it will be slower)
• The above setup should work.. though we'd need to make an 8-bit quant for that.

You can make the 8-bit quant like so once ml-explore/mlx-lm#32 lands.

mlx_lm.convert --hf-path deepseek-ai/DeepSeek-R1 -q --q-bits 8 --upload-repo mlx-community/DeepSeek-R1-8bit

• If each machine has 512 you only need 2 of them (and I wouldn't recommend using more because it will be slower)
• The above setup should work.. though we'd need to make an 8-bit quant for that.

Why does 4 machines slower than 2 machines? Because the connection speed of Thunderbolt 5 is too slow? I thought that more machines mean more kvcache space and it would be faster.

With pipeline parallelism (which is used here).. if you have enough RAM to fit the model then you only are adding communication latency as you add more machines.

With pipeline parallelism (which is used here).. if you have enough RAM to fit the model then you only are adding communication latency as you add more machines.

Does mlx-lm support tensor parallelism?

Sort of. There is a PR for it.. but it's too slow right now to be practical. Something we are working on / hoping to improve for the future.

404 for the url
https://raw.githubusercontent.com/ml-explore/mlx-examples/refs/heads/main/llms/mlx_lm/examples/pipeline_generate.py

The run script can be downloaded at
https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/examples/pipeline_generate.py

Good News

pipeline_generate.py work very well with other DeepSeek model "DeepSeek-V2.5-1210-3bit "

mlx.launch --hosts mac1,mac2 --backend mpi "pipeline_generate.py" --max-tokens 12800 --model mlx-community/DeepSeek-V2.5-1210-3bit --prompt "Generate a python script"

==========
Prompt: 21 tokens, 85.378 tokens-per-sec
Generation: 776 tokens, 17.794 tokens-per-sec
Peak memory: 55.234 GB

mlx.launch --hosts mac1,mac2 --backend mpi "pipeline_generate.py" --max-tokens 12800 --model mlx-community/DeepSeek-V2.5-1210-4bit --prompt "Generate a python script"

==========
Prompt: 21 tokens, 80.473 tokens-per-sec
Generation: 901 tokens, 17.410 tokens-per-sec
Peak memory: 70.257 GB

Less good News

1°) When I run mlx_distributed_deepseek.py
error message :

except statement is broken in "distributed_run.py"

Edit around line 175. Find:
in the file "except e:"
replace with
"except Exception as e:"

2°) And when I run this command: mlx.distributed_config --verbose --hosts
error message :

/miniconda3/envs/mlxmpi/lib/python3.11/site-packages/mlx/distributed_run.py", line 507, in prepare_tb_ring
connected_to = items[0]["domain_uuid_key"]
~~~~~~~~^^^^^^^^^^^^^^^^^^^
KeyError: 'domain_uuid_key'

Does mlx support gguf format?

To set up a machine cluster for MLX generation, install OpenMPI and mlx-lm on each machine using conda and pip. Download the pipeline script and place it in the same path on all machines. Create a hosts.json file listing the machines, ensuring SSH access between them. Set memory limits with sudo sysctl. If you’ve ever had to finish an essay overnight, you know the pressure. I’ve relied on UKWritings in such cases their write my essay service at https://ukwritings.com/write-my-essay is perfect for fast, high-quality academic writing that meets urgent deadlines.

Also works for DeepSeek v2. But not for other models.

would it be possible, to make a smaller model have this option of distributed, just to be able to get it to work on less expensive hardware haha. I have a few base model m4 mac minis and before i fork out for the m3 ultras, I'd love to see it actually worked. tried many ways to achieve it with grpc, tensor parallelism, etc. most of the time things got corrupted in the kv cache synchronization.

You can try one of the smaller DSV2 models like mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit

You can try one of the smaller DSV2 models like mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit

Thank you so much Awni, I was trying to convince Claude Code to have it work with Qwen models but it just would not get it to work properly. Then I guess it had too much memory of that approach to accept your guidance. Finally got it to work and I asked Claude to do a beginner friendly writeup, please review/comment as i think this would be a nice example that many people would be able to run.
https://github.com/georgiedekker/mlx_distributed_ring_inference

@awni made another attempt at speeding things up. seems to work please review, feel free to use any as example in the mlx repo. https://github.com/georgiedekker/mlx_distributed_ring_inference_v2

Hi @awni,

I have been using the pipeline generate with mpi distributed (ring, two/three nodes, TB network) since may but starting from version "mlx-lm==0.27.0", I am seeing this crash (gpu timeout)

libc++abi: terminating due to uncaught exception of type std::runtime_error: [METAL] Command buffer execution failed: Caused GPU Timeout Error (00000002:kIOGPUCommandBufferCallbackErrorTimeout)

This is a very deterministic issue on my system in multiple scenarios (servers mpi, single text generation...), basically all the pipeline distributed generation seems to stop working from mlx-lm v0.27.0, while I have still no issue if I remain to previous versions, e.g. python -m pip install -U "mlx-lm==0.26.4". The nodes I am running are M4, M2 Pro and a mix of them.

For debugging, I was able to reproduce the issue using the pipeline_generate.py from the example directory and the following command - note I am using the smaller model due to hw ram sizes, but this used to be supported by pipeline_generate as part of the deepseek_v2 family:

clear && mlx.launch \
    --hostfile hosts_2.json \
    --backend ring \
    pipeline_generate.py \
    --prompt "What number is larger 6.9 or 6.11?" \
    --max-tokens 512 \
    --model mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit

Below the two logs produced on my conda env with the old version of mlx-lm (up to 0.26.4) and the new versions (either 0.27.0 or 0.27.1):

python -m pip install -U "mlx-lm==0.26.4":

python -m pip install -U "mlx-lm==0.27.0":

Fix is incoming here: ml-explore/mlx-lm#483