1e20366a26
Re-apply the subcommand restructure that was lost during the earlier
rebase conflict (the linter pre-commit on the remote re-formatted the
single-function version, so my checkout --ours kept the wrong copy).
Adds:
* argparse subcommands `train` and `reload --format X --dir D` so
each reload runs in a FRESH Python process the way real users
hit the cold-start path.
* Per-phase Phase() context manager records elapsed wall-clock,
peak GPU memory (mx.metal.get_peak_memory), and peak RSS
(resource.getrusage) into a metrics dict written to
{train,lora_reload,merged_reload,gguf_reload}_metrics.json
next to the saved dir for cross-CI regression detection.
* batch_size=2, gradient_accumulation_steps=3 (was 2/1) so the
7-step run sees 42 sequences total.
* GGUF save is best-effort. unsloth-zoo#627 fixed the
NotImplementedError on Apple Silicon, but llama.cpp's
convert_hf_to_gguf currently asserts on the gemma-3-270m
tokenizer vocab (`max(vocab IDs) >= vocab_size`). That's a
downstream llama.cpp limitation, not an unsloth_zoo bug, so the
train step records gguf_supported=false + the reason instead of
raising, and the GGUF reload step emits a workflow warning and
exits 0. The LoRA + merged_16bit reload assertions remain the
gating signal.
The earlier-draft LoRA workaround that copied base config.json into
the LoRA save dir is removed; unsloth-zoo#627 makes
FastMLXModel.from_pretrained(lora_dir) work on the saved adapter
directory directly (the failing run before #627 confirmed the bug,
the run after #627 lands shows the adapter is detected and the base
model is pulled from adapter_config.json:base_model_name_or_path).
291 lines
13 KiB
YAML
291 lines
13 KiB
YAML
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
# Focused PR gate for the MLX dispatch surface, running on a real
|
|
# Apple Silicon runner.
|
|
#
|
|
# Runner: macos-14 (M1, 3 vCPU / 7 GB / Apple Silicon standard runner
|
|
# -- FREE for public repositories per the GitHub Actions billing
|
|
# reference; larger variants like macos-14-large/-xlarge are paid so
|
|
# we deliberately avoid those).
|
|
#
|
|
# Why a single Mac job (no Linux+spoof leg): the dispatch tests are
|
|
# 100% spoofed monkeypatches and run identically on any host, so the
|
|
# Linux leg was duplicating the matrix tests already covered on Mac
|
|
# while missing everything Apple-specific. The Mac job runs the SAME
|
|
# spoofed matrix PLUS three things only a real Apple Silicon host
|
|
# can prove:
|
|
#
|
|
# 1. unsloth._IS_MLX flips True on Darwin+arm64 with mlx genuinely
|
|
# installed (no spoof).
|
|
# 2. Every PR-A MLX-only unsloth_zoo module (mlx_loader, mlx_trainer,
|
|
# mlx_compile, mlx_utils, mlx_cce, gated_delta_vjp) imports
|
|
# against the real `mlx` + `mlx-lm` + `mlx-vlm` PyPI wheels --
|
|
# each does `import mlx.core as mx` at module top level, so this
|
|
# catches a future change that breaks the real wheels without
|
|
# needing a Mac developer in the loop.
|
|
# 3. The hardware-dispatch spoofs do not collide with the real
|
|
# environment (the test fixture installs a MetaPathFinder that
|
|
# blocks `import mlx.core` for "no-mlx" profiles, faithfully
|
|
# simulating a Mac without mlx even when mlx IS installed).
|
|
# 4. End-to-end MLX training + inference smoke test:
|
|
# run_real_mlx_smoke.py trains unsloth/gemma-3-270m-it for 7
|
|
# deterministic LoRA steps on a single repeated text row, then
|
|
# verifies the trained model can complete the prompt and that
|
|
# losses + grad norms are finite and well-behaved. This is the
|
|
# only place in CI that exercises a real MLX backward pass +
|
|
# optimizer step + inference call.
|
|
#
|
|
# Three dispatch test files documented in tests/studio/README.md:
|
|
# - test_hardware_dispatch_matrix.py parametrized 7-profile matrix
|
|
# + 2 dispatch-priority canaries
|
|
# - test_is_mlx_dispatch_gate.py AST + runtime guard on
|
|
# unsloth._IS_MLX
|
|
# - test_mlx_training_worker_behaviors.py AST contract checks on
|
|
# studio/backend/core/training/worker.py
|
|
#
|
|
# Surfaces a single PR check ("MLX CI on Mac M1 / dispatch").
|
|
#
|
|
# Security audit footprint: every package this workflow installs is
|
|
# already covered by .github/workflows/security-audit.yml -- the deps
|
|
# come from studio/backend/requirements/studio.txt and unsloth-zoo's
|
|
# pyproject (resolved transitively). The git+ install of unsloth-zoo
|
|
# is intentionally skipped by the audit (pip-audit cannot resolve a
|
|
# git URL through PyPI metadata; the audit comment in security-audit.yml
|
|
# documents this). No new package is introduced solely by MLX CI.
|
|
|
|
name: MLX CI on Mac M1
|
|
|
|
on:
|
|
pull_request:
|
|
paths:
|
|
- 'unsloth/__init__.py'
|
|
- 'unsloth/_gpu_init.py'
|
|
- 'studio/backend/utils/hardware/**'
|
|
- 'studio/backend/core/training/worker.py'
|
|
- 'studio/backend/core/inference/mlx_inference.py'
|
|
- 'tests/studio/test_hardware_dispatch_matrix.py'
|
|
- 'tests/studio/test_is_mlx_dispatch_gate.py'
|
|
- 'tests/studio/test_mlx_training_worker_behaviors.py'
|
|
- 'tests/studio/run_real_mlx_smoke.py'
|
|
- 'tests/conftest.py'
|
|
- '.github/workflows/mlx-ci.yml'
|
|
push:
|
|
branches: [main, pip]
|
|
|
|
concurrency:
|
|
group: ${{ github.workflow }}-${{ github.ref }}
|
|
cancel-in-progress: true
|
|
|
|
permissions:
|
|
contents: read
|
|
|
|
jobs:
|
|
dispatch:
|
|
name: dispatch
|
|
runs-on: macos-14
|
|
# 25 min: dispatch + spoofed matrix + 7-step real LoRA training is
|
|
# under 2 min; GGUF export builds llama.cpp via cmake on Apple
|
|
# Silicon (~5-7 min), so we budget headroom.
|
|
timeout-minutes: 25
|
|
steps:
|
|
- uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1
|
|
|
|
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
|
with:
|
|
python-version: '3.12'
|
|
cache: 'pip'
|
|
|
|
# macOS install ladder, validated locally against a Linux
|
|
# mac-sim venv (platform spoofed + mlx_simulation shim + real
|
|
# datasets/transformers/structlog).
|
|
#
|
|
# 1. studio/backend/requirements/studio.txt brings structlog,
|
|
# fastapi, etc. The hardware probe imports structlog at
|
|
# module top level.
|
|
# 2. Same pytest / numpy / httpx stack the rest of the repo CI
|
|
# uses.
|
|
# 3. torch is explicitly installed: unsloth-zoo's pyproject
|
|
# deliberately excludes torch on darwin+arm64 (mlx replaces
|
|
# it for runtime use), but the dispatch tests spoof
|
|
# torch.cuda / torch.xpu / torch.backends.mps via monkeypatch
|
|
# and so the test process needs torch importable. We pull
|
|
# from the PyTorch CPU index so Apple Silicon gets the
|
|
# explicit cpu+MPS arm64 wheel rather than something the
|
|
# default PyPI resolver might pick up. The CPU index hosts
|
|
# macosx_*_arm64 wheels alongside the Linux x86_64 ones.
|
|
# 4. unsloth-zoo from git main (NOT PyPI), WITH deps. PR-A's
|
|
# MLX support landed after the most recent unsloth-zoo PyPI
|
|
# release; the wheel still raises NotImplementedError on
|
|
# Apple Silicon when device_type.get_device_type() runs
|
|
# unguarded. Studio's own install.sh overlays unsloth-zoo
|
|
# from git main for the same reason. Pulling deps lets pip
|
|
# resolve the platform-conditional MLX-only wheels (mlx,
|
|
# mlx-lm, mlx-vlm gated on darwin+arm64 in unsloth-zoo's
|
|
# pyproject) AND the shared deps (datasets, transformers,
|
|
# sentencepiece, ...) that unsloth's MLX branch loads via
|
|
# dataprep/raw_text.py.
|
|
# 5. unsloth -e . --no-deps so the editable install does not
|
|
# fight the unsloth-zoo dep set.
|
|
#
|
|
# All explicit pip installs are version-pinned to a single
|
|
# released version (the latest as of 2026-05-07 within each
|
|
# project's existing constraint range). bump alongside the rest
|
|
# of the security audit when a new release lands.
|
|
- name: Install deps
|
|
run: |
|
|
python -m pip install --upgrade pip
|
|
pip install -r studio/backend/requirements/studio.txt
|
|
pip install \
|
|
'python-multipart==0.0.27' \
|
|
'aiofiles==25.1.0' \
|
|
'sqlalchemy==2.0.49' \
|
|
'cryptography==48.0.0' \
|
|
'pyyaml==6.0.3' \
|
|
'jinja2==3.1.6' \
|
|
'mammoth==1.12.0' \
|
|
'unpdf==1.0.0' \
|
|
'requests==2.33.1' \
|
|
'typer==0.25.1' \
|
|
'numpy==2.4.4' \
|
|
'pytest==9.0.3' \
|
|
'pytest-asyncio==1.3.0' \
|
|
'httpx==0.28.1'
|
|
pip install --index-url https://download.pytorch.org/whl/cpu \
|
|
'torch==2.10.0'
|
|
# Pin to the unslothai/unsloth-zoo#627 branch until that PR
|
|
# merges -- it carries the GGUF NotImplementedError + LoRA
|
|
# local_path fixes that the export round-trip phase below
|
|
# depends on. Flip back to bare `unsloth-zoo @ git+...` once
|
|
# #627 lands on main.
|
|
pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@fix/mlx-export-roundtrip-on-apple-silicon"
|
|
pip install -e . --no-deps
|
|
|
|
# Real Apple Silicon sanity: confirm _IS_MLX activates on real
|
|
# hardware with no platform spoof.
|
|
- name: Verify _IS_MLX flips True on real Apple Silicon
|
|
run: |
|
|
python -c "
|
|
import platform
|
|
assert platform.system() == 'Darwin', platform.system()
|
|
assert platform.machine() == 'arm64', platform.machine()
|
|
import unsloth
|
|
assert unsloth._IS_MLX is True, f'expected _IS_MLX=True on real Apple Silicon, got {unsloth._IS_MLX}'
|
|
print('OK: _IS_MLX activated on real Apple Silicon')
|
|
"
|
|
|
|
# Real Apple Silicon sanity: confirm every PR-A MLX-only module
|
|
# loads against real mlx + mlx-lm + mlx-vlm wheels.
|
|
- name: Smoke-import every MLX-only unsloth_zoo module
|
|
run: |
|
|
python -c "
|
|
import importlib
|
|
for name in [
|
|
'unsloth_zoo.mlx_loader',
|
|
'unsloth_zoo.mlx_trainer',
|
|
'unsloth_zoo.mlx_compile',
|
|
'unsloth_zoo.mlx_utils',
|
|
'unsloth_zoo.mlx_cce',
|
|
'unsloth_zoo.gated_delta_vjp',
|
|
]:
|
|
importlib.import_module(name)
|
|
print('OK:', name)
|
|
from unsloth_zoo.mlx_loader import FastMLXModel
|
|
from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig
|
|
assert hasattr(FastMLXModel, 'from_pretrained')
|
|
print('OK: FastMLXModel + MLXTrainer surface present')
|
|
"
|
|
|
|
# Spoofed dispatch matrix. Runs on the real Mac too -- the
|
|
# test fixture installs a MetaPathFinder that blocks
|
|
# `import mlx.core` for "no-mlx" profiles, so the spoofs
|
|
# faithfully simulate every supported hardware combo regardless
|
|
# of whether mlx is installed for real.
|
|
- name: MLX dispatch tests (3 files, 36 tests)
|
|
env:
|
|
PYTHONPATH: ${{ github.workspace }}/studio
|
|
UNSLOTH_COMPILE_DISABLE: '1'
|
|
run: |
|
|
python -m pytest -v --tb=short \
|
|
tests/studio/test_hardware_dispatch_matrix.py \
|
|
tests/studio/test_is_mlx_dispatch_gate.py \
|
|
tests/studio/test_mlx_training_worker_behaviors.py
|
|
|
|
# Real MLX training + inference smoke test. Trains
|
|
# unsloth/gemma-3-270m-it for 7 deterministic LoRA steps
|
|
# (batch_size=2, gradient_accumulation_steps=3) on a single
|
|
# repeated row ("<<HELLO!!>> My name is Unsloth!"), then saves
|
|
# the trained model in 3 export formats. The `train` subcommand
|
|
# captures per-phase timing + peak GPU + peak RSS into
|
|
# train_metrics.json so we can detect regressions across CI runs.
|
|
- name: MLX export round-trip — TRAIN + SAVE 3 formats
|
|
env:
|
|
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
|
UNSLOTH_COMPILE_DISABLE: '1'
|
|
run: |
|
|
mkdir -p mlx_workdir
|
|
python tests/studio/run_real_mlx_smoke.py train \
|
|
--workdir "$PWD/mlx_workdir"
|
|
|
|
# Each reload step runs in a FRESH Python process to confirm
|
|
# the cold-start path users would hit in production also works
|
|
# (not just the in-memory continuation of a still-running
|
|
# trainer). FastMLXModel.from_pretrained gets called from
|
|
# scratch; mx.random is re-seeded; per-step timing + peak
|
|
# memory are emitted to {format}_reload_metrics.json next to
|
|
# the saved dir.
|
|
- name: MLX export round-trip — RELOAD LoRA (fresh process)
|
|
env:
|
|
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
|
UNSLOTH_COMPILE_DISABLE: '1'
|
|
run: |
|
|
python tests/studio/run_real_mlx_smoke.py reload \
|
|
--format lora \
|
|
--dir "$PWD/mlx_workdir/lora"
|
|
|
|
- name: MLX export round-trip — RELOAD merged_16bit (fresh process)
|
|
env:
|
|
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
|
UNSLOTH_COMPILE_DISABLE: '1'
|
|
run: |
|
|
python tests/studio/run_real_mlx_smoke.py reload \
|
|
--format merged \
|
|
--dir "$PWD/mlx_workdir/merged_16bit"
|
|
|
|
# GGUF reload uses the llama-cli binary that save_pretrained_gguf
|
|
# built. If save_pretrained_gguf was skipped during train (e.g.
|
|
# llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer
|
|
# vocab -- a downstream llama.cpp limitation, not an unsloth_zoo
|
|
# bug), this step emits a workflow warning and exits 0 so the
|
|
# LoRA + merged_16bit assertions remain the gating signal.
|
|
- name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process)
|
|
env:
|
|
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
|
run: |
|
|
if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then
|
|
python tests/studio/run_real_mlx_smoke.py reload \
|
|
--format gguf \
|
|
--dir "$PWD/mlx_workdir/gguf"
|
|
else
|
|
REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')")
|
|
echo "::warning title=GGUF round-trip skipped::${REASON}"
|
|
echo "GGUF export was skipped during the train phase. Reason:"
|
|
echo " ${REASON}"
|
|
echo "Continuing without failing the job; the LoRA + merged_16bit"
|
|
echo "reload assertions are still gating this PR."
|
|
fi
|
|
|
|
# Print all metrics JSON files so regressions are visible in the
|
|
# job log. always() so we get telemetry even if a reload step
|
|
# asserted gibberish.
|
|
- name: MLX export round-trip — aggregate metrics
|
|
if: always()
|
|
run: |
|
|
for f in mlx_workdir/train_metrics.json \
|
|
mlx_workdir/lora_reload_metrics.json \
|
|
mlx_workdir/merged_reload_metrics.json \
|
|
mlx_workdir/gguf_reload_metrics.json; do
|
|
echo "=== $f ==="
|
|
cat "$f" 2>/dev/null || echo "(missing)"
|
|
echo
|
|
done
|