LoRA: Increased volatility of train loss

When using the latest LoRA training, the volatility of Loss became larger. When I further analyzed the cause, I suspected that it might have been introduced by commit #528 due to an internal reordering of the dataset.

https://github.com/ml-explore/mlx-examples/blob/e2205beb668abb2334a75f405de094436cd323fe/llms/mlx_lm/tuner/trainer.py#L78-L86

I revert this commit and retrain, and the loss curve returns to stability. The following is the situation during my training. The fluctuation in this part will reduce the accuracy of the final model by about 10%.

I'm not quite sure what the reason for this adjustment is, or how I should ignore this part of the logic for sorting the dataset in this case. If this part of the logic is really needed, the way I think of it is to enlarge it in the Dataset definition and allow adding your own data sets.

Start with llms/mlx_lm/tuner/trainer.py lines 78-86 and compare the current dataset ordering with commit #528 and its reverted behavior. Reproduce the reported LoRA training case, then determine whether the ordering change causes the loss volatility and confirm that the resulting training curve and final accuracy are stable.