MLX-LM distributed fine-tuning and gradient synchronization
Parent: MLX-LM DoRA and distributed training extensions · Published reference · snapshot 2026-10-23 · viability 4.25 / 5 (research)
↓ Facts as markdownall context files
Source-anchored research for MLX-LM distributed fine-tuning and gradient synchronization. Run status: COMPLETED.
These notes link each claim to its source. Some citations identify archived research reports that are no longer hosted on this site. A published reference means the content is available; it does not certify independent review or accuracy.Read the editorial policy and follow the sources before relying on a claim.
Structure and components
- The v3 research paper describes a non-overlapped data-parallel baseline with gradient all-reduce after backpropagation. Its modified JACCL path queues gradients as layers finish and uses a separate CPU consumer thread to reduce ready gradients while GPU backpropagation continues. This is the paper's research implementation; the stock MLX-LM trainer source alone does not establish this overlap. (medium confidence) (verified as of 2026-10-06) [source]
How it works
- The standard distributed fine-tuning arrangement is data parallel: ranks keep model replicas, process different data batches, and average gradients so replicas can apply matching updates; Apple's WWDC26 walkthrough and an independent multi-Mac guide describe this workflow. (high confidence) (verified as of 2026-10-23) [source]
- The MLX-LM trainer adds local gradient trees across accumulation microbatches, synchronizes them with `average_gradients` only when an update is due, divides by the accumulation-step count, and then updates the optimizer; MLX's distributed guide recommends `average_gradients` because it batches many small communications into a larger one. (high confidence) (verified as of 2026-10-23) [source]
- Distributed fine-tuning and model-parallel inference are separate paths: Apple's WWDC26 walkthrough presents `mlx_lm.lora` fine-tuning as data parallel, while Matt Beton describes pipeline-parallel training as a separate `mlx-train` implementation with custom stage-boundary gradient flow. (high confidence) [source]
Measurements and reference values
- The October 4, 2026 v3 paper reports 936 tokens per second for four-node Qwen3-9B fine-tuning at sequence length 17,408, or 3.5x its single-node throughput; this is a weak-scaling result for that M3 Ultra setup with local batch size one, not a general speedup guarantee. (low confidence) (verified as of 2026-10-06) [source]
- A 2026 research paper describes extending Apple's distributed fine-tuning implementation with multi-trunk communication, persistent worker threads, and CPU-side gradient overlap; for a four-node Mac Studio Qwen3-9B run, its authors report up to 1.6× weak-scaling throughput over a single-trunk, non-overlapped baseline. (low confidence) (verified as of 2026-10-04) [source]
Comparisons and alternatives
- At training report time, train() all-sums each rank's accumulated losses on the CPU stream and divides by steps*world_size. It separately all-sums n_tokens, computes tokens_per_second from that global count divided by local measured train_time, and adds the global token count to trained_tokens. This training-loss report averages per-step losses across ranks rather than using the token-weighted validation formula. (low confidence) (verified as of 2026-10-23) [source]
Facts and statements
- The current trainer treats `batch_size` as the distributed global batch: it partitions examples by rank and creates each worker's local batch with `batch_size // world_size`, rejecting batch sizes that are not divisible by the worker count. Apple's walkthrough accordingly scales `--batch-size` with the device count to preserve the same per-worker batch. (high confidence) (verified as of 2026-10-23) [source]
- Apple's WWDC26 walkthrough launches distributed fine-tuning with `mlx.launch` and `mlx_lm.lora`; an independent community repository also documents multi-node `mlx-lm` LoRA runs through `mlx.launch --hostfile`, with Ring and JACCL setup examples. (medium confidence) (verified as of 2026-10-23) [source]
- An issue report opened 2026-09-08 records intermittent failures in `mlx==0.32.2` ring-backend gradient reductions at three or more ranks, including CPU-stream deadlocks and a GPU-stream watchdog timeout; it documents a reporter's reproduction and does not establish behavior for other releases or backends. (low confidence) (verified as of 2026-09-08) [source]
- `mlx_lm.tuner.trainer.train()` initializes an MLX distributed group, computes local gradients with `nn.value_and_grad`, calls `average_gradients` on update steps, and passes the synchronized gradient tree to the optimizer. (high confidence) [source]
- The captured MLX distributed guide describes a closed TCP ring and lists JACCL, MPI and NCCL backends. It says JACCL requires a fully connected mesh, while the captured WWDC26 transcript says JACCL supports both mesh and ring topologies. These sources disagree about JACCL topology support; neither establishes an unqualified restriction across versions. (high confidence) [source]
- iterate_batches first sorts dataset indices by sequence length. For each complete global batch starting at i, rank r of W workers takes idx[i+r:i+r+batch_size:W]; it then permutes batch order, using the supplied seed if present. It allocates batch_size//W local rows and omits the incomplete trailing global batch. (low confidence) (verified as of 2026-10-23) [source]
- Only rank zero writes periodic adapter snapshots and the final adapter weights. Periodic saves write both args.adapter_file and a seven-digit iteration-named *_adapters.safetensors checkpoint from model.trainable_parameters(); the final save writes args.adapter_file. (low confidence) (verified as of 2026-10-23) [source]
- evaluate() accumulates loss*tokens and token counts for local validation batches. It all-sums both accumulators across workers on the CPU stream and returns total_weighted_loss/total_tokens, giving a token-weighted validation loss across the distributed evaluation. (low confidence) (verified as of 2026-10-23) [source]
Children
- MLX distributed Ring and JACCL communication backends (frontier)
- Pipeline and tensor parallel training (frontier)
- MLX collective backends and topology selection (frontier)
- Data-parallel versus pipeline-parallel backpropagation (frontier)
Frontier under this node: Data-parallel versus pipeline-parallel backpropagation, MLX collective backends and topology selection, MLX distributed Ring and JACCL communication backends, Pipeline and tensor parallel training