add peak_memory info to training callback (#572)

2025-09-01 04:14:38 +08:00 · 2024-03-14 11:17:10 +08:00
parent 376bb9cc44
commit d4e1de1d5b
1 changed files with 1 additions and 0 deletions
--- a/llms/mlx_lm/tuner/trainer.py
+++ b/llms/mlx_lm/tuner/trainer.py
@@ -250,6 +250,7 @@ def train(
                    "iterations_per_second": it_sec,
                    "tokens_per_second": tokens_sec,
                    "trained_tokens": trained_tokens,
+                    "peak_memory": peak_mem,
                }
                training_callback.on_train_loss_report(train_info)