# MLX-LM DoRA and distributed training extensions > Source-anchored research for MLX-LM DoRA and distributed training extensions. Run status: COMPLETED. Parent: [MLX and mlx-lm on Apple silicon](https://llms-explorer.com/tree/mlx-and-mlx-lm-on-apple-silicon/) · 5 facets · 27 facts · page: https://llms-explorer.com/tree/mlx-lm-dora-and-distributed-training-extensions/ ## 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) - As of 2026-10-23, MLX-LM exposes fine_tune_type=dora and maps nn.Linear and nn.QuantizedLinear layers to DoRALinear; its dispatcher raises for SwitchLinear and QuantizedSwitchLinear, so native DoRA conversion does not cover those MoE layer types. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/utils.py) - For an nn.QuantizedLinear base, DoRALinear.from_base recovers the logical input dimension from the packed weight width using 32 // bits; DoRALinear then dequantizes weight, scales, and biases using the layer's group_size, bits, and mode for its forward calculation. (high confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/tuner/dora.py) - MLX-LM exposes `--fine-tune-type dora` and maps nn.Linear and nn.QuantizedLinear targets to DoRALinear; its DoRALinear.fuse merges the low-rank delta and magnitude scaling, then requantizes a quantized layer with its original group_size, bits, and mode unless dequantize=True. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md) - The fine-tuning entry point calls model.freeze() before converting selected layers to DoRA. The new DoRA wrapper introduces trainable lora_a, lora_b and magnitude m while retaining the frozen pretrained base layer; full fine-tuning instead explicitly unfreezes selected base layers. (high confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/lora.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) - For ordinary LoRA on a quantized base, MLX-LM retains the wrapped base layer and calls self.linear(x) before adding the low-rank branch; its guide calls LoRA training on a quantized model QLoRA. This preserves the base module's own forward implementation, unlike DoRALinear's explicit full-weight dequantization path. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/lora.py) - The adapter converter routes nn.QuantizedEmbedding to DoRAEmbedding, but DoRAEmbedding.from_base explicitly checks for nn.QuantizedLinear rather than nn.QuantizedEmbedding; the independent guide flags that the packed embedding path can therefore miss the intended guard and fail later on shape mismatch, so quantized embedding targets need version-specific verification. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/utils.py) - DoRAEmbedding.from_base contains the exact error message "DoRAEmbedding does not yet support quantization." Its guard tests nn.QuantizedLinear rather than nn.QuantizedEmbedding, although the converter selects DoRAEmbedding for nn.QuantizedEmbedding. Do not interpret the message as proof that every quantized embedding is rejected cleanly. (high confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/dora.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) - DoRA reparameterizes pretrained weights into direction and magnitude, applying a low-rank LoRA update to direction while learning magnitude separately; MLX-LM's DoRALinear implements this with lora_a, lora_b, and a trainable per-output magnitude vector m. (high confidence) — [source](https://arxiv.org/abs/2402.09353) - For a QuantizedLinear base, MLX-LM's DoRALinear forward calls _dequantized_weight(), computes the base output with a dense x @ w.T, constructs the full adapted weight matrix, and computes its row norms on each call. A quantized base is therefore accepted, but the unmerged DoRA forward path does not preserve the base QuantizedLinear forward path; an independent practitioner report describes the same dequantization and dense-normalization behavior. (high confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/tuner/dora.py) - A practitioner report dated 2026-08-08 measured 4.1 tokens/s for one DoRA adapter on a 4-bit Qwen base on an M5 Max and about 33 tokens/s after fusing the adapter or switching to LoRA. These are setup-specific measurements, not a portable speed ratio. (low confidence) (verified as of 2026-08-08) — [source](https://justinbarias.io/blog/notes-from-fine-tuning-a-tekken-coach/) - In MLX-LM's DoRALinear, the learned base magnitude is stored per output row; each forward pass adds the scaled low-rank update, computes the adapted row norms with stop_gradient, and rescales the output by the ratio of base magnitude to adapted norm. (high confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/main/mlx_lm/tuner/dora.py) - DoRALinear uses scale directly: z=(dropout(x) @ lora_a) @ lora_b and out=x @ w.T + scale*z before magnitude normalization. The adapted weight is w + (scale*lora_b.T) @ lora_a.T. This code does not divide scale by rank; its quantized base w comes from mx.dequantize. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/dora.py)