🚀 MLX Distributed Training Framework

Train and deploy AI models across multiple Apple Silicon devices with automatic hardware optimization and a seamless developer experience. A user-friendly CLI framework for distributed training, powered by MLX.

╔══════════════════════════════════════════════════════════╗
║  __  __ _     __  __  _____           _                  ║
║ |  \/  | |    \ \/ / |_   _| __ __ _ (_) _ __            ║
║ | |\/| | |     \  /    | | | '__/ _` || || '_ \          ║
║ | |  | | |___  /  \    | | | | | (_| || || | | |         ║
║ |_|  |_|_____|/_/\_\   |_| |_|  \__,_||_||_| |_|         ║
║                                                          ║
║          Distributed Training on Apple Silicon           ║
╚══════════════════════════════════════════════════════════╝

┌─────────────────┐      ┌─────────────────┐
│   MacBook Pro   │  ←→  │   MacBook Air   │  Training
│    Node 1       │      │    Node 2       │  Cluster
└────────┬────────┘      └────────┬────────┘
         │                        │
         └──────────┬────────────-┘
                    ▼
         ┌──────────────────────┐
         │    MLX Training      │  Distributed
         │   ═══════════ 100%   │  Progress
         └──────────────────────┘
                    ▼
         ┌──────────────────────┐
         │  🚀 Trained Model    │  Local
         │  localhost:8000      │  Deployment
         └──────────────────────┘

⚡️ Quick Start & Interactive Script

# Clone the repository
git clone https://github.com/jbarnes850/mlx-train
cd mlx-train

# Install dependencies
pip install mlx  # Core MLX framework
pip install -r requirements.txt  # Additional dependencies

# Make the script executable
chmod +x scripts/mlx-train.sh

# Run the interactive training script
./scripts/mlx-train.sh

That's it! The script will:

• 🛠️ Set up your environment automatically
• 🔍 Detect your Apple Silicon devices
• 🤝 Help connect multiple devices if available
• 🚀 Guide you through model training
• 🌐 Deploy your model locally

See quickstart for more details.

Watch your model train with real-time visualizations:

Device Utilization:
Device 1: [====================] 80.5% (12.8GB)
Device 2: [===================] 75.2% (12.1GB)

Training Progress:
[====================] 50%

Metrics:
• Throughput: 1250.32 tokens/sec
• Loss: 0.0234
• Memory Usage: 25.6GB

✨ Features

• 🚄 Distributed Training: Seamlessly scale across multiple Apple Silicon devices
• 🔧 Hardware Optimization: Automatic detection and configuration for optimal performance
• 🎯 Zero-Config Setup: Automatic environment setup and dependency management
• 📊 Training Visualization: Real-time metrics and progress tracking
• 🧠 Model Development: Build custom architectures or use pre-built components
• 🔄 Export & Serve: Deploy models locally or export for other platforms

Scale When Ready:

• 🔄 Start with any Apple Silicon Device (M1/M2/M3/M4)
• 🔄 Add more devices anytime
• 🔄 Automatic distributed training
• 🔄 No code changes needed

🛠 Installation & Requirements

Requirements

• Python 3.10 or higher
• Apple Silicon Mac (M1/M2/M3)
• 8GB RAM minimum (16GB+ recommended)
• macOS 12 or higher

Quick Install

pip install mlx-train

Development Setup

# Clone repository
git clone https://github.com/jbarnes850/mlx-train
cd mlx-train

# Create virtual environment
python3 -m venv venv
source venv/bin/activate

# Install dependencies
pip install -e ".[dev]"

The mlx-train.sh script will automatically check and install all requirements.

🔧 Advanced Configuration

Hardware Optimization

• Automatic device discovery
• Memory-aware batch size selection
• Gradient accumulation optimization

Training Options

• Mixed precision training
• Gradient checkpointing
• Custom learning rate schedules

🤝 Contributing

We welcome contributions! Please see our Contributing Guide for details.

1. Fork the repository
2. Create your feature branch (git checkout -b feature/amazing-feature)
3. Commit your changes (git commit -m 'Add amazing feature')
4. Push to the branch (git push origin feature/amazing-feature)
5. Open a Pull Request

📚 Documentation

Documentation Installation Examples Core Concepts Roadmap

🔬 Examples

Check out our examples directory for:

• Custom model training
• Multi-device distributed training
• Dataset preprocessing
• Model export and serving

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

📫 Support

• 🐛 Issue Tracker

---

Made with ❤️ for the MLX community