Skip to content

mlx-bun v0.0.5

Choose a tag to compare

View all tags
@joshuarossi joshuarossi released this 19 Jun 16:06
v0.0.5
887df99

mlx-bun v0.0.5

Headline: ORPO LoRA training on Apple Silicon, with an [M,vocab]-free flash-CCE head that makes large-vocab preference fine-tuning feasible on a single Mac. Plus the new mlx-bun.dev docs site and one-command install.

✨ ORPO training stack (new)

Train LoRA adapters with ORPO (Odds Ratio Preference Optimization) — reference-free preference tuning that reuses the DPO {prompt, chosen, rejected} data format at half the forwards/step.

  • Flash-CCE head (forward + backward). The response-head logits are computed inside one Metal kernel built on MLX's verbatim steel quantized GEMM, so neither the [M,vocab] logits nor a dequantized head ever touch HBM. Backward 3687 → 754 ms on e4b (exact), peak 0.93 GB flat at M=8192 — large vocab (e4b 262k, MiniCPM5 130k) no longer dominates memory.
  • Prefix-sharing (orpo_prefix_shared) — one forward over [prompt; chosen; rejected] (block-sparse mask + block-wise RoPE), so a shared prompt is encoded once. Big win on prompt-dominant data (e.g. long-document preference pairs). Per-row two-forward fallback when chosen/rejected prompts differ.
  • Segmented backward (segment_size) — gradient-checkpointed layer activations for long context, composed with both the flash head and prefix-sharing.
  • Warm-start — continue a run from a checkpoint's LoRA weights (RESUME=<adapter-dir> on the launcher; warm_start_adapter in the job API).
  • Launcher — scripts/train-orpo.ts: full stack on by default, auto-detects e4b and sets its env, checkpoints every N steps, RESUME. See docs/reference/orpo-quickstart.md.

Supported scope (important): the full stack targets OptiQ-quantized (affine 4/8-bit) MiniCPM5 (Llama-arch) and Gemma-4 (e4b / 12B / 26B) models. Other architectures aren't wired for the segmented/prefix paths yet.

Validation: parity vs autograd — dh 0.40% (e4b) / 0.28% (MiniCPM5), bf16 class. Integration suite (tests/train-orpo-fused-ce.test.ts): flash / segmented+flash / prefix+flash / segmented+prefix+flash all train end-to-end. e4b @ 8192 full stack ≈ 13 GB / ~70 s/step on an M1 Max — the prior "e4b OOMs ≥ 2048" wall is gone.

📖 Docs site + reference

  • mlx-bun.dev — an Astro Starlight documentation site, deployed via Pages, with reference docs generated from source at build (no drift).
  • README corrected against source; experimental features labeled; KV-flag / library-API / benchmark sections fixed.
  • New training reference (docs/reference/training.md) + ORPO quickstart.

📦 Install / distribution

  • One-command install — mlx-bun.dev/install.sh (the recommended path), plus bunx and a direct-download option.
  • Homebrew tap — one-command publish pipeline that auto-syncs the tap.
  • Fix: disable-library-validation so brew-relocated dylibs load.

⚠️ Notes

  • This release ships the ORPO trainer (the recipe), not a pre-trained adapter — you bring your own preference data and a supported base model.
  • Apple-CCE backward skips (MLX_BUN_CCE_BWD_FILTER_EPS / _BLOCK_EPS) are off by default — exact gradients; opt in only on genuinely peaked data.
  • Gemma e4b training requires MLX_BUN_PERF_KERNEL=0 + MLX_BUN_FUSED_GELU=0 (the launcher sets these automatically).

Full changelog: v0.0.4...v0.0.5

Assets 5
Loading