Implement dpo tuning

[eghuzefa]

Add Direct Preference Optimization (DPO) Support

Summary

This PR implements Direct Preference Optimization (DPO) for MLX-LM, enabling users to fine-tune language models using human preference data without requiring a separate reward model.

What is DPO?

DPO is a simpler alternative to RLHF that directly optimizes on preference pairs (chosen vs rejected responses), avoiding the complexity of training reward models and using PPO. It's mathematically equivalent to RLHF but more stable and efficient.

Key Features Added

Core Implementation

• DPO Loss Function: Bradley-Terry preference model with KL regularization
• Preference Dataset: Handles both simple prompt format and chat conversation format
• Dual Forward Passes: Policy and reference model evaluation during training
• Beta Parameter: Configurable temperature for preference strength

Integration

• CLI Interface: python -m mlx_lm dpo with full argument parsing
• Fine-tuning Support: Compatible with LoRA, DoRA, and full fine-tuning
• Memory Efficient: Works with quantized models (QLoRA)
• MLX Native: Follows existing MLX-LM patterns and conventions

Data Format Support

{"prompt": "What is AI?", "chosen": "Detailed explanation...", "rejected": "Short answer."}
{"messages": [{"role": "user", "content": "Hello"}], "chosen": "Detailed response", "rejected": "Hi."}

Documentation & Testing

• Comprehensive Documentation: mlx_lm/DPO.md with usage examples and best practices
• Unit Tests: Full test coverage for loss functions, datasets, and integration
• End-to-end Validation: Successfully trained models with real preference data

Usage Example

# Basic DPO training
mlx_lm.dpo \
    --model mlx-community/Meta-Llama-3-8B-Instruct-4bit \
    --train \
    --data preference_data/ \
    --beta 0.1 \
    --fine-tune-type lora

Files Added/Modified

• mlx_lm/dpo.py - Main DPO module
• mlx_lm/tuner/losses.py - DPO loss implementation
• mlx_lm/tuner/datasets.py - PreferenceDataset class
• mlx_lm/tuner/trainer.py - DPO training functions
• mlx_lm/__main__.py - CLI registration
• mlx_lm/DPO.md - Documentation
• tests/test_dpo.py - DPO-specific tests
• tests/test_losses.py - Enhanced with DPO loss tests

Testing

• ✅ All existing tests pass
• ✅ New DPO unit tests pass (9/9)
• ✅ End-to-end training validation completed
• ✅ Memory usage verified with quantized models

Benefits for MLX-LM Users

1. Simplified Preference Training: No need for complex RLHF pipelines
2. Better Alignment: Train models to follow human preferences effectively
3. Resource Efficient: Works with existing LoRA/QLoRA infrastructure
4. Production Ready: Stable, well-tested implementation

This implementation enables MLX-LM users to easily train more helpful, harmless, and honest models using preference data.

[eghuzefa]

Hi maintainers! @awni @Goekdeniz-Guelmez This PR adds DPO support for preference-based fine-tuning. It's been open for ~2 months and I'd appreciate any feedback. Happy to address any concerns or make changes. Thanks!

[Goekdeniz-Guelmez]

Hey Mohammad, merging will take time, since this is a big addition, you'd have to wait so that it passes all the reviews and has the needed code quality.

[Totorong-Bot]

@angeloskath @awni

Dear MLX teams. Some users(including me) have been awaited adoption of DPO into MLX framework. We'll highly appriciate if you consider merging of ths PR.