Reusable segmented-backward driver for one (model, lora, segmentation). Build
ONCE per training run and call .step(batch) each iteration — the per-segment
value_and_grads are created up front and reused (rebuilding them every step
leaks at the mlx level). Returns the loss scalar and the LoRA gradients in
flatParams(lora) order ([...A, ...B] over targets), drop-in for
ValueAndGrad.apply(flatParams(lora)). Caller owns the returned arrays.
The active adapter must be attached (attachForTraining) so the layer forward
picks up the swapped LoRA leaves via loraState. B=1 SFT only. The graph must
declare segmented training.
Reusable segmented-backward driver for one (model, lora, segmentation). Build ONCE per training run and call
.step(batch)each iteration — the per-segment value_and_grads are created up front and reused (rebuilding them every step leaks at the mlx level). Returns the loss scalar and the LoRA gradients inflatParams(lora)order ([...A, ...B] over targets), drop-in forValueAndGrad.apply(flatParams(lora)). Caller owns the returned arrays.The active adapter must be attached (attachForTraining) so the layer forward picks up the swapped LoRA leaves via loraState. B=1 SFT only. The graph must declare
segmentedtraining.