# MLX-LM API and adapter contracts in external post-training integrations > Source-anchored research for MLX-LM API and adapter contracts in external post-training integrations. Run status: COMPLETED. Parent: [MLX and mlx-lm on Apple silicon](https://llms-explorer.com/tree/mlx-and-mlx-lm-on-apple-silicon/) · 6 facets · 33 facts · page: https://llms-explorer.com/tree/mlx-lm-api-and-adapter-contracts-in-external-post-training-integrations/ ## Definitions - Tokenizer equality is a requirement of direct token-level teacher/student comparison, not of text-only teacher-response SFT. AutoContext rejects a vocabulary mismatch for its local OPD path; a Hugging Face TRL issue records that passing student token IDs directly into a differently-tokenized teacher gives incorrect teacher probabilities. Cross-tokenizer logit training needs explicit alignment or another objective; `mlx-go-lm` exposes a token-logit remapper and counts unmapped tokens. (high confidence) (verified as of 2026-10-23) — [source](https://autocontext.ai/docs/guides/distillation-mlx) - MLX-LM's DWQ is a scoped teacher-guided path: the documentation says it tunes quantized-model parameters such as scales and biases against a non-quantized teacher, and the implementation computes KL targets while updating the quantized model. This is weight-quantization distillation, not a general task-level teacher/student training interface; the September 2026 feature request explicitly distinguishes DWQ from general knowledge distillation. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LEARNED_QUANTS.md) ## Structure and components - AutoContext maps variant=gspo to external --grpo-loss-type grpo --importance-sampling-level sequence; ordinary grpo selects token. The extension sequence option averages masked token log-ratios over each completion before exponentiation/clipping, while token returns per-token log-ratios. Its pinned variant guide labels this GSPO-style ratios and explicitly warns that the option does not reproduce every component of a complete GSPO recipe. This setting belongs to the external extension, not native fine_tune_type. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/greyhaven-ai/autocontext/f72c154152397f54d55407508b153d28861888a5/autocontext/src/autocontext/training/autoresearch/grpo_backend.py) - AutoContext OPD writes adapters/adapters.safetensors from student.trainable_parameters() and adapter_config.json containing fine_tune_type=lora, num_layers, and lora_parameters. Pinned native load(base,adapter_path=that_directory) reads those two files, recreates adapter layers, and validates allowed names and shapes before loading. This is a concrete file-contract match for the matching base/config, not a live interoperability test; a standalone weights file or MLXSmith ADAPTER.txt placeholder does not satisfy it. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/greyhaven-ai/autocontext/f72c154152397f54d55407508b153d28861888a5/autocontext/src/autocontext/training/autoresearch/on_policy_distill.py) ## How it works - mlx_lm.generate.generate_step yields each sampled token with a vocabulary-wide log-probability vector computed after logits processors but before the sampler runs. make_sampler can then apply temperature, top-p, min-p, or top-k filtering. Therefore, the returned vector does not automatically represent the filtered behavior distribution when those sampler transforms are active; an adapter should define and consistently store the behavior-policy log probability used in its importance ratios. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/generate.py) - The current MLX-LM tuner accepts a caller-supplied loss; its default loss applies shifted-token, length-masked cross-entropy, and its update path differentiates the supplied loss with nn.value_and_grad(model, loss) before calling optimizer.update(model, grad). An external GRPO or GSPO integration must provide rollout preparation and its RL objective around this hook because the shown tuner loop does not implement reward grouping or policy-ratio logic. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py) - MLX-LM stream_generate yields a token ID and a log-probability vector for each generation step, while generate returns decoded text. The streamed vector is normalized from logits after explicit logits_processors and before the sampler, so it is not guaranteed to equal the sampling-distribution log probability when the sampler changes its input distribution. An external adapter should retain token IDs and completion masks and verify that rollout and training scores use matching policy conventions. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/generate.py) - Current TRL documentation requires a custom rollout_func to return prompt_ids, completion_ids, and logprobs, with logprob_token_ids optional. Its current GSPO recipe uses importance_sampling_level="sequence" with loss_type="grpo" and says the sequence setting matters when training goes slightly off-policy; the Paper Index also says the experimental GSPO-token implementation was removed after v1.13. An adapter should pin its trainer version and keep core sequence-level GSPO distinct from the removed experimental path. (medium confidence) (verified as of 2026-10-23) — [source](https://huggingface.co/docs/trl/grpo_trainer) - AutoContext run_grpo_training at f72c1541 launches python -m mlx_lm_lora.train, not native mlx_lm.lora. It writes train/valid JSONL and autocontext_reward.py; the registered reward accepts prompts=None, completions=None, **kwargs, reads singular answer (or answers), and forwards them to score_completions. For agent-task scenarios, evaluate_output receives the decoded per-completion answer state when available, otherwise shared_state. Game scenarios instead call execute_match(strategy,seed=seed) without passing that decoded state. Throwing verifiers score zero. The wrapper passes the reward file/name and selected GRPO-family flags to the external process. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/greyhaven-ai/autocontext/f72c154152397f54d55407508b153d28861888a5/autocontext/src/autocontext/training/autoresearch/grpo_backend.py) - The extension entry point loads the reference separately from reference_model_path or the base model path and calls freeze(); policy adapters belong to the distinct trainable model. grpo_loss scores the exact prompt-conditioned completion IDs with policy and reference, stop_gradient on reference scores, and skips reference scoring when beta=0 or no reference exists. In this implemented once-per-rollout objective, log_ratio=logps-stop_gradient(logps) has forward value zero but retains the policy gradient: reference probabilities are only for KL, not pi_old. Native streamed pre-sampler logprobs are therefore not a substitute for externally defined behavior probabilities when rollouts are reused or sampler filters change the distribution. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/Goekdeniz-Guelmez/mlx-lm-lora/311058c7a6a1d3c62ddd8cdf2c795bbb6d0bd632/mlx_lm_lora/train.py) - Current external integrations separate offline teacher-text SFT from online token-level distillation. AutoContext describes its `mlxlm` route as SFT over a fixed teacher-written dataset and runs OPD in a separate in-process loop with a frozen teacher and LoRA student; MLXSmith likewise separates generated-response SFT from its OPD scoring mode. (high confidence) (verified as of 2026-10-23) — [source](https://autocontext.ai/docs/guides/distillation-mlx) ## Problems, failure modes and limitations - Token-level KL cannot compare raw teacher and student logits as if token IDs shared meaning when their tokenizers or vocabulary sizes differ. In TRL issue 4562, the maintainer says this constraint follows from GKD/MiniLLM's token-level objective; TRL's experimental GOLD trainer instead aligns text spans and uses ULD for cross-tokenizer cases. An MLX-LM integration must preserve a shared token coordinate system or implement an explicit alignment method; passing student token IDs directly to a different-tokenizer teacher is unsafe. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/huggingface/trl/issues/4562) ## Comparisons and alternatives - GRPO uses online groups of sampled completions for each prompt, scores the completions, derives group-relative advantages, and optimizes token-level current-versus-old policy ratios; implementations may also include a reference-policy KL term. The rollout adapter therefore needs to preserve each prompt's group membership and align rewards and advantages with the corresponding completion tokens. (high confidence) — [source](https://arxiv.org/abs/2402.03300) - To adapt an external distillation objective to the pinned native trainer, supply train(student,optimizer,...,loss=my_loss,iterate_batches=my_batches): standard my_loss receives (model,token_batch,lengths) and returns (scalar_loss,token_count). Build an explicit completion/padding mask and token denominator rather than assuming the inclusive default_loss boundary excludes every padded target. The native loop differentiates only the supplied student module, sums gradients, averages across workers, divides by args.grad_accumulation_steps on updates, and calls the supplied optimizer; values below one are rejected. AutoContext instead calls nn.value_and_grad(student,loss_fn) in its own loop with teacher targets stopped. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/tuner/trainer.py) - Native --resume-adapter-file calls model.load_weights(...,strict=False), then builds a new optimizer and starts the normal loop; it does not restore teacher identity/weights, optimizer moments, scheduler progress, RNG state, or rollout state. AutoContext OPD saves adapter tensors/config rather than those states, while MLXSmith may apply a source adapter but initializes a fresh optimizer and random.Random(seed). Exact external-training continuation therefore needs a separate checkpoint for that omitted state and resolved teacher/student revisions. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/ml-explore/mlx-lm/0d4d3d04f8ab7ce4a51e7209e8615755a48ba94d/mlx_lm/lora.py) ## Facts and statements - On the upstream main snapshot checked on 2026-10-23, mlx_lm.tuner.trainer.train accepts custom loss and iterate_batches functions, computes gradients with nn.value_and_grad(model, loss), and applies them through optimizer.update(model, grad). Its signature has no named rollout-generation or reward callback, so an online GRPO/GSPO adapter must supply that logic through custom batching and loss code or an outer loop. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py) - GRPO samples a group of G completions for each prompt and computes advantages from rewards relative to that group. Its clipped policy update uses per-token current-policy to rollout-policy probability ratios; implementations may also include a reference-policy KL term. If an adapter reuses rollouts for multiple optimizer updates, it must retain the rollout log probabilities or a snapshot that can recompute them. (high confidence) — [source](https://arxiv.org/abs/2402.03300) - GSPO defines an importance ratio from the current-to-old likelihood of a complete response, raised to the inverse of its valid completion length, and applies clipping and optimization at sequence level. An adapter must aggregate the per-token log-probability differences over the valid completion span before computing the GSPO ratio; reusing token-level GRPO clipping changes the objective. (high confidence) — [source](https://arxiv.org/abs/2507.18071) - MLX-LM's stock iterate_batches pads token rows with zeros and returns offset/end length boundaries; default_loss uses those boundaries for its token mask. An RL custom loss needs explicit completion boundaries and a denominator over valid completion tokens, or prompt and padding positions can enter the policy objective. The adapter should also decide explicitly how to handle truncated completions. (medium confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py) - A separate MLX ecosystem guide lists GRPO and GSPO for the mlx-lm-lora fork but labels its mode and flag descriptions as unverified README claims. Soup's current training guide lists grpo_variant: gspo, yet its native backend: mlx page rejects GRPO and says that backend supports only SFT. These examples show support is project- and backend-specific; they do not define an upstream mlx-lm integration contract. (medium confidence) (verified as of 2026-10-23) — [source](https://hbmartin.github.io/Foundation-Models-and-Core-AI-and-MLX-skills/part-12-mlx-python/references/06-finetuning-and-porting-models/) - In mlx-lm-lora generate_grpo at 311058c7, each prompt index is repeated group_size times, BatchGenerator.insert assigns rollout UIDs, and sampled IDs are collected by UID before returning completions/texts/batch_indices in insertion order. It uses native make_sampler(temperature,top_p,min_p,top_k), EOS stop sequences plus an optional encoded answer delimiter, and bounds completion/prefill concurrency separately from group expansion. EOS remains in policy token IDs but is removed from reward text; generator.close() and restoration of the model training mode run in finally. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/Goekdeniz-Guelmez/mlx-lm-lora/311058c7a6a1d3c62ddd8cdf2c795bbb6d0bd632/mlx_lm_lora/trainer/grpo_trainer.py) - The external reward callback receives prompts, completions, answer (singular), and types, each expanded with batch_indices so rewards map to the originating prompt group; it must return one reward per completion. Group-relative advantages are calculated from the weighted rewards. train_grpo passes exact completion IDs/texts, batch_indices, advantages and reward_metrics to grpo_loss. Its completion mask starts at prompt_length-1 and spans exactly completion_length shifted targets, excluding right padding; grpo_loss returns (loss,valid_token_count,metrics). This triple and grouped metadata require an explicit wrapper/custom-batch design to fit native trainer.train, whose standard custom loss receives (model,batch,lengths) and returns only (loss,token_count). (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/Goekdeniz-Guelmez/mlx-lm-lora/311058c7a6a1d3c62ddd8cdf2c795bbb6d0bd632/mlx_lm_lora/trainer/grpo_trainer.py) - For standard LoRA, extension from_pretrained creates adapter_config.json with lora_parameters and num_layers; train_grpo writes trainable tensors to adapters.safetensors and numbered weight snapshots. Native load(...,adapter_path=...) expects those files, defaults missing fine_tune_type to lora, and requires compatible base-model names/shapes. Extension --resume-adapter-file only loads weights before initializing the optimizer. These adapter artifacts do not restore optimizer state, schedule position, RNG, reward-module/scenario state, frozen-reference revision, rollout groups, or stored behavior scores; preserve those separately for exact continuation. No artifact load was executed here. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/Goekdeniz-Guelmez/mlx-lm-lora/311058c7a6a1d3c62ddd8cdf2c795bbb6d0bd632/mlx_lm_lora/utils.py) - The documented MLX-LM LoRA CLI consumes JSONL chat, tools, completions, or text records. Its `--mask-prompt` option computes loss on completion tokens for chat and completion datasets; for chat data the final message is treated as the completion. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md) - The Python `mlx_lm.tuner.trainer.train` function accepts a loss callable, differentiates that callable with `nn.value_and_grad`, and updates the supplied optimizer. External code can provide a student-logit objective at this lower-level seam; the stock LoRA CLI contract remains supervised fine-tuning. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py) - Hosted teacher APIs can feed the offline text path: AutoContext documents API or hosted teacher trace collection before training, while the MedTech and Distill-SQL projects describe teacher-generated reasoning or labels rendered as training records and then learned with MLX-LM LoRA. (medium confidence) (verified as of 2026-10-23) — [source](https://autocontext.ai/docs/guides/distillation-mlx) - `mlx-go-lm` documents a teacher-logit cache contract that writes precomputed targets in the same order used by training and streams cached logits batch by batch. An external MLX-LM integration that adopts cached logits therefore needs to preserve sample order and vocabulary coordinates across precompute and student training. (high confidence) (verified as of 2026-10-23) — [source](https://pkg.go.dev/github.com/tmc/mlx-go-lm/mlxlm/llm/training) - The documented `mlx_lm.lora` data contract accepts local JSONL with `chat`, `tools`, `completions`, or `text` records, one example per line. An external teacher's generated response can enter the student's ordinary SFT path after the caller stores it as a prompt/completion or message record. `--mask-prompt` supports completion-only loss for chat and completion datasets; the documented default computes loss for every token. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md) - The Python `mlx_lm.tuner.trainer.train` contract accepts caller-supplied `loss` and `iterate_batches` functions plus a `training_callback` object implementing report methods. Its default loss shifts the token batch, computes masked next-token cross-entropy, and returns loss plus a token count. The current signature has no teacher-model argument, so direct logit distillation needs a custom loss and batch path or a separate training loop. (high confidence) (verified as of 2026-10-23) — [source](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/trainer.py) - A hosted teacher that supplies response text gives supervised examples, not token-probability targets. AutoContext documents provider-generated rationales as SFT data and says token-level distillation needs local or open-weight teacher logits; MLX-LM's JSONL loader and custom-loss hooks provide the student-side data and training extension points for those two distinct paths. (high confidence) (verified as of 2026-10-23) — [source](https://autocontext.ai/docs/guides/distillation-mlx) - AutoContext run_on_policy_distillation at f72c1541 imports mlx_lm.load and linear_to_lora_layers. It loads/freeze()s the teacher, obtains student and tokenizer from load(student_model), freezes the student base, checks _model_logit_vocab(student/teacher) via a one-token forward, then installs LoRA on num_layers with rank=8, scale=20.0, dropout=0.0. Prompts use the student tokenizer chat template. This validates output-logit dimension equality; it does not validate identical token-to-ID mappings between tokenizers. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/greyhaven-ai/autocontext/f72c154152397f54d55407508b153d28861888a5/autocontext/src/autocontext/training/autoresearch/on_policy_distill.py) - AutoContext full_kl OPD feeds [B,T,V] student and teacher raw logits into its own reverse_kl_per_token, applies temperature-scaled log-softmax internally, stop_gradient to teacher logits, and computes KL(student||teacher). Its fixed-length student rollout has no early EOS and masks prediction positions P-1 through P+max_tokens-2; normalization uses the mask sum with an epsilon floor. Native kl_div_loss/js_div_loss also accept raw logits and return per-position losses, but substituting them requires matching the intended divergence orientation and completion mask, not merely matching tensor shapes. (medium confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/greyhaven-ai/autocontext/f72c154152397f54d55407508b153d28861888a5/autocontext/src/autocontext/training/autoresearch/on_policy_distill.py) - MLXSmith run_distill at 3737578f uses a different OPD route: the student generates token IDs and behavior logprobs, the teacher scores that completion (or supplies a sequence-average fallback), and importance_sampling_loss uses a teacher-minus-behavior advantage. It aligns/truncates completion scores and accumulates gradients over cfg.train.grad_accum before applying averaged student gradients. Its offline branch writes teacher text then calls run_sft. If loading raises BackendNotAvailable, it writes ADAPTER.txt and returns; that placeholder is not a trained native adapter. (low confidence) (verified as of 2026-10-23) — [source](https://raw.githubusercontent.com/Hmbown/MLXSmith/3737578f24c55ccbd56acfa41e7dab5b63a7a63b/src/mlxsmith/train/distill.py)