Visualization for MLX Training logs

[Gautam-Galada]

Description

This pull request adds visualization feature with 6 plots currently , keeping the current MLX training setup in mind. The image contents are as follows :

models/visualization/{timestamp}/plots/
├── 1_loss_curves.png
├── 2_learning_rate.png
├── 3_training_speed.png
├── 4_throughput.png
├── 5_memory_usage.png
└── 6_token_progress.png

How to Test

# Default: Creates 6 separate files
python scripts/visualize_mlx_training.py

# Legacy: Single combined plot
python scripts/visualize_mlx_training.py --combined

# High resolution
python scripts/visualize_mlx_training.py --dpi 300

# Custom directory
python scripts/visualize_mlx_training.py --output my_analysis

# Custom timestamp
python scripts/visualize_mlx_training.py --timestamp 2026-02-01_01-28-37

Checklist

Added matplotlib to pyproject.toml
Unit-testing is included by pytest tests/test_visualize_mlx_training.py -v

Additional Information

The visualization code is folder friendly, it creates the plots for the latest timestamp, or specific if provided as an argument.

[Copilot]

Pull request overview

This pull request adds a comprehensive visualization feature for MLX training logs. The implementation includes a script that parses training logs, extracts various metrics (loss, learning rate, throughput, memory usage, etc.), and generates informative plots to help analyze training runs.

Changes:

• Added scripts/visualize_mlx_training.py with log parsing, run management, and visualization capabilities
• Added unit tests for log parsing and run management functionality
• Added matplotlib as a dependency in pyproject.toml

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 4 comments.

File	Description
scripts/visualize_mlx_training.py	New script implementing MLX training log parser, run manager, and visualizer with support for both individual and combined plot generation
tests/test_visualize_mlx_training.py	Unit tests for log parsing and run management functionality
pyproject.toml	Added matplotlib dependency for visualization support

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.