Train MLX

Status: WIP / Experimental

The goal of this project is to provide MLX training on Apple Silicon (and CUDA devices) with semi-drop-in support for Axolotl configurations (best-effort).

Setup

Sync dependencies using uv:

uv sync

Usage

Run the training scripts using uv run. For example:

uv run scripts/mlx_trainer.py --config config.example.yml

Development

We use black for formatting and ruff for linting.

# Format code
uv run black .

# Lint code
uv run ruff check .

Limitations on CUDA

Due to MLX limitations, some features are not supported on CUDA devices:

• QLoRA (4bit and 8bit)
• Possibly other stuff!