Reviewed this end to end. Merging as is; no code changes needed from me. Four things below I would like you to address in the description or a follow-up, none of them blocking.

Method: MLX runs on Linux CPU (pip install "mlx[cpu]" mlx-lm mlx-vlm), so I executed the merge base (7ad13f6f) and head (a45ab3b7) against real mx.compile and real HF processors, plus a macos-14 run from a staging replica.

The bug is real and the fix works

Counting actual mx.compile traces on a long-tail schedule of 1000 rows, widths 44 to 1043:

Head reports action=exact reason=schedule_within_cap raw_sig=32 planned_sig=32. The negative control is the important one: with the plan hunk out, head reproduces the unbounded growth exactly like base, so this is not an invented issue.

The mechanism checks out at the source. mlx/backend/metal/device_info.cpp reads sysctlbyname("iogpu.rsrc_limit", ...) and falls back to rsrc_limit = 499000, and mlx/compile.cpp keys the cache on per-input (ndim, shape, dtype) with entries.push_back and no eviction and no size bound, each entry retaining its inputs / outputs / tape. One caveat worth a word in the description: 499000 is MLX's own conservative fallback, not a hard Metal limit, so "the Metal 499,000-resource limit" overstates it slightly.

No correctness regressions

• Real SmolLM-135M-4bit LoRA, 6 steps, real weights, base versus head: bit identical. Loss 3.563464085261027, per-step [3.75, 3.67708325, 3.89930558, 3.67567563, 3.41145825, 2.96726179], 235 tokens, compile_enabled=true both sides.
• Planned versus unplanned VLM trainer runs across grad accum 1/2/4 by batch size 1/2, with and without CCE and gradient checkpointing: 9 of 9 pass at LOSS_REL_DIFF=0.000e+00 with equal trained_tokens.
• create_vlm_batches eager wrapper across 14 configurations, and _prepare_vlm_batch_for_compile(phase=None) across all ten model-type branches: byte-identical digests.
• Whole MLX suite in a single process: head 39 failed / 573 passed, base 39 failed / 570 passed, identical failure set both ways, all pre-existing torch-shim contamination. Per file, 32 files, zero failures on either side; the four changed test files are strictly additive.
• 300-config randomized fuzz over rows, batch size, grad accum, cap, mode, sidecars, model types and orders: 297 planned, 3 declined, zero failures.
• pyproject.toml untouched, and zero torch / transformers / trl references anywhere in the diff. No mlx-lm or mlx-vlm API is touched. No Studio assets. No platform.system(), path joins, symlinks, locking or multiprocessing changes, so Linux / macOS / Windows / WSL are unaffected.

Four things for the description

1. This does change compile eligibility, and the inline comment says otherwise. trainer.py:827 says "Compilation is disabled later for this configuration; do not pay the survey", but compile_ineligible_global_norm does not exist at merge base at all. Measured: VLM with max_grad_norm=1.0 and gradient_accumulation_steps=2 was compile_enabled=True with 8 compiled calls at base, and is compile_enabled=False with 0 calls at head. Losses stay bit identical so it is perf only, and max_grad_norm defaults to 0.0 and Studio's MLX worker sets 0.0, so it is off by default. But it is a regression rather than pre-existing behaviour and the comment should say so.

2. Left-padding processors get no compression. The real Qwen2VLProcessor reports padding_side: left, so the mask-contiguity guard correctly declines those batches. Correct and safe, but it means Qwen2-VL family models at per_device_train_batch_size > 1 see zero signature reduction, which is the second reporter's configuration in [Bug] Metal buffer leak in Studio training loop → [metal::malloc] Resource limit (499000) exceeded at ~step 20 unsloth#6990. At batch size 1 it works well: 10 raw widths (32 to 87) collapse to 5 signatures, action=exact. Worth stating as a known coverage limit.

3. "Adds no padding below the ceiling" is measured against canonical widths, not raw ones. planned_event_widths() always rounds through _finite_text_pad_width(pad_to_multiple=32), so padding_work_fraction reports 0.0 while real added token work is +2.99 percent on 1000 rows (44 to 1043) and +14.89 percent on 100 rows (44 to 143). Suggest reporting raw-relative padding, or rewording.

4. The per-epoch data-pipeline multiplier is not disclosed, and it applies with compile=False. Processor calls over 4 batches: base 4 at 4 steps and 4 at 12 steps (3 epochs); head eager 4 and 12; head compiled 8 and 16. Related, formatting_func is now memoized (12 calls at base, 4 at head over 3 epochs), which is identical for deterministic functions but changes behaviour for randomizing or augmenting ones. Also worth a line: the plan retains dataset[idx] source rows (decoded full-resolution PIL images for HF Image features) for the whole run, where base retained materialized pixel_values instead.

One more, minor: the headline resource numbers were sampled via MLX's private allocator counter, which has no Python binding, so they cannot be reproduced through any public API. The compile-trace count above is reproducible and makes the same point.

Review thread

Thumbs down on the pad-token thread. The premise is right in isolation and I reproduced it: with pad_token_id=None a fixed-width schedule degrades to eager/vlm_pad_token_unavailable, or aborts under strict. But it is unreachable, because utils.py:4890 passes padding=True unconditionally to the processor, and transformers raises "Asking to pad but the tokenizer does not have a padding token" in tokenization_utils_base.py for a single sequence, equal-length pairs and unequal pairs alike. I confirmed end to end that a pad-token-less processor raises that ValueError in plan.materialize(0) at both batch size 1 and 2. So the "fixed-width, no padding needed" schedule the comment describes cannot be constructed, the repo already treats this as fatal elsewhere in _vlm_pad_token_id, and even if it were reachable the consequence is an eager fallback or a strict-mode abort, not wrong data.

I also went through the Copilot threads and rejected all four: the while width in forbidden loop is bounded by the size of forbidden and endpoints stay at or above every member's raw width (verified with forbidden=[1,33], all 10 planned, no collision); labels.tolist()[0] is unreachable because _to_mx_vlm_batch reshapes 1-D input_ids to (1,-1), and a deliberately 1-D processor returns ([0,1], 0) identically on base and head; the np.asarray(attention_mask) sync is on an already host-materialized array; and the type-tag cache grows only with distinct type objects.

@codex review

Second round adjudicated. All three rejected, verdict unchanged at MERGE, still no code changes needed from me.

Exact-width plans without a pad token. This restates the earlier thread with a new framing (an exact plan needs no padding, so it should not need a pad id) but no new reachable path. The framing does not hold, because transformers rejects on the padding strategy, not on whether padding is actually needed. Measured on real GPT-2 tokenization with the pad token removed:

padding=True   single sequence       -> ValueError: Asking to pad but the tokenizer does not have a padding token
padding=True   two EQUAL-length seqs -> ValueError: (same)
padding=True   two ragged seqs       -> ValueError: (same)
padding=False  all three             -> ok

The equal-length case is exactly the scenario the comment is built on, and it still raises. Since _processor_vlm_inputs sends padding=True unconditionally at utils.py:4890, a run with pad_token_id is None cannot materialize any batch at all, so exact-plan versus eager is moot: it dies at the first batch either way with the standard actionable transformers message. I also confirmed this on a real Idefics3Processor (SmolVLM-500M) at batch size 1 and 2. Relaxing the guard would only change behaviour for a processor that silently ignores padding=True, which no real one does.

List-backed position metadata. The mechanism is real but the path is not reachable. Three separate checks: grep -rn position_ids transformers/models --include='processing_*.py' on 5.14.1 returns zero hits, so no HF processor emits it (gemma4's image_position_ids is a different key and a torch stack); ProcessorMixin.__call__ wraps output in BatchFeature(data, tensor_type=return_tensors) and _is_tensor_or_array_like recurses into lists, so under return_tensors="np" a list position_ids becomes an ndarray and then an mx.array (measured: [[0,1,2,3]] -> ndarray (1,4) -> mlx (1,4)); and no real processor lists a position-like key in skip_tensor_conversion. The one real emitter is mlx-vlm's hunyuan_vl, which is not in _VLM_POSITION_GENERATING_MODEL_TYPES and therefore hits this exact traversal, and head already handles it: disposable_keys=set(), survey width=8 padable=False forbidden=[1,4,8], declined. Reproducing the synthetic case needs a processor that invents a key no HF processor produces and also violates the return_tensors contract.

Formatting per scheduled occurrence. The description of the diff is accurate but the framing as a bug is not. Measured over 4 unique rows and 3 epochs:

base, response_mask_fn=None : 12 calls / 12 batches
base, response_mask_fn=set  :  4 calls / 12 batches
head, either                :  4 calls / 12 batches

So the base was already split-brained, and the branch that already deduped is the train_on_responses_only path, which is the dominant Unsloth VLM configuration. TRL, the reference here, calls dataset.map(_func, batched=False, ...) once in _prepare_dataset during __init__, with datasets fingerprint-caching the result to disk across runs, so per-epoch re-invocation is not a semantic TRL offers either. What a user loses is confined to the intersection of a stochastic or stateful formatter, an index scheduled more than once, and response_mask_fn is None: the same row yields the same text each epoch. Nothing is dropped, truncated or mis-argued, batch counts are identical at 12 versus 12, and deterministic formatters are bit identical, consistent with the LOSS_REL_DIFF=0.000e+00 from round one. The suggested remedy would also be a design regression: storing per occurrence expands one row into N plan rows across N epochs, which defeats the plan's compaction, grows retained memory linearly in epochs, and reintroduces exactly the per-occurrence width variability this PR exists to remove. Worth a line in the description, which I already asked for, rather than a code change.

The red MLX check on this PR is not yours

tests/test_mlx_batching_and_decay.py::test_compiler_review_guards_are_present fails with TypeError: unsupported operand type(s) for |: '_Noop' and 'type' on macos-14. I confirmed it fails identically on a branch that is upstream main plus nothing but an unrelated CI fix:

tests/test_mlx_batching_and_decay.py::test_compiler_review_guards_are_present FAILED [ 81%]
E   TypeError: unsupported operand type(s) for |: '_Noop' and 'type'
1 failed, 10 passed

The test came in with #684 and lives on main, and this PR touches neither it nor unsloth_zoo/compiler.py where _Noop is defined. It also fails on Linux for a different reason (ModuleNotFoundError: No module named 'triton'), so it really wants a CUDA Linux box with triton present. I am handling that separately.

For the record, on a macos-14 runner with the rest of this PR's tests, tests/test_mlx_shape_guard.py, tests/test_mlx_trainer_internals.py, tests/test_mlx_vlm_label_masks.py and tests/test_mlx_training_e2e_metal.py all pass.

Genuine, and fixed in a275ea99.

An arch or backend override can select strict for a qualified VLM while the base compile_mode stays best_effort. _resolve_effective_policy_mode returns the override verbatim and resolve_training_compile derives policy_mode and fallback_allowed from it, so such a run reaches the shape planner with strictness in force. Probed directly with MLXVLMCompilePolicy(mode="best_effort", arch_overrides=(("qwen2_vl", "strict"),)):

base mode            : best_effort
decision.enabled     : True
decision.policy_mode : strict
fallback_allowed     : False
should_raise         : False

should_raise is False and enabled is True, so planning proceeds past both guards and into the failure branches, which read the base compile_policy.mode and therefore see best_effort. Reproduced on the unplannable-family branch:

before  best_effort base + arch override strict : NO RAISE, action=eager, reason=vlm_unplannable_family, allowed=False
        strict base                             : RAISED
after   best_effort base + arch override strict : RAISED
        strict base                             : RAISED

So a user who asked for strict on this architecture silently got an eager run instead of the abort they configured.

Fixing the four planner branches alone would not have been enough on the DDP path: _coordinate_text_shape_guard swallows the local planner error under a best_effort base, so the coordination call now carries the resolved mode too. The helper reads policy_mode off the decision and falls back to the base policy mode when a decision does not carry one, so existing callers and their tests are unaffected.

On newness: the base-mode pattern does predate this PR in the text planner and in _coordinate_text_shape_guard, but neither has a decision plumbed through. _plan_single_process_vlm_shapes is added by this PR and already holds compile_decision, using it for should_raise, so the inconsistency between the two reads is introduced here rather than inherited.

Test: test_vlm_shape_planning_follows_the_resolved_override_mode, which also asserts the reverse override (strict base with a best_effort arch override) still degrades to vlm_unplannable_family rather than aborting. It is red without the source change. Across every tests/test_mlx*.py file the failure set is byte identical to current origin/main, 40 failed on both.

@codex review

Reversing my position on the pad-token guard, which I had rejected twice. The strict-abort one still does not hold up.

Reversed: allow exact VLM plans without a pad token (4e0f2e32)

You were right, and the argument that changes it is one I never checked in either of my earlier replies: this is a merge-base regression, the same shape as the max_seq_length item I conceded in 45478b4b.

At 7ad13f6f, _plan_single_process_text_shapes returned ("not_applicable", "vlm", ...) for is_vlm, FiniteVLMBatchPlan did not exist, and no pad id was required anywhere in VLM planning. Same model, same custom processor, both trees:

merge base 7ad13f6f   compile_enabled=True   scope=full_step       loss 6.082592
this branch           compile_enabled=False  scope=fallback_eager  loss 6.082592
this branch, strict   RuntimeError at setup

So a configuration that compiled before now degrades, and under strict aborts outright. That is a backwards-compatibility break rather than a preference, and it is what my "corner of a corner" framing missed.

My safety argument does not survive either. The branch already accepts the mirror-image abort. With a pad id present and a processor that rebuilds wider than the admitted endpoint:

pad_id=0  rebuild 32 -> 36   ValueError: VLM width plan endpoint 32 is below this batch's
                             prepared width 36                    utils.py:2579
pad_id=0  rebuild 32 -> 30   materialize OK, padded back to 32

_is_compile_exception returns False for that message too and trainer.py:3662 re-raises unconditionally, so "non-idempotent processor means an unrecoverable mid-run abort" was already the contract here. The guard defended only the narrowing half of it while charging the entire pad-id-less population up front.

What does survive from my rejection is that your stated criterion, every scheduled batch already having the same width, is too weak. The forbidden-extent bump widens uniform raw widths: (40,32,32,32,32,32) plans to (40,33,33,33,33,33), and an all-32 schedule plans to (32,33,33,33,33,33) once one batch injects 32 into the union. Keying on the built plan, which is what you actually asked for, handles that correctly, so the fix requires a pad id only when some admitted endpoint exceeds that batch's raw width. Uniform-width and declined-only schedules now stay compilable, and _finalize_vlm_batch_width returns early on not padable so the pad id is provably never consulted there.

Cost, stated honestly: a pad-less plan that does widen now pays one survey pass before degrading, processor_calls going from 0 to len(plan), measured at 6 and 48. The survey runs under _preserved_preprocessing_rng(), so c9d6f65f's determinism invariant is intact, and that pass is already paid unconditionally for vlm_unplannable_family and planning-failure fallbacks. My "degrade to eager for free" line no longer applies to that one case.

red   (guard up front)   5 failed, 26 passed
green (guard on plan)    1 failed, 30 passed, the pre-existing one only
full tests/test_mlx*.py  40 failures, byte identical to origin/main, 0 new

Rejected: coordinate strict aborts before raising

Third time on this one, and the re-raise names no new mechanism and no new input, so here is the full sweep rather than another summary.

Structurally, every should_raise= in resolve_training_compile is either False or strict_requested, which is policy_mode == "strict", and the single branch that sets it from compiled_policy.mode sets policy_mode from that same value. Exhaustively at this head, args.compile by 13 compile_mode inputs by arch overrides by backend-override pairs by patch_mode by 6 architecture scenarios:

15792 combinations, 13536 decisions, 1512 with should_raise=True
effective-mode histogram: {'strict': 1512}     violations: 0

Cross-process determinism: 36 subprocesses over 9 cases varying PYTHONHASHSEED, import order, patch installation and device touch, 0 mismatches, and no environ, getpid, time, random, hash, metal or distributed reads anywhere in the transitive graph.

A two-thread rendezvous driving the real _coordinate_text_shape_guard with a 1.5 s barrier timeout, so a strand raises rather than being assumed: 243 uniform pairs under automatic caps and 162 under a fixed cap, 0 stranded. Force the modes apart by hand and 134 of 162 strand, which is the mechanism you describe, but reaching it needs divergent args or checkpoints across ranks.

The auto-tune angle is closed too: _apply_compile_recommendations mutates exactly one field, args.gradient_checkpointing, and runs at trainer.py:2505, after the decision resolves at 2498 and before compile_policy is rebuilt at 2526. And a275ea99 does cover this path, since 1792 of the swept rows have compile_policy.mode strictness differing from the resolved strictness and trainer.py:2659 passes the resolved mode, so both ranks move together.

One byproduct worth recording rather than acting on: decision_abort at 2667 is unreachable, because should_raise implies strict implies every rank already raised at 1468, and a single process raises at 1462 first.

Splitting this into its three assertions, because they do not all land the same way. One is a regression I introduced and have fixed, one I am declining with measurements, and one I am conceding outright.

Genuine, and my fault: later formatter calls rewrote earlier visits (45e3a364)

You are right, and this is worse than the bug my previous commit was fixing. With a mutate-and-return formatter (example["text"] = ...; return example, the common TRL shape) over a plain list dataset, every slot for a given index held the same object, so last write won. First token id per batch, 6 rows, batch 2, 3 epochs:

merge base 7ad13f6f    1000..1005    2000..2005    3000..3005
origin/main            1000..1005    2000..2005    3000..3005
9ebab255 (before)      1000..1005    1000..1005    1000..1005
4e0f2e32 (after)       3000..3005    3000..3005    3000..3005
with this fix          1000..1005    2000..2005    3000..3005

The old behaviour froze at epoch 1, which was at least correct for epoch 1. Mine made epoch 1 train on epoch-3 data, so no epoch was right. Neither ensure_descriptors nor check_family_drift can see it, since the survey observes the same already-mutated object.

Fixed by snapshotting the container of each formatted row. Only the container, never the payloads, so images and tensors are not duplicated: measured cost is +0.5 MB peak in the worst case, 768 extra dict headers. Exact dict/list matching leaves str returns and non-dict Mappings such as a datasets LazyRow untouched, and nothing depends on row object identity.

Declined: defer the formatter into materialization

I prototyped the literal ask on this head and it does not survive contact with the survey.

ensure_descriptors() builds every scheduled batch, so deferring formatting into materialize runs the formatter 36 times against the eager builder's 18. Any stateful formatter then advances twice per slot, and this PR's own test_vlm_plan_refreshes_visit_dependent_formatting_per_occurrence would read {0: 4, ...} where it asserts {0: 2, ...}. A formatter whose rendered width varies between survey and pass also hard-fails mid-run with ValueError: VLM width plan endpoint 3 is below this batch's prepared width 40 out of _finalize_vlm_batch_width.

Two corrections to reasoning I gave earlier in this thread, both found by executing it rather than arguing it. check_family_drift does not catch a width-changing formatter, because the text axis is symbolic for padable batches and the drift check returns None; the hard failure comes from the width finalizer instead. And _preserved_preprocessing_rng does roll back a formatter drawing from the global generators, so the RNG window was not the blocker I claimed it was. The call doubling is the real blocker.

On ordering, the current head is already at parity for the formatter itself. The formatter-only subsequence is 012345 012345 012345, byte identical to the merge base, and a stateful counter formatter gives [1..18] on both where 9ebab255 gave [1..6] repeated. Head and base diverge only when the processor also draws from the same global RNG, and that divergence predates my commit; it came in with the lazy plan.

Conceded: peak memory, and my earlier wording was wrong

I wrote that the extra rows stay "far below the merge base". That is true only for a formatter that shares its payload. For one that allocates a fresh payload per call it is false, and my commit erased the PR's memory win entirely. 256 rows carrying 1.354 MB each, batch 4, 192 batches, peak RSS:

formatter                       base 7ad13f6f   9ebab255   4e0f2e32
none                              1503.3 MB      515.9 MB   513.6 MB
fresh dict sharing the payload    1498.8 MB      515.9 MB   514.1 MB
allocates a payload per call      1502.5 MB      845.8 MB  1508.3 MB

So for that class my commit took peak from 845.8 MB to 1508.3 MB, a 1.78x increase that lands 5.8 MB above the eager merge base this plan is meant to beat. Retained-after-construction goes 744.8 MB to 1407.2 MB, which is exactly the 512 extra rows times 1.354 MB.

That cost is intrinsic to per-slot correctness while the survey requires user code to be consumed up front, so I am not fixing it here, but it should be recorded rather than argued away. The common payload-sharing case is unaffected at 514.1 MB against 515.9 MB.

Verification

tests/test_mlx_batching_and_decay.py 31 passed against the pre-existing test_compiler_review_guards_are_present, which is red on origin/main too. Red without the fix at assert [20, 21, 22, 20, 21, 22] == [10, 11, 12, 20, 21, 22]. Across every tests/test_mlx*.py the failure set is byte identical to current origin/main, 40 on both, 0 new.

Both genuine, fixed in 31ec13d3. The second one is a hole in the fix I pushed an hour ago.

Genuine: preserve per-occurrence dataset sampling

Literally correct. The compact-remap branch read dataset[idx] once per unique index, while the merge-base eager builder called it once per scheduled slot inside every batch build, so a revisited index replayed the first draw. Dataset returning {"text": str(10*idx + visit)}, 3 rows, batch 1, 6 batches over 2 epochs:

merge base 7ad13f6f   [[0], [10], [20], [1], [11], [21]]   6 __getitem__ calls
9ebab255              [[0], [10], [20], [0], [10], [20]]   3
45e3a364              [[0], [10], [20], [0], [10], [20]]   3
fixed                 [[0], [10], [20], [1], [11], [21]]   6

At batch size 2 over 4 rows, base is [[0,10],[20,30],[1,11],[21,31]] with 8 reads against head's 4.

Blame is not my recent commits: git log -S puts this in 0bf116cb, the first commit of the PR. My 6ef852f3 fixed the sibling formatter branch and walked past this one.

It also reaches a path the item does not mention. With response_mask_fn set and formatting_func unset, formatted_items is None, so the stored row was dataset[idx] and froze the same way: base [[1],[11],[21],[2],[12],[22]] over 9 reads against head's [[1],[11],[21],[1],[11],[21]] over 6. The fix restores base exactly, labels included.

One case is correctly deduped and stays that way: response_mask_fn and formatting_func both set, where _filter_trainable_vlm_indices builds the formatted rows once per index and the eager builder reused those same objects. All three trees already agree there, so the compact remap is kept for that branch alone.

Genuine, and mine: snapshot nested formatted rows recursively

45e3a364 fixed the call count but not the aliasing, because copy.copy does not protect nested containers. A formatter mutating row["messages"][0]["content"] in place, which is the ordinary chat-template shape:

merge base 7ad13f6f   [[1], [2], [3], [4], [5], [6]]
9ebab255              [[1], [2], [3], [1], [2], [3]]   formatter frozen, 3 calls
45e3a364              [[4], [5], [6], [4], [5], [6]]   6 calls, last write wins
fixed                 [[1], [2], [3], [4], [5], [6]]

So every stored slot in both epochs read epoch 2's values, which is the exact bug 45e3a364 was written to kill, surviving one level of nesting. Top-level mutation was already correct at that head.

The snapshot now rebuilds nested dict and list containers, bounded at depth 8 so a self-referential row cannot recurse without end. Payloads are still only referenced, never copied, which is why deepcopy was not an option: with 6000 slots each holding a 64x64x3 image the copy costs about 0.6 KB of container skeleton per slot instead of 12 KB of pixels. Plan build over 2000 rows and 6000 slots goes from 0.60 MB retained and 161 ms to 4.22 MB and 260 ms, where a deepcopy would have duplicated roughly 72 MB. copy.copy on the mapping keeps subclasses such as LazyRow and BatchFeature intact, and only dict and list recurse, so tuples, strings, arrays and PIL images pass through untouched.

Folding these two together also let the branch derive its supervision flag again rather than hardcoding True, which matters now that the per-slot path admits the response_mask_fn-without-formatter case.

Verification

tests/test_mlx_batching_and_decay.py 31 passed before, 33 after, against the pre-existing test_compiler_review_guards_are_present, which is red on origin/main too. Red without the fix on all three assertions, including assert [0, 1, 2, 3] == [0, 1, 2, 3, 0, 1] and assert ['20','21','22','20','21','22'] == ['10','11','12','20','21','22']. Across every tests/test_mlx*.py the failure set is byte identical to current origin/main, 40 on both, 0 new.

One existing test was asserting the bug: test_vlm_plan_formats_only_scheduled_rows_and_compacts_without_a_formatter pinned (len(plain), len(plain.rows)) == (6, 3). It now counts reads and asserts [0, 1, 2, 3, 0, 1] with (6, 6), plus that materialize_all() adds no further reads. The nested test also asserts payload identity per row, so the no-copy contract is pinned rather than assumed.

The previous head verified green on Apple Silicon (31/31 batching, 25/25 shape guard, 41/41 VLM label masks, 118/118 trainer internals, 10/10 e2e Metal, 1/1 DDP, nothing skipped); re-running that for this head.

@codex review

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: 31ec13d3c9

@codex review

Both halves confirmed and fixed in 1f0fc188. This is a merge-base regression rather than a preference, the same shape as the max_seq_length and pad-token items I conceded earlier in this thread.

I ran the real training loop on this head and on the merge base with a processor that augments from the process-global RNGs and stamps its draw into the ids it returns, so a batch's contents show which position of the preprocessing stream built it. Sweep: gradient_accumulation_steps in {1,2,3} by formatting_func set/unset by eval_dataset present/absent by response_mask_fn present/absent, checkpoint at optimizer step 2 of 6, plus a second pass varying whether the descriptor survey ran.

Resume

On the merge base a resumed run reproduced the uninterrupted run byte for byte in 24 of 24 configurations, because create_vlm_batches had built every scheduled batch up front, so batch j always carried the j-th processor call. On this head the resumed run differed from the uninterrupted run in 24 of 24. The mechanism is exactly as you name it: batch_idx starts at _resume_step * grad_accum, and the only fast-forward is gated on batch_iter is not None, which is streaming, so a finite FiniteVLMBatchPlan never replays the skipped materialize() calls.

grad_accum=1, batch 2, 8 rows, seed 3407, resuming at step 2, stamp is (random, numpy, torch):

uninterrupted, both trees, and merge-base resumed
  step 3   rows (6,7)   aug [979, 1770, 2552]
  step 4   rows (4,3)   aug [364, 1969, 3457]
this head, resumed at step 2
  step 3   rows (6,7)   aug [958, 1505, 3135]    the draw step 1 used
  step 4   rows (4,3)   aug [383, 1839, 2548]    the draw step 2 used

Row selection is correct in every case; only the augmentation draws are rewound. At grad_accum=3 the whole resumed stream is the baseline stream shifted back by exactly the 6 skipped micro-batches.

Eval ordering

Also confirmed. VLM eval batches are built eagerly before the loop, so with a lazy train plan those processor calls now land ahead of the first training batch. With an eval split of 4 rows, 2 eval batches, this head's first training batch is shifted forward by exactly 2 draws relative to the merge base, in 12 of the 12 configurations that set eval_dataset. The merge base was immune only because the training batches were already built by the time the eval ones were.

The survey is clean

Worth recording, since it is the obvious suspect and it is not the culprit: ensure_descriptors() under _preserved_preprocessing_rng() left random, numpy.random and the global torch generator exactly where it found them in 24 of 24 configurations, with batch contents byte identical with and without it. Private generators are not restored, which that docstring already declares out of scope. The gap is the two orderings above.

The fix

FiniteVLMBatchPlan gains advance_preprocessing(visits): it rebuilds and discards the first visits scheduled visits, dropping the MRU and holding one batch at a time, and deliberately does not preserve RNG state, since advancing it is the point. train() calls it right after the streaming fast-forward, for finite VLM plans only, with the same distributed failure coordination that branch already uses. The count is clamped to the schedule length, because the eager builder ran the processor at most once per scheduled batch however far the killed run got. Cost is strictly below what the eager builder paid on every run, resumed or not, and peak memory is unchanged at one batch.

VLM eval batch construction is wrapped in _preserved_preprocessing_rng(), so building an eval split no longer consumes draws the first training batch is owed.

Trade-off stated up front: with that wrap, eval batches now draw from stream positions 1..E rather than N+1..N+E. That is unavoidable once training is lazy, it moves only eval-set augmentation, which is a fixed sample either way, and it is the only ordering that keeps the training stream identical to the eager baseline.

Verification

Two regression tests in tests/test_mlx_batching_and_decay.py drive the real training loop rather than the plan in isolation: one asserts a resumed run reproduces the uninterrupted run's tail, one asserts adding an eval split does not move the training batches. Both fail without the change and pass with it. After the fix, all three predicates hold in 24 of 24: resumed equals uninterrupted, uninterrupted equals merge base, and resumed equals merge-base resumed.

tests/test_mlx_batching_and_decay.py 35 passed against the pre-existing test_compiler_review_guards_are_present, which is red on origin/main too. Across every tests/test_mlx*.py the failure set is byte identical to current origin/main, 61 on both, 0 new and 0 fixed.

One harness note for anyone reproducing this: the first batch build lazily imports rich, which calls random.getrandbits(24) at import time. Without warming that up first, it manufactures a survey-shaped difference that is not real.

@codex review