[Proposal] RDMA weight streaming — eliminate per-node model storage for distributed inference

[guruswami-ai]

Summary

A proposal for RDMA weight streaming: a master node reads model weights from disk and broadcasts only the slices each rank needs via all_sum, eliminating the requirement to store the full model on every node. This builds on the existing mlx_lm.share and sharded_load infrastructure.

Motivation

Currently, distributed inference in MLX requires the full model to be stored on every participating node:

• TP (Tensor Parallel): sharded_load downloads all weight files to every node, then model.shard() selects the needed slices at eval time. Each node stores the full model on disk even though it only uses 1/N of the weights in memory.
• PP (Pipeline Parallel): sharded_load already optimizes this — it downloads only the files needed for the local pipeline stage. This is the closest to what we're proposing.
• EP (Expert Parallel, MoE): No built-in support. Each node stores and loads the full model.

For large models this is a significant constraint:

Model	FP16 Size	4-Node TP4 Disk Usage (current)	With Weight Streaming
Llama 405B	810 GB	3.2 TB (810 GB × 4)	810 GB (master only)
DeepSeek V3	670 GB	2.7 TB	670 GB
Kimi K2.5	612 GB	2.4 TB	612 GB

With mlx_lm.share demonstrating 5+ GB/s broadcast throughput over Thunderbolt 5 RDMA, streaming weights at startup is fast enough to be practical.

What Exists Today

Component	Status	What It Does
mlx_lm.share	Working	Broadcasts entire directories via all_sum at 5+ GB/s
sharded_load (PP)	Working	Downloads only the files each pipeline stage needs
sharded_load (TP)	Working but redundant	Downloads full model to every node, shards in memory
model.shard()	Working	Knows how to slice weight matrices for TP
model.pipeline()	Working	Knows which layers belong to which pipeline stage
MoE expert placement	Not implemented	No built-in expert-level distribution

Proposed Approach

Phase 1: TP Weight Streaming (Low-Hanging Fruit)

The master node (rank 0) reads each safetensor file, and instead of broadcasting the entire file, slices the weight matrices according to the TP sharding plan and broadcasts only the relevant slice to each rank.

Current TP flow:
  Every node: read full model from disk → shard() → eval (loads 1/N into RAM)
  Disk: N copies of full model

Proposed TP flow:
  Rank 0: read full model from disk → slice weights → broadcast slices via all_sum
  Rank 1-N: receive only their slices → load into RAM
  Disk: 1 copy on rank 0 only

This requires:

1. Rank 0 inspects the model's shard plan (which model.shard() already computes)
2. For each weight file, rank 0 reads it, extracts each rank's slice
3. Broadcasts each slice via all_sum (rank 0 sends real data, others send zeros — same pattern as mlx_lm.share)
4. Each rank receives only its portion and loads directly into the model

The safetensors.index.json file already maps weight names to files, and model.shard() already knows the slicing dimensions. The main engineering work is connecting these two systems.

Phase 2: PP Weight Streaming

PP is partially solved — sharded_load already knows which files each stage needs and only downloads those. The missing piece is sourcing files from a master node via RDMA instead of from disk/HuggingFace:

Current PP flow:
  Each node: download only needed files from HF → load
  Disk: partial model per node (already efficient)

Proposed PP flow:
  Rank 0: read files from disk → broadcast only the files each stage needs
  Rank 1-N: receive only their stage's files → load
  Disk: 1 copy on rank 0 only

This is simpler than TP streaming since it operates at the file level (no matrix slicing). sharded_load already computes the file-to-stage mapping.

Phase 3: MoE Expert Placement

This is the most impactful but also most complex. For MoE models (Mixtral, DeepSeek V3, Kimi K2.5), experts are independent weight blocks that can be distributed across nodes:

Proposed EP flow:
  Rank 0: read model → distribute expert subsets to each node
  Each node: holds 1/N of experts in memory
  Inference: router selects active experts → each node runs its local experts → all_sum collects results

Benefits:

• A 4-node cluster with 512 GB each can serve a model with up to ~2 TB of expert weights
• Only active experts compute per token — idle experts consume zero GPU time
• Could extend to demand-paging: evict cold experts, stream hot ones at 5 GB/s

This requires deeper integration with the model architecture — the router needs to know which experts are on which nodes, and the forward pass needs all_sum collectives after expert computation.

Phase 4: Dynamic Expert Demand-Paging (Speculative)

The most advanced version: experts are streamed on-demand during inference based on routing patterns. Hot experts stay resident, cold experts are evicted, and when a cold expert is needed, it's streamed from the master at 5 GB/s (~50ms for a typical expert block).

This is speculative and may not be practical for real-time inference, but for batch/offline workloads the latency could be acceptable.

Benefits Beyond Disk Space

1. Model switching without pre-staging — swap from one model to another by streaming from the master node. No need to copy hundreds of GB to every node first.

2. Serve models larger than any single node's storage — a node with 1 TB NVMe can participate in serving a 2 TB model if it only holds its shard.

3. Heterogeneous clusters ([BUG] Distributed inference OOMs on machines with different RAM size #1804) — nodes with different RAM sizes can receive proportional weight slices. A 64 GB node gets fewer layers/experts than a 512 GB node, with the master orchestrating the placement.

4. Faster cold start for PP — pipeline stages can begin processing as soon as their layers arrive. Stage 0 starts inference while stages 2-3 are still receiving weights.

5. Dynamic reconfiguration — switch from TP2 to TP4 by re-streaming with different slicing. No file management, no pre-staging, just a different sharding plan.

6. Single source of truth — model updates, quantization changes, and fine-tune checkpoints only need to exist on the master node.

Hardware Context

Tested on a 5-node cluster of Mac Studio M3 Ultra (512 GB each) connected via Thunderbolt 5 full mesh (4 cables per node):

• mlx_lm.share measured at 5.2 GB/s broadcast to 3 nodes simultaneously
• Our independent all_sum-based broadcast measured 3.7 GB/s (without async_eval pipelining)
• Successfully broadcast 812 GB (Llama 405B FP16) to 3 nodes with verified checksum integrity
• RDMA throughput does not degrade with additional receivers (each TB5 cable is an independent 80 Gbps link)

At these speeds, streaming a 200 GB TP4 slice takes ~38 seconds — a one-time startup cost that eliminates terabytes of redundant storage.

Questions for the Team

1. Is this direction something the MLX team is already exploring or interested in?
2. For TP weight streaming, would it make sense to extend sharded_load or create a new loading path?
3. For MoE expert placement, are there plans to add expert parallelism to MLX's distributed primitives?
4. The Group.split() proposal ([Enhancement] Group.split() support for JACCL and Ring backends (parity with NCCL #3172) #3205) would be useful for hybrid TP+EP — is that on the roadmap?

Happy to contribute implementation work if this aligns with the project's direction.

[guruswami-ai]

Prior Art: Exo's Memory-Aware Heterogeneous Placement

We investigated Exo, which is popular for running LLMs across heterogeneous hardware (different Mac models, different RAM sizes). Their approach to the "mixed cluster" problem is relevant context for this proposal.

What Exo Does

Exo implements memory-aware pipeline shard placement (src/exo/master/placement.py):

• Proportional sharding: Distributes pipeline stages across nodes proportionally to available RAM. A 192 GB node gets ~3× more layers than a 64 GB node.
• Topology-aware cycle detection: Discovers which nodes are directly connected (including RDMA links) and finds valid subgraphs that can serve a model.
• Full model download per node: Despite the smart placement, every node downloads the entire model from HuggingFace. resolve_allow_patterns() in download_utils.py has per-shard filtering code but it's disabled (return ["*"]) — the comment notes "Tensor parallel requires all files." Each node stores the full model on disk even if it only uses a subset of layers.
• RDMA for inference: Uses JACCL for tensor parallel inference (mx.distributed.init(backend="jaccl")) with full TB5 topology detection and bridge fix handling.

This is one of Exo's most popular features — users can combine any spare Apple Silicon hardware into a working inference cluster without needing identical machines.

How This Relates to Our Proposal

Aspect	Exo	This Proposal
Shard assignment	Memory-proportional pipeline placement	Similar goal, extends to TP slicing and MoE expert placement
Weight delivery	Every node downloads full model from HuggingFace	Master node broadcasts only needed slices over RDMA
RDMA usage	Inference only (TP all_sum via JACCL)	Inference + model distribution
Storage	Full model on every node	Zero per-node storage — weights stream from master, live only in RAM
Model switching	Re-download full model per node from HF	Re-stream from master (~38s for 200 GB slice over TB5)
Heterogeneous support	Yes (pipeline parallelism)	Yes (pipeline + tensor + expert parallelism)

Topology: Ring vs Full Mesh

This proposal does not require a full mesh topology. The all_sum collective that powers mlx_lm.share works on any connected topology — JACCL handles routing across hops. A simple daisy-chain / ring is sufficient.

This matters because full mesh gets expensive fast:

Nodes	Full Mesh Cables	Ring Cables
2	1	1
4	6	4
5	10	5
8	28	8

TB5 cables run $70-130+ each depending on length. A 4-node full mesh is $400-800 in cables alone; 8 nodes is impractical. Most users will daisy-chain a few Macs with one cable each — and that's fine. Ring topology throughput is lower (data hops through intermediate nodes, so effective bandwidth ≈ link_speed for broadcast), but still in the multi-GB/s range over TB5 and orders of magnitude faster than each node downloading independently from HuggingFace over the internet.

The key point: RDMA weight streaming is accessible to anyone with a single TB5 cable between two Macs, and scales gracefully as you add more nodes in a ring. Full mesh is a luxury for maximum throughput, not a requirement.

Takeaway

The demand for heterogeneous cluster support is validated by Exo's popularity. MLX already has the building blocks — RDMA broadcast via mlx_lm.share for delivery, model.shard() for TP slicing. Combining Exo-style memory-aware placement with MLX's RDMA broadcast could deliver the best of both: smart shard assignment for mixed hardware + fast local delivery that doesn't require every node to download and store the full model.

[guruswami-ai]

Update: Expert Parallelism Already In Progress (#3158)

Phase 3 of this proposal (MoE expert placement) is being actively implemented in #3158 by @0xDaizz — a full Expert Parallelism implementation with all_to_all dispatch, fused C++ primitives, and Metal GPU backends.

Their benchmarks show EP is 3.1× faster than TP for decode on MoE models (E=384, which matches Kimi K2.5's expert count). The all_to_all primitive is already merged (#3164).

This validates the core idea that MoE models benefit from distributing experts rather than sharding them uniformly. The remaining novel pieces from this proposal are:

• Phase 1 (TP weight streaming) — broadcast only needed slices from a master node, eliminating per-node storage
• Phase 2 (PP weight streaming) — same idea for pipeline stages
• Phase 4 (demand-paging) — stream hot experts on-demand (speculative, may not be practical)

Great to see EP moving forward. We'll be testing it on our 5-node TB5 cluster with Kimi K2.5 and DeepSeek V3 once the remaining sub-PRs land.