…unslothai#931/unslothai#933/others)

Merge origin/main into the streaming branch. Only unsloth_zoo/mlx/trainer.py
conflicted, in MLXTrainingConfig.__init__'s wholesale-copy detection: unslothai#922
(true grad-norm reporting + compiled global-norm clipping) replaced this
branch's _MLX_CONFIG_OPTIONAL_COPY_FIELDS-based check with its own
_appended_fields set {compile_max_variants, report_grad_norm}, and neither set
alone covers the merged config. Reconciled to the union: warmup copy-detection
now tolerates every appended field (the positional optional-copy suffix plus
the mid-order compile_max_variants). unslothai#922's report_grad_norm landed as the new
last field, so it joins _MLX_CONFIG_OPTIONAL_COPY_FIELDS to keep every
post-image_size field copy-optional and preserve the positional-copy suffix
invariant. Validation on the merged tree: 242 unit + 16 metal + 25
recorded-validation checks pass (grad-norm and streaming/prefetch loops
coexist).

The main merges left this branch's streaming config fields (max_eval_batches, streaming_text_length_window_batches, streaming_prefetch_batches) before label_smoothing_factor (unslothai#930) and report_grad_norm (unslothai#922), which each intentionally declared themselves last to preserve positional indices. A positional copy from a main-lineage config therefore bound its label_smoothing_factor / report_grad_norm values to this branch's fields, silently disabling label smoothing or failing lazy-eval validation. Move the streaming fields to the very end so every pre-existing (including main-appended) field keeps its positional index; _MLX_CONFIG_OPTIONAL_COPY_FIELDS again holds exactly this branch's trailing suffix, while the wholesale-copy warmup detection tolerates all appended fields (compile_max_variants, label_smoothing_factor, report_grad_norm) so legacy dumps missing them are still recognized.

Resolve the trainer.py conflict in the gradient-norm reporting region.

The branch carried the optimizer-state grad-norm helpers from the original
MLX path (_grad_norm_from_optimizer_state / _optimizer_v_total, unslothai#634), which
recover ||g|| from an Adam-family optimizer's second moment. main's unslothai#922
replaced that approach with true grad-norm reporting computed directly from
the gradient tree (_global_grad_norm_fp32 / report_grad_norm), which is
optimizer-agnostic and covers the rmsprop/adamax/adagrad/adadelta optimizers
this PR adds. Took main's side: the branch's call sites were already removed
by main's changes in the auto-merged region, so dropping the helper
definitions leaves no orphaned callers, and all four new optimizers still
build.

Note for @danielhanchen: your 35e8a10 helper _mlx_optimizer_state_norm_supported
and its test test_optimizer_state_norm_supported_excludes_adamax are left in
place but are now unused in production, since unslothai#922's direct grad-norm reporting
supersedes the optimizer-state path they gate. Left untouched rather than
deleting your work -- happy to remove them if you'd prefer, or keep them for a
future optimizer-state reporting path.

Resolve the trainer conflicts against the batch-plan preflight (unslothai#918) and
the true grad-norm reporting rework (unslothai#922):

- union the dataclasses import (field plus replace)
- keep main's preflight / _resolve_training_steps structure and re-derive
  _epoch_stop_total_microbatches from batches after grad_accum is known
- keep the committed/pending metric window split and fold it into main's
  single mx.eval boundary, dropping the optimizer-state grad-norm fallback
  that main removed
- keep both appended test blocks in test_mlx_trainer_internals.py

…-grouped batching, and background prefetch (#927)

* feat(mlx): lazily batch streaming text rows

* feat(mlx): define text stream replay semantics

* fix(mlx): preserve lazy text tokenization state

* fix(mlx): manage lazy stream iterator lifecycle

* fix(mlx): validate lazy text stream boundaries

* feat(mlx): prepare lazy response-masked streams

* feat(mlx): stream text eval and sized epochs

* test(mlx): cover streaming eval contracts

* feat(mlx): dispatch lazy text streams from rank zero

* feat(mlx): expose prepared streaming text parity

* feat(mlx): length-group lazy text stream batches in a bounded window

* feat(mlx): rebuild lazy VLM streaming on the text lifecycle contracts

* refactor(mlx): stage lazy text batches host-side behind one MLX finalizer

* feat(mlx): prefetch host-staged text stream batches on a producer thread

* refactor(mlx): stage VLM label placement host-side behind the finalizer

* feat(mlx): prefetch host-staged VLM stream batches on the producer thread

* fix(mlx): stage processor-owned MLX outputs opaquely for VLM prefetch

* fix(mlx): materialize opaque VLM payloads and stage prompt/completion opaquely

Lazy MLX graphs cannot be evaluated from another thread, so opaque processor-owned payloads now cross the prefetch boundary through a producer-side materialization barrier (mx.eval of the processor's own pending graphs on the thread that issued them; Unsloth still introduces no new MLX computation producer-side). Prompt/completion streams from MLX-returning processors previously rejected in producer mode; they now stage via a pc_opaque carrier whose concat/truncation/completion-mask/label decisions all run in the consumer-thread finalizer through the extracted _combine_vlm_prompt_completion_inputs helper, bit-identical to the synchronous legacy route.

* fix(mlx): reject MLX-valued text rows before producer-side parsing

Raw dataset rows and formatter results carrying MLX arrays previously reached truthiness probes and row parsing before any provenance guard, so a prefetch producer crashed with a misleading conversion error (an Unsloth-issued MLX evaluation off the consumer thread). In producer mode the lazy text row iterator now scans each row on entry and each formatter result immediately, rejecting with the actionable streaming_prefetch_batches=0 remedy before anything can touch the values. Synchronous behavior is unchanged.

* fix(mlx): carry tokenizer scalars instead of the live processor in pc_opaque

The pc_opaque carrier stored the live processor and the consumer finalizer read its tokenizer padding side and pad id, racing the producer's next processor call (the queue slot is released before finalization), which violated the producer's exclusive processor ownership. The collator now snapshots flush_side and pad_id immediately after the processor calls while the producer still exclusively owns the processor, and the combine helper consumes those immutable scalars without ever dereferencing the processor; a poisoned-tokenizer regression locks the non-dereference. Host-valued non-integer ids also get their own diagnostic instead of the misleading MLX-array provenance message, and the opaque finalize comment now describes the permitted materialization barrier instead of claiming no producer-side MLX op ran.

* test(mlx): route the P/C parity regression through the unsized producer path

The prompt/completion parity block fed a plain list, which routes to the sized synchronous branch where prefetch_batches is ignored, so the committed test never crossed the pc_opaque producer/consumer boundary it claimed to lock. The block now streams a replayable unsized source (verified to create the prefetcher), covers both completion_only_loss=None and False with label-dtype parity, and keeps the poisoned-tokenizer ownership assertion.

* fix(mlx): harden streaming eval metadata sync and iterator teardown

Rank-0 metadata resolution now synchronizes any BaseException (KeyboardInterrupt and SystemExit included) through the status collective before re-raising on rank 0, so peer ranks cannot hang in the all_sum when rank 0 is interrupted mid-resolve. Infinite-stream detection reads the private datasets attributes (ex_iterables, stopping_strategy) with getattr fallbacks so a rename in a future release degrades detection instead of failing eval for finite streams whose wrapper class matches by name. The lazy eval batch view closes its underlying iterator in a finally, releasing the pass's owned source cursors when max_eval_batches truncates or the consumer exits early.

* test(mlx): pin the DDP probe interpreter via launch argv

The ring launcher bash-execs the command array verbatim and does not apply --python to it, so the probe relied on a shebang plus chmod to select the venv interpreter. Passing sys.executable as the command's first element removes that dependency and handles interpreter paths the shebang mechanism cannot (spaces, length limits).

* fix(mlx): synchronize owner-dispatch failures for any BaseException

The rank-0 contract check and the per-batch owner pull both guarded only Exception, so a KeyboardInterrupt or SystemExit raised by the source, formatter, or tokenizer on rank 0 unwound past the dispatch collective while peer ranks blocked in all_sum indefinitely. Both handlers now capture BaseException, broadcast the failure status through the collective, and re-raise the original error on rank 0 — the same contract the rank-0 metadata resolver already follows.

* fix(mlx): reject non-integer max_eval_batches values

int() coercion silently accepted fractional values and booleans (1.9 and True both became 1), changing how much eval data is scored while the error contract promises a positive integer. Non-integral values now reject with the existing actionable message; exact integral values (including integer-valued floats and numpy integers) still pass.

* fix(mlx): keep interrupted ranks inside the failure consensus

Rank-local work that feeds the rank-wide failure-status collectives caught only Exception, so a KeyboardInterrupt or SystemExit on one rank unwound past the consensus while surviving ranks blocked in it. All consensus-feeding handlers now capture BaseException: batch fetching, evaluation fetch and loss, resume fast-forward, the DDP local-step pair, best-model/checkpoint/final saves, and compiled-step setup (where ordinary failures keep the eager fallback while interrupts abort through the consensus instead of silently downgrading the run). Best-model restore gains an unconditional consensus before the diagnostics collective so an interrupt there cannot strand peers either. The consensus raiser re-raises interrupts unwrapped on the failing rank without setting stop_requested (a reused trainer must not inherit a stop request from an interrupt); peers still get the context-carrying RuntimeError, and a committed regression locks both behaviors. Single-process behavior is unchanged.

* fix(mlx): resolve completion mode from prompt/completion rows before locking

With completion_only_loss=None, a lazy stream that emitted a plain-text row before a prompt/completion row locked the inherited full-sequence mode, silently training prompt tokens, while the reverse order rejected — behavior depended on row order. A prompt/completion row now always resolves the default completion mask from its own shape before the mode locks, so the existing schema and label validators reject plain/prompt-completion mixes consistently in both orders. Homogeneous streams and explicit completion_only_loss settings are unchanged.

* fix(mlx): abort streaming eval before the first pull when stopped

A stop requested just before a scheduled evaluation (for example a step callback firing on an eval step) still entered the eval loop, which pulled the first batch before consulting the stop state. For an unsized eval source whose next row blocks, cancellation could never take effect. The eval accumulator now performs one rank-synchronized stop check before constructing the iterator, so every rank returns together without a blocking pull. A stopped evaluation therefore returns empty totals without building its iterator for sized sources too (previously it pulled one already-materialized batch before breaking); the accumulated result is identical. Non-stopped evaluation is unchanged.

* fix(mlx): drain queued prefetch batches once the producer is dead

Terminal close stopped the producer but left staged batches in the envelope queue, and the consumer generator kept the prefetcher alive until the training run's outer teardown. The queued payloads therefore survived the final model save; for VLM processors that emit opaque MLX arrays those are materialized device tensors, adding avoidable memory pressure at the peak-memory save.

Close now runs the join, orphan classification, and drain under one lock and returns this caller's own post-join observation, so a True return guarantees the queue was drained under that lock before being reported — concurrent closers cannot observe it half-drained. Every close() caller gates on that outcome rather than a separate liveness read: the terminal-save path refuses to serialize when not clean, and the run-teardown path persists a not-clean producer as an orphan in a finally so a propagating signal cannot bypass tracking. The orphaned flag is set on entry, so a signal raised anywhere in close() leaves a conservative orphan rather than a falsely clean producer; interrupts propagate (the abort is honored) rather than being swallowed. Persisted orphans that have since terminated are closed (and thus drained) when cleared, releasing their queues deterministically at the next gate. The mid-training pause/resume path is untouched (it never closes).

* fix(mlx): treat horizontally-concatenated streaming eval as finite when any child is

_mlx_stream_declares_infinite classified both vertical and horizontal HF concatenations with any(child_infinite), but a horizontally-concatenated iterable advances its children in lockstep and terminates with the shortest — so it is bounded whenever any child is finite. A streaming eval built from one infinite and one finite child was therefore wrongly rejected as infinite unless max_eval_batches was set. Horizontal concat now requires all children to be infinite; vertical (sequential) concat keeps the any() rule.

* fix(mlx): keep label_smoothing_factor copy-optional so the config suffix invariant holds

The main merge added label_smoothing_factor (a post-original config field) after this branch's streaming fields, so the _MLX_CONFIG_OPTIONAL_COPY_FIELDS entries were no longer the declaration suffix and a positional copy from an older config mismapped its trailing argument (image_size). Add label_smoothing_factor to the optional-copy set — it is likewise a field added after the original public surface — restoring the invariant that every post-image_size field is copy-optional.

* fix(mlx): append streaming config fields after main's trailing fields

The main merges left this branch's streaming config fields (max_eval_batches, streaming_text_length_window_batches, streaming_prefetch_batches) before label_smoothing_factor (#930) and report_grad_norm (#922), which each intentionally declared themselves last to preserve positional indices. A positional copy from a main-lineage config therefore bound its label_smoothing_factor / report_grad_norm values to this branch's fields, silently disabling label smoothing or failing lazy-eval validation. Move the streaming fields to the very end so every pre-existing (including main-appended) field keeps its positional index; _MLX_CONFIG_OPTIONAL_COPY_FIELDS again holds exactly this branch's trailing suffix, while the wholesale-copy warmup detection tolerates all appended fields (compile_max_variants, label_smoothing_factor, report_grad_norm) so legacy dumps missing them are still recognized.

* Treat horizontally concatenated streaming eval as infinite when any child is

HorizontallyConcatenatedMultiSourcesExamplesIterable.__iter__ removes each
child from its live iterator list as that child raises StopIteration and
breaks only once the list is empty, padding the dropped columns. It therefore
ends with the longest child, not the shortest, so an axis=1 concat holding one
infinite child never terminates. Requiring every child to be infinite made the
detector return False for it, the lazy eval guard was skipped, and evaluation
ran unbounded. Cover it with a real datasets-backed test alongside the mocks.

Verified against datasets 3.6.0 and 4.3.0; the method body is unchanged from
3.4.1 through main.

* Tighten comments in the MLX streaming changes

* Tighten docstrings in the MLX streaming changes

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>