AFM Trainer

NOTE TO USERS: The WanDB intergration remains untested until I get a chance. Apologies for the inconvenience.

Requires Python 3.11+

The LoRA adapter can be used and tested on MacOS Beta 26+ with the afm command: afm -a output/my_adapter.fmadapter -s "Write a story about Einstein'

Get afm here: https://github.com/scouzi1966/maclocal-api

On MacOS: $$\color{blue}{python \space run.py}$$ (all packages will be downloaded automatically with uv and will use Pytoch with Metal)

On Linux (Ubuntu tested): $$\color{blue}{python\space run-linux.py}$$ (all packages will be downloaded automatically with uv and will use Pytoch with CUDA)

Note: On Linux, the CUDA python packages will be downlaoded as part of the virtual environment and can take a lot of time depending on your bandwidth. I suggest not cleaning up your environemnt after every quit to save time. Clean it up when you done or on your last run when you quit.

A comprehensive GUI wrapper application for Apple's Foundation Models Adapter Training Toolkit, providing an intuitive interface for training LoRA adapters for Apple's on-device foundation models.

⚠️ Disk Space Requirements

Important: The UV package manager will download and cache dependencies that require approximately ~1 GB of available disk space. This includes:

• PyTorch and related ML libraries (~800 MB)
• CoreML tools and dependencies (~150 MB)
• Additional Python packages (~50 MB)

Ensure you have sufficient free disk space before installation. You can clean up the UV cache later using the built-in cleanup options when quitting the application.

🧹 Disk Space Management

• Automatic Cleanup: When quitting, the app offers to clean up the UV environment automatically
• Manual Cleanup: If you decline, detailed instructions are provided for manual cleanup
• Cache Location: UV stores packages in a system cache directory (typically ~/.cache/uv/)
• Safe to Clean: The cache will be recreated automatically when you next run AFM Trainer

🌟 Features

• 🎯 Modern GUI: Beautiful, professional interface with dark/light theme switching
• 🎨 Enhanced Theming: Modern Sun Valley theme for a polished, forest-inspired appearance
• 📊 Real-time Monitoring: Live training progress with loss visualization and log streaming
• 🔄 Complete Workflow: From dataset validation to .fmadapter export
• ⚙️ Flexible Configuration: All Apple toolkit parameters exposed with sensible defaults
• 📈 WandB Integration: Optional Weights & Biases logging for advanced metrics tracking
• 🛡️ Comprehensive Error Handling: User-friendly error messages and crash reporting
• 🔧 UV Environment Management: Transparent dependency handling and installation
• 📁 Smart File Management: Automatic .gitignore updates and dataset validation
• 🚀 Draft Model Training: Optional speculative decoding support
• 💾 Configuration Profiles: Save and load training configurations
• 🌓 Theme Customization: Real-time switching between dark and light themes

📋 Requirements

• Operating System: macOS or Linux (Apple Silicon and Intel Macs supported)
• Python: 3.11 or higher
• Package Manager: UV (will be installed automatically if missing)
• Apple Toolkit: Foundation Models Adapter Training Toolkit v26.0.0+ (requires entitlements)

🚀 Quick Start

1. Get the Application

# Option 1: Clone the repository
git clone <repository-url>
cd AFMTrainer

# Option 2: Download and extract the release package

2. Get Apple's Toolkit

⚠️ Important: You must have Apple Developer Program entitlements to access the toolkit.

1. Visit the official Apple documentation: Apple Foundation Models Adapter Training
2. Download the toolkit from Apple (requires entitlements)
3. Place it in one of these locations:
   • .adapter_training_toolkit_v26_0_0/ (recommended for auto-detection)
   • adapter_training_toolkit_v26_0_0/

Note: The toolkit is not included with AFM Trainer and must be obtained directly from Apple.

3. Launch the Application

# Simple launcher (recommended)
# For Linux users, it is recommended to use the dedicated Linux launcher:
# python run-linux.py
# For macOS, use the universal launcher:
python run.py

# Or use UV directly
uv run afm-trainer

# Or run the GUI module directly
uv run python -m afm_trainer.afm_trainer_gui

The launcher will automatically:

• Check Python version compatibility
• Install UV if missing
• Set up the environment
• Detect the Apple toolkit
• Launch the GUI with modern theming

🎨 Visual Experience

AFM Trainer features a modern, professional interface with:

🌓 Theme Options

• Dark Theme (default): Professional dark interface inspired by modern development tools
• Light Theme: Clean, bright interface for different preferences
• Real-time Switching: Toggle between themes instantly from the header

✨ Visual Enhancements

• Modern Styling: Sun Valley theme providing a forest-inspired, Excel-like appearance
• Professional Typography: Enhanced fonts and spacing for better readability
• Icon Integration: Meaningful icons throughout the interface for improved navigation
• Visual Hierarchy: Clear organization with styled sections and grouped controls
• Enhanced Controls: Modern buttons, entries, and interactive elements

🎯 User Experience

• Intuitive Layout: Logical flow from setup to training to export
• Visual Feedback: Clear status indicators and progress visualization
• Accessibility: High contrast and readable typography in both themes
• Consistency: Unified styling across all interface elements

Try the Theme Demo: Run uv run python theme_demo.py to see the visual improvements!

⚡ Performance Mode

For ultra-fast tab switching, enable Performance Mode:

# Enable high-performance mode
AFM_TRAINER_PERFORMANCE_MODE=true python run.py

Performance Mode Features:

• ⚡ Instant Tab Switching: 2-3x faster tab switching performance
• 🎯 Optimized Theme: Fast native theme instead of sv-ttk
• 🚀 Reduced Overhead: Minimal theme updates and visual effects
• 💡 Smart Fallbacks: Maintains full functionality with speed priority

When to Use Performance Mode:

• Slow tab switching on your system
• Working with large datasets requiring frequent tab changes
• Older hardware or systems with limited graphics performance
• Preference for speed over visual styling

Theme Comparison:

• Normal Mode: Beautiful dark theme, ~0.06s tab switching
• Performance Mode: Clean native theme, ~0.03s tab switching (50% faster)

📖 User Guide

Setup Tab

1. Toolkit Configuration

   • Browse and select your Apple toolkit directory
   • The app will automatically validate the toolkit and update .gitignore

2. Dataset Configuration

   • Select your training JSONL file (required)
   • Optionally select evaluation JSONL file
   • The app will validate dataset format and show preview

3. Output Configuration

   • Choose output directory for checkpoints and exports
   • Directory will be created automatically if it doesn't exist

Training Tab

Configure all training parameters with real-time validation:

Basic Parameters:

• Epochs: Number of training passes (default: 2)
• Learning Rate: Step size for parameter updates (default: 1e-4)
• Batch Size: Samples per training batch (default: 4)

Advanced Parameters:

• Warmup Epochs: Learning rate warmup period (default: 1)
• Gradient Accumulation: Steps to accumulate gradients (default: 1)
• Weight Decay: Regularization coefficient (default: 1e-2)
• Precision: Training precision (bf16-mixed, f16-mixed, bf16, f32)
• Activation Checkpointing: Memory optimization (trades compute for memory)
• Model Compilation: Performance optimization (CUDA only)

Export Tab

Configure adapter metadata and export options:

• Adapter Name: Name for your exported adapter
• Author: Your name or organization
• Description: Detailed adapter description
• License: License information
• Draft Model: Enable speculative decoding training

WandB Integration:

• Enable/disable Weights & Biases logging
• Automatic login status detection
• Real-time training metrics tracking

Monitor Tab

Real-time training monitoring:

• Progress Bar: Visual training progress
• Live Logs: Streaming training output with timestamps
• Loss Tracking: Real-time loss values and trends
• Time Estimates: Estimated completion times

📊 Dataset Format

Training data must be in JSONL (JSON Lines) format. Each line represents one training sample:

Basic Format

[{"role": "user", "content": "Tell me about cats"}, {"role": "assistant", "content": "Cats are fascinating animals..."}]

With System Instructions

[{"role": "system", "content": "You are a helpful pet expert"}, {"role": "user", "content": "Tell me about cats"}, {"role": "assistant", "content": "Cats are fascinating animals..."}]

Multi-turn Conversations

[{"role": "system", "content": "You are a helpful assistant"}, {"role": "user", "content": "What's 2+2?"}, {"role": "assistant", "content": "4"}, {"role": "user", "content": "What about 3+3?"}, {"role": "assistant", "content": "6"}]

Alternative Wrapped Format

{"messages": [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}]}

The app includes comprehensive dataset validation and will report:

• Format errors