Fine-Tune LLMs on Apple Silicon with mlx-lm LoRA

Apple’s mlx-lm is a Python package for generating text and fine-tuning large language models on Apple silicon with the MLX framework. It supports both low-rank and full model fine-tuning, including for quantized models, and integrates with the Hugging Face Hub. The package installs with pip install mlx-lm, and the training extras come from pip install "mlx-lm[train]".

LoRA and QLoRA Training

The main command for fine-tuning is mlx_lm.lora. According to the documentation at https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md, the command can fine-tune a model using low-rank adaptation (LoRA) or quantized LoRA (QLoRA). The choice is automatic: if --model points to a quantized model, the training uses QLoRA, otherwise it uses regular LoRA.

A basic run uses --train, --data, and --iters. The --fine-tune-type flag selects between lora (the default), dora, and full. Adapters are written to a path set by --adapter-path, which defaults to adapters/, and --resume-adapter-file continues from an existing adapter. The --mask-prompt flag computes loss only on completions for chat and completion datasets, and --test computes test-set perplexity. Options can also be supplied through a YAML config with -c or --config.

The documented model families include Mistral, Llama, Phi2, Mixtral, Qwen2, Gemma, OLMo, MiniCPM, and InternLM2.

Data Formats

Training data uses .jsonl files in four supported formats: chat, tools, completions, and text. The loader expects a train.jsonl file for training, an optional valid.jsonl for validation loss, and test.jsonl for evaluation, with each example on a single line. Hugging Face datasets are supported after installing the datasets package, configured through hf_dataset keys such as prompt_feature, completion_feature, text_feature, and chat_feature.

Working Within Memory Limits

The documentation lists several ways to reduce memory use during training. QLoRA quantizes a model with the -q flag. The --batch-size flag defaults to 4 and can be lowered to 2 or 1, optionally combined with --grad-accumulation-steps. The --num-layers flag defaults to 16 and can be reduced to 8 or 4. Examples can also be broken into shorter sequences, and --grad-checkpoint trades memory for additional computation.

As a concrete example, the docs describe fine-tuning Mistral-7B on a 32 GB machine with batch size 1, four layers, and the wikisql data, reporting roughly 250 tokens per second on an M1 Max.

Using the Result

After training, mlx_lm.generate runs the model with --model, --adapter-path, and --prompt. The mlx_lm.fuse command merges adapters back into the base model, saving to fused_model/. Fuse supports uploading through --upload-repo and --hf-path, and can export to GGUF with --export-gguf, though GGUF export is limited to Mistral, Mixtral, and Llama style models in fp16 precision.

The package also documents distributed inference and fine-tuning through mx.distributed, and logging to Weights and Biases or SwanLab with --report-to. For teams already working on Mac hardware, mlx_lm.lora provides a documented path to adapter-based fine-tuning without leaving Apple silicon.