Iterate SFT batches. Examples are length-sorted (mlx-lm iterate_batches),
grouped into contiguous runs of batchSize, and the batch order is
shuffled with seed each epoch. Examples longer than maxSeqLen are
right-truncated (preserving the prompt-boundary clamp). loop=false
yields one epoch then stops.
B=1: yields one example per step with NO padding and no pad-to rounding —
bit-identical to the original single-example trainer.
B>1: faithful to mlx-lm. Sort ascending by length, take contiguous
windows idx[i : i+B] for i in 0, B, 2B, … (the trailing remainder
that does not fill a batch is dropped, matching mlx-lm's range bound),
permute batch order, and pad every row to
min(1 + PAD_TO·ceil(max_len/PAD_TO), maxSeqLen) with padId. Each row
carries its true (truncated) length so the loss + attention masks ignore
padding. padId defaults to 0 (mlx-lm SFT uses np.zeros); pass the eos id
to pad with a real sentinel — the padded positions are masked out either
way, so the choice does not change the loss.
Iterate SFT batches. Examples are length-sorted (mlx-lm iterate_batches), grouped into contiguous runs of
batchSize, and the batch order is shuffled withseedeach epoch. Examples longer thanmaxSeqLenare right-truncated (preserving the prompt-boundary clamp).loop=falseyields one epoch then stops.B=1: yields one example per step with NO padding and no pad-to rounding — bit-identical to the original single-example trainer.
B>1: faithful to mlx-lm. Sort ascending by length, take contiguous windows
idx[i : i+B]fori in 0, B, 2B, …(the trailing remainder that does not fill a batch is dropped, matching mlx-lm's range bound), permute batch order, and pad every row tomin(1 + PAD_TO·ceil(max_len/PAD_TO), maxSeqLen)withpadId. Each row carries its true (truncated) length so the loss + attention masks ignore padding.padIddefaults to 0 (mlx-lm SFT uses np.zeros); pass the eos id to pad with a real sentinel — the padded positions are masked out either way, so the choice does not change the loss.