# MLX-LM distillation loss kernels and custom training loss > Source-anchored research for MLX-LM distillation loss kernels and custom training loss. Run status: COMPLETED. Parent: [MLX and mlx-lm on Apple silicon](https://llms-explorer.com/tree/mlx-and-mlx-lm-on-apple-silicon/) · 3 facets · 15 facts · page: https://llms-explorer.com/tree/mlx-lm-distillation-loss-kernels-and-custom-training-loss/ ## How it works - The current MLX-LM tuner loss module implements `kl_div_loss` and `js_div_loss` using `mx.fast.metal_kernel` forward kernels and hand-written VJPs registered with `mx.custom_function`; it uses composed MLX loss operations when Metal is unavailable. The dynamic quantization implementation imports and calls `kl_div_loss`. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/losses.py) - In the source snapshot checked on 2026-10-23, `mlx_lm.tuner.losses` implements per-last-axis KL and JS losses with custom Metal forward and backward kernels when `can_run_metal()` is true; `kl_div_loss` otherwise normalizes both logits with log-sum-exp and calls `mlx.nn.losses.kl_div_loss` with no reduction. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/losses.py) - Independent MLX tooling documents broader distillation orchestration: MLXSmith describes offline teacher-generated responses followed by student SFT and online preference distillation using student-generated candidates, teacher scoring, and an importance-sampled loss; AutoContext describes an in-process path with student rollouts, a frozen teacher, and per-token reverse KL. These are external workflows, not evidence of a built-in MLX-LM teacher-student trainer. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/Hmbown/MLXSmith/blob/main/docs/cli/distillation.md) - trainer.default_loss uses inputs=batch[:,:-1] and targets=batch[:,1:]. It constructs one-based token positions and includes only lengths[:,0] <= position <= lengths[:,1], multiplies per-token cross-entropy by that mask, and returns the float32 sum divided by ntoks together with ntoks=mask.sum(). This is the literal supplied-boundary rule, not a guarantee that every padded target is excluded: iterate_batches records an L-token sequence with upper length L, so when that row is padded, shifted target position L (original index L, the first padding token) is still included by the inclusive upper comparison. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py) ## Comparisons and alternatives - The Metal `kl_div_loss(logits_q, logits_p)` path computes KL(softmax(logits_p) || softmax(logits_q)) and returns a gradient for `logits_q` but zeros for `logits_p`. AutoContext documents its on-policy distillation objective as reverse KL(student || teacher), so callers must account for this gradient direction instead of assuming the helper directly implements that reverse-KL objective. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/losses.py) - MLX-LM's DWQ path uses teacher logits computed on demand or optionally precomputed top-logit targets, and computes KL between the quantized model's logits and those targets. dwq_quantize explicitly unfreezes scales and biases on affine quantized layers below 8 bits, then optimizes the full model.trainable_parameters() collection; exclusivity to those scales and biases depends on which other parameters were already frozen. This is teacher-guided quantization fine-tuning rather than general student-architecture distillation. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/quant/dwq.py) - An open MLX-LM issue filed on 2026-09-17 describes DWQ as present but general structural knowledge distillation as missing, and proposes off-policy distillation followed by on-policy rollouts. The issue page showed no maintainer response when checked on 2026-10-23, so it records a user request rather than a repository roadmap commitment. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/issues/1900) - kl_div_loss(logits_q,logits_p) and js_div_loss(logits_q,logits_p) take unnormalized logits with vocabulary on the last axis; normalization is internal. Their public results omit that vocabulary axis and retain the preceding batch/token positions, so the caller supplies token masking and reduction rather than expecting a scalar training loss. The Metal custom VJPs compute the logits_q gradient and explicitly return zeros_like(logits_p) for the target gradient. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/losses.py) ## Facts and statements - `mlx_lm.tuner.trainer.train` and `evaluate` accept a `loss` callable that defaults to masked next-token cross-entropy; training differentiates the callable with `nn.value_and_grad(model, loss)`, and the callable returns a loss value and token count. This is a Python-level custom-loss hook; separate teacher loading and rollout orchestration are supplied by caller or external code. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py) - `kl_div_loss(logits_q, logits_p)` returns one loss value per position (`logits_q.shape[:-1]`); its custom VJP computes a gradient for `logits_q` and returns zeros for `logits_p`. MLX-LM's dynamic quantization sensitivity path computes targets with the original model and differentiates the quantized model against them, using the second KL input as fixed reference logits. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/losses.py) - An upstream issue opened on 2026-09-17 asks for a native general structural knowledge-distillation framework. This is scoped: MLX-LM already has low-level KL/JS loss kernels and uses teacher-reference KL in quantization sensitivity estimation, but its generic trainer hook is not itself a general teacher/student training pipeline. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/issues/1900) - MLX's `value_and_grad` pattern differentiates a scalar objective. In its pipeline-training example, Matt Beton uses `sum(output * stop_gradient(upstream_gradient))` as an auxiliary scalar for intermediate-stage backpropagation; MLX-LM's KL kernel similarly provides an explicit VJP for its custom Metal computation. (medium confidence) — [source](https://mattbeton.com/blog/pipeline-parallel-apple-silicon.html) - Independent MLX-related implementations document higher-level distillation patterns: MLXSmith separates offline teacher-generated SFT from online preference distillation, while `mlx-go-lm` exposes weighted cross-entropy/KL/JS terms and APIs for precomputing teacher logits. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/Hmbown/MLXSmith/blob/main/docs/cli/distillation.md) - A Python caller supplies train(...,loss=my_loss). train builds nn.value_and_grad(model,loss), then invokes it as loss_value_and_grad(model,*batch); the standard batch iterator yields (token_batch,lengths), so my_loss receives (model,token_batch,lengths) and must return (scalar_loss,token_count). The scalar drives differentiation; token_count supports reporting and weighted evaluation. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py) - At 0d4d3d04f8ab, the native mlx_lm.lora CLI parser exposes fine-tuning and optimizer settings but no generic teacher/student or custom-loss selector. Its train_model call invokes train without overriding loss, so the default next-token objective applies. A general distillation workflow must supply teacher loading, targets, alignment and a custom Python loss/integration; low-level KL/JS primitives or the separate DWQ path do not establish such a CLI workflow. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/lora.py)