<!-- llms-explorer concept facts · https://llms-explorer.com/tree/mlx-lm-distributed-fine-tuning-and-gradient-synchronization/ · pack 2026-10-23 · ~1787 tokens -->

# MLX-LM distributed fine-tuning and gradient synchronization

> Source-anchored research for MLX-LM distributed fine-tuning and gradient synchronization. Run status: COMPLETED.

Parent: [MLX-LM DoRA and distributed training extensions](https://llms-explorer.com/tree/mlx-lm-dora-and-distributed-training-extensions/) · 5 facets · 15 facts · page: https://llms-explorer.com/tree/mlx-lm-distributed-fine-tuning-and-gradient-synchronization/

## 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](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py)

## 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](https://developer.apple.com/videos/play/wwdc2026/233/)
- 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](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py)
- 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](https://developer.apple.com/videos/play/wwdc2026/233/)

## 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](https://arxiv.org/abs/2609.18066)
- 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](https://arxiv.org/abs/2609.18066)

## 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](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py)

## 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](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py)
- 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](https://developer.apple.com/videos/play/wwdc2026/233/)
- 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](https://github.com/ml-explore/mlx/issues/4475)
- `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](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py)
- 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](https://github.com/ml-explore/mlx/blob/main/docs/src/usage/distributed.rst)
- 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](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py)
- 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](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py)
- 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](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py)
