custral

computer-use × Mistral: fine-tuning an open-weight Mistral vision model into a Computer-Use Agent, entirely locally on Apple Silicon (MLX). Milestone 1 — the focus of this repo — is GUI grounding: given a screenshot and an instruction, output where to click.

Base model: Ministral 3 14B Instruct (4-bit) · Toolchain: mlx-vlm · Hardware: one MacBook (M4 Max).

Trained adapters (SFT + RFT) are on Hugging Face: ukanwat/custral-14b-lora — no retraining needed. Setup: pip install mlx-vlm datasets pillow (Apple Silicon).

The task

input :  screenshot  +  "Click on: add to cart button"
output:  click(x=786, y=194)        # x,y normalized to 0-1000 (resolution-independent)

Action format is a compact function-style DSL (click(x,y), type(text=…), scroll(…), press(…)) — the field standard for GUI agents (UI-TARS / OS-Atlas / Aguvis), not JSON tool-calls.

Results (ScreenSpot, point-in-bbox, n=150)

Stage	ScreenSpot	Notes
Ministral 3 14B baseline (un-tuned)	22.0%	format-parses 100% out of the box
+ web-grounding SFT (Wave-UI, 1400 iters)	22.0%	learned the DSL, not yet the pointing
+ mix SFT (13k web+desktop+action, cum. 2200)	34.0%	text 41.4 / icon 23.8
+ mix SFT continued (cum. 3000 / 4000)	30.7% / 30.0%	plateau — extra iters redistribute, don't add
+ RFT (rejection-sampling fine-tune, 1400 iters)	28.0%	null result — within noise of the SFT plateau; sharpening didn't transfer

The headline finding: with a frozen vision encoder, 13k examples, and LoRA on the language model only, this recipe converges to low-30s ScreenSpot — a real +8–12 pt lift over baseline that then plateaus. The levers that would change the regime (grounding-pretrained base like Qwen-VL, unfrozen vision tower, 100× more data) all live outside the "quantized Mistral on a laptop" constraint this project set out to explore.

Repo layout

prep/     dataset builders (one-time online; everything after is offline)
  prep_waveui.py       Wave-UI-25K -> balanced grounding set in the DSL
  prep_osatlas.py      OS-Atlas desktop grounding subset -> data_mix
  prep_aguvis_web.py   Aguvis stage-2 web ACTION configs (click/type/scroll/press) -> data_mix
  prep_rft.py          rejection sampling: best-of-K in-box clicks from a trained ckpt -> data_rft
  audit_data.py        visual + statistical sanity check of prepared data
train/
  resume_train.py      warm-start QLoRA trainer (the workaround that makes 4-bit resumable)
  train_resilient.sh   crash-proof auto-resuming training wrapper
eval/
  eval_grounding.py    point-in-bbox scoring on ScreenSpot or any prepped imagefolder
tools/                 probes, inspectors, and ops scripts used during the runs
docs/
  AUTODRIVE.md         artifact: the autonomous train->eval->RFT runbook an agent ran for ~3 days

All scripts assume they are run from the repo root and default to a Python env with mlx-vlm, datasets, and Pillow installed (PY=<venv python> overrides in the shell scripts).

The recipe (empirically derived — see findings)

python train/resume_train.py \
  --model mlx-community/Ministral-3-14B-Instruct-2512-4bit \
  --dataset data_mix --iters 4000 \
  --resume-from <checkpoint or fresh> --output-path adapters_mix/adapters.safetensors \
  --grad-checkpoint --steps-per-save 200
# defaults: LoRA r=32 a=64, lr 2e-5, batch 1, 1024x1024, train-on-completions

Findings (all measured locally)

Training mechanics

• 4-bit base ⇒ LLM-only QLoRA. --train-vision fails on a quantized base (QuantizedMatmul::vjp: no gradient wrt quantized weights); the vision encoder stays frozen. For a grounding task this is the binding constraint — the eyes never learn.
• --grad-checkpoint is mandatory — without it batch>1 exceeds unified memory.
• --batch-size 1 is required — mlx-vlm's collation does a naive mx.stack (no padding), so variable-length VLM examples can't batch. Costs nothing: the GPU is saturated at batch 1 (~70–110 tok/s; batch-4 is no faster).
• Built-in resume OOMs on 4-bit — apply_lora_layers dequantizes the base (8 GB → ~28 GB). train/resume_train.py works around it: load the base the fresh (quantized) way, add LoRA layers, then inject the saved adapter weights. A crash costs ≤ steps-per-save steps.
• Eval must pause training — two 14B copies exceed unified memory. One model at a time.

What the numbers mean

• Train loss is the wrong gauge for grounding. Completions are ~76% boilerplate (click(x=…, y=…)) that hits zero loss immediately, and ~24% coordinate digits that are irreducibly unpredictable at the token level. Loss pins at ~2.4 forever while point-in-bbox accuracy climbs from 22% to 34%. Eval on the metric, every ~1000 iters; ignore the loss.
• The mix-SFT plateaus early. Accuracy peaked around cum. iter 2200 and stayed flat (±noise) for 1800 more iterations. More of the same data at lr 2e-5 buys nothing after convergence.
• Best-of-4 sampling beats greedy (that gap is what RFT distills): at temp 0.8 the SFT model solves ~17–19% of held-out mix grounding best-of-4. Rejection sampling keeps the in-box click closest to box center; retaining action rows prevents forgetting. In practice the distilled sharpening did not transfer to ScreenSpot (28.0% vs 30.0% SFT — a null result): with only ~300 self-solved examples and a frozen vision tower, there was too little new signal to move the benchmark. The plateau is a recipe/data ceiling, not an inference-sharpness problem.
• Generation is prefill-bound. A 1024×1024 screenshot is ~2.5k image tokens per sample; K=4 sampling costs ~2–2.5 min/example on M4 Max. Budget data-gen accordingly (or rent an A100: the whole pipeline here ≈ $10–15 of cloud time).

Data

• Wave-UI's instruction field is just the element type ("button") — useless; name ("add to cart button") is the correct grounding target.
• Mix (13k): Wave-UI web grounding + OS-Atlas desktop grounding + Aguvis web actions (~18k click, 1.6k type, 1.3k scroll, 0.2k press across targets).
• Square-resize to 1024×1024 distorts aspect ratio per screenshot — a suspected accuracy tax; native-aspect processing is the fix (future work).

Resilience (hard-won)

• Everything offline (HF_HUB_OFFLINE=1), checkpoint every 200 steps, distinct-named backups at milestones, and warm-resume on crash. Survived two battery-death shutdowns with ≤200 iters lost.
• Long generation jobs must flush results incrementally (prep_rft.py appends each kept row to metadata_incremental.jsonl) — all-or-nothing output files turn any crash into a total loss.

Reproduce

python prep/prep_waveui.py             # one-time online download; offline afterwards
python prep/prep_osatlas.py            # desktop grounding -> data_mix
python prep/prep_aguvis_web.py         # web actions -> data_mix
python prep/audit_data.py              # sanity-check the data
python eval/eval_grounding.py --model mlx-community/Ministral-3-14B-Instruct-2512-4bit --n 300   # baseline
./train/train_resilient.sh data_mix adapters_mix 4000                                            # resilient SFT
python eval/eval_grounding.py --model <…> --lora-checkpoint adapters_mix/adapters.safetensors --data screenspot --n 150
python prep/prep_rft.py --adapter adapters_mix/adapters.safetensors --n-ground 500 --k 4 --n-act 500
python train/resume_train.py --model <…> --dataset data_rft --resume-from adapters_mix/adapters.safetensors \
  --iters 1500 --output-path adapters_rft/adapters.safetensors --grad-checkpoint --steps-per-save 200

Use --lora-checkpoint (quant-preserving) for evaluating on a 4-bit base — --adapter-path dequantizes to ~28 GB and thrashes.

Roadmap to a full CUA

1. Grounding (this repo) — where to click. Done: 22% → low-30s; ceiling identified.
2. Actions — richer trajectory SFT (drag/double_click/hover need desktop trajectory data).
3. Runtime — screenshot → model → execute loop + demo.

The identified next lever for grounding itself: a grounding-pretrained base (Qwen2.5-VL scores ~80% ScreenSpot zero-shot at half the size), an unfrozen vision tower (bf16), and 10–100× data — i.e., a cloud-GPU experiment, not more laptop iterations on this recipe.