# Pipeline parallel backpropagation on Apple silicon with mlx-train > Source-anchored research for Pipeline parallel backpropagation on Apple silicon with mlx-train. Run status: COMPLETED. Parent: [Mac local LLMs: Clusters, RDMA, exo and ds4](https://llms-explorer.com/tree/mac-local-llms-clusters-rdma-exo-ds4/) · 4 facets · 16 facts · page: https://llms-explorer.com/tree/pipeline-parallel-backpropagation-on-apple-silicon-with-mlx-train/ ## Structure and components - mlx-train splits the model into contiguous transformer-layer PipelineSlice modules: rank 0 embeds token IDs, later ranks accept hidden states, and the final slice applies normalization and the language-model head before producing logits for the loss. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/MattBeton/mlx-train/blob/main/mlx_train/ppp.py) ## How it works - The published mlx-train schedule processes one batch at a time. The author leaves multi-minibatch schedules such as 1F1B and GPipe as future work because overlapping minibatches requires preserving the right activations and model state for each backward pass. (medium confidence) (verified as of 2026-10-23) — [source](https://mattbeton.com/blog/pipeline-parallel-apple-silicon.html) - This pipeline-training implementation belongs to external MattBeton/mlx-train at ee8aea37ffd3adb95bf6c03c936afb3ea79d3fe7. Its load_configure_model slices a loaded MLX-LM model into PipelineSlice when distributed="pp", and its own train loop calls train_step_pp. Native mlx_lm.lora instead calls mlx_lm.tuner.trainer.train; these sources do not make the external prototype a supported native pipeline-training option or prove current-version compatibility. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/lora.py) ## Comparisons and alternatives - Inside build_graph, nonfinal local_loss is sum(y_loc*stop_gradient(dy_loc)); differentiation seeds the local reverse pass with the received downstream cotangent. Nonzero ranks use mx.value_and_grad(local_loss,argnums=(0,1)) to obtain parameter gradients and the stage-input gradient dx, then send dx to rank-1. Rank zero in a multi-stage pipeline differentiates parameters only; the final rank differentiates masked_loss instead of a downstream-cotangent surrogate. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/MattBeton/mlx-train/ee8aea37ffd3adb95bf6c03c936afb3ea79d3fe7/mlx_train/ppp.py) - Stage allocation in load_configure_model and _rank_bounds uses per=total_layers//world_size, start=rank*per, and end=(rank+1)*per except that the last rank receives the remainder through total_layers. PipelineSlice owns that contiguous layer slice; only start=0 owns token embeddings and only end=total owns norm/output projection. build_graph assumes these rank roles and compatible (B,L,H) boundaries rather than validating arbitrary layer assignments. The pinned allocation has no explicit rejection for world_size greater than total_layers, which can produce empty slices. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/MattBeton/mlx-train/ee8aea37ffd3adb95bf6c03c936afb3ea79d3fe7/mlx_train/model.py) - write_adapters_distributed in mlx_train/model.py assembles stage adapters. It maps each slice-local layers/h index to a canonical (start+local_layer_index,parameter_tail) identity, gathers LoRA tensors neighbor-to-neighbor toward rank zero, and matches those identities to a lazy reference model's official parameter keys. Rank zero asserts no reference LoRA key is missing, saves adapters.safetensors and adapter_config.json, and all ranks synchronize. The reference-key mapping preserves model-wide layer identities instead of reusing each slice's local layer numbering. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/MattBeton/mlx-train/ee8aea37ffd3adb95bf6c03c936afb3ea79d3fe7/mlx_train/model.py) ## Facts and statements - Pipeline parallel training divides a sequential model across stages. Activations move downstream during the forward pass, then boundary gradients move upstream during backprop so each stage can compute gradients for its parameters and input. (high confidence) — [source](https://mattbeton.com/blog/pipeline-parallel-apple-silicon.html) - MLX's core mx.value_and_grad requires a scalar loss (or a tuple whose first value is scalar) and accepts argnums to select input positions for gradients. A stage can use it to differentiate with respect to both model parameters and its input activation; mlx.nn.value_and_grad is the parameter-oriented helper. (high confidence) (verified as of 2026-10-23) — [source](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.value_and_grad.html) - For a nonfinal stage, ppp.py defines a scalar local loss as the sum of the stage output multiplied by the stopped received cotangent. Differentiating this surrogate seeds the local reverse pass with that cotangent; mx.value_and_grad with argnums=(0,1) returns gradients for trainable parameters and the stage input. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/MattBeton/mlx-train/blob/main/mlx_train/ppp.py) - The blog explains that an unconstrained lazy graph may schedule a blocking backward recv before the forward send that enables it. Its single-eval proof of concept adds a forward-before-backward dependency; the large-model path instead evaluates forward and backward phases in order. In ppp.py at ee8aea37ffd3, the send expression precedes the cotangent recv expression in Python, but the _depends_like call is commented out, so expression construction alone does not guarantee their lazy execution order. The README documents two-step forward/backward evaluation; the stage code also computes an input cotangent for the prior rank. (medium confidence) (verified as of 2026-10-23) — [source](https://mattbeton.com/blog/pipeline-parallel-apple-silicon.html) - Apple's WWDC26 material documents MLX-LM's --pipeline option for distributed inference and demonstrates distributed fine-tuning by replicating the model and processing separate batches. It does not demonstrate a native pipeline-parallel backprop training interface, so inference support alone does not establish that training feature. (medium confidence) (verified as of 2026-10-23) — [source](https://developer.apple.com/videos/play/wwdc2026/233/) - Pipeline backpropagation reverses the message direction: every rank except the final rank receives the output cotangent from the next rank; the final rank instead differentiates masked_loss. The stages compute local parameter gradients, and nonzero ranks also compute and send their input-boundary cotangent to rank-1. Rank 0 has no preceding stage. (high confidence) — [source](https://github.com/MattBeton/mlx-train/blob/main/mlx_train/ppp.py) - The current repository exposes activation and gradient send/receive tokens and documents separate forward and backward evaluation phases. The project article reports that a single evaluation graph could leave ranks waiting on blocking receives long enough to time out on a large model; the staged evaluation limits each wait to the relevant communication phase. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/MattBeton/mlx-train) - The project article describes a one-batch-at-a-time schedule. GPipe uses batch splitting to overlap microbatches, and PipeDream admits multiple minibatches into the pipeline to reduce idle stalls; that comparison shows the mlx-train example's pipeline-parallel label alone does not establish high utilization or a speedup. (high confidence) — [source](https://mattbeton.com/blog/pipeline-parallel-apple-silicon.html) - build_graph(model:PipelineSlice,tokens,targets,lengths) constructs lazy stage-local forward/loss/gradient arrays and a tokens_out communication dictionary, then returns (loss,g_s,tokens_out). Noninitial ranks recv (B,L,H) bfloat16 activations from rank-1; nonfinal ranks send their (B,L,H) outputs to rank+1 and receive downstream cotangents. It does not by itself execute a forward-before-backward schedule. Asserts check nonfinal output rank and hidden width, and the dependency-like ordering call is commented out. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/MattBeton/mlx-train/ee8aea37ffd3adb95bf6c03c936afb3ea79d3fe7/mlx_train/ppp.py) - train_step_pp shifts a batch into inputs/targets, calls build_graph, and groups its y/x and dx/dy tokens into forward and backward evaluation lists. It loops stages forward with barriers to evaluate receive/send work, loops stages backward with barriers to evaluate cotangent receive/send work, then applies optimizer.update(model,g_s) and evaluates model/optimizer state. This explicit evaluation schedule supplies ordering; Python send-expression construction alone does not. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/MattBeton/mlx-train/ee8aea37ffd3adb95bf6c03c936afb3ea79d3fe7/mlx_train/train.py)