Fix MLX multi-batch validation memory growth

[yhn112]

Summary

• materialize each MLX validation batch loss before accumulating it, instead of building one lazy graph across the entire validation split
• keep lightweight validation progress logging so long validation passes do not look hung

Why this happens

This only shows up when validation spans multiple batches. In eval_val(), total_loss was accumulated as an mx.array, so MLX kept extending a lazy graph until the final mx.eval(...). Memory usage then grew steadily during validation, and the script looked like it hung.

This is especially easy to hit in the MLX smoke flow because validation always runs over the full fineweb_val_* split, while the README example uses a small VAL_BATCH_SIZE. In practice the first visible symptom is often after the final training step, which makes it look like a post-training hang.

Single-batch validation does not exhibit the problem.

Logging

The progress logging is intentional here: once validation is split across many batches, a long final validation pass can otherwise look indistinguishable from a freeze.

Validation

• python -m py_compile train_gpt_mlx.py
• reproduced locally with a short MLX run where memory growth started only once final validation began, then verified the fix stopped the growth

[0hq]

Thanks! Missed this.