MLX-LM integration contracts for external GRPO and GSPO training

Parent: MLX-LM API and adapter contracts in external post-training integrations · Published reference · snapshot 2026-10-23 · viability 3.55 / 5 (research)

↓ Facts as markdown all context files

Source-anchored research for MLX-LM integration contracts for external GRPO and GSPO training. Run status: COMPLETED.

These notes link each claim to its source. Some citations identify archived research reports that are no longer hosted on this site. A published reference means the content is available; it does not certify independent review or accuracy.Read the editorial policy and follow the sources before relying on a claim.

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]

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]
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]
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]
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]
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]
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]

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]

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]
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]
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]
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]
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]
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]
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]
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]

Children

MLX-LM rollout generation and old-policy log-probability contract (frontier)
GRPO completion masks and group-relative advantage contract (frontier)
GSPO sequence ratio and clipping contract (frontier)
MLX-LM rollout token alignment and completion masks (frontier)
MLX-LM GSPO scoring validation (frontier)

Frontier under this node: GRPO completion masks and group-relative advantage contract, GSPO sequence ratio and clipping contract, MLX-LM GSPO scoring validation, MLX-LM rollout generation and old-policy log-probability contract, MLX-LM rollout token alignment and completion masks

← the whole tree · 3D view· how to read this page