vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-08-25 08:42:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29
unsloth/.github/workflows/mlx-ci.yml
Long Yixing 0998656891
Studio: honor the OpenAI sampling parameters on the MLX backend (#9262) 
* feat(studio): honor per-request seed on the MLX backend

Studio accepts an OpenAI `seed` in the chat-completion schema but only ever forwarded it to llama-server, so a seeded request against an MLX model — including a safetensors model, which on Apple Silicon is served by the MLX backend — sampled non-deterministically with no indication the field had been dropped.

Seeding uses a request-scoped `mx.random.key` rather than `mx.random.seed`. Global seeding mutates thread-local state that later requests inherit, so a request that sets no seed but follows a seeded one would silently become reproducible; `mx.random.state` cannot be reliably saved and restored either. A per-request key keeps determinism inside the request that asked for it, and works identically on text and vision without depending on the mlx-vlm version.

The seeded sampler reuses mlx-lm's own top-p/min-p/top-k stages instead of reimplementing them. Supplying a sampler suppresses the chain mlx-lm and mlx-vlm would otherwise build — mlx-vlm constructs its temperature/top-p/min-p/top-k sampler only when no sampler is passed — so anything not reused here would be silently dropped from seeded requests alone.

Seed values are reduced modulo 2**64 on the MLX path. The field is shared with llama-server, which forwards every non-null seed including -1, while `mx.random.key` rejects negatives and anything at or above 2**64; normalizing at the point of use keeps the shared schema unnarrowed so no request that works today starts failing.

Forwarding is gated on the backend declaring a seed parameter. The transformers backend has neither `seed` nor `**kwargs`, so forwarding unconditionally would turn its existing "ignores the seed" behavior into a TypeError. External providers are untouched.

Validated against mlx-community/Qwen2.5-0.5B-Instruct-4bit and mlx-community/Qwen2-VL-2B-Instruct-4bit on mlx-vlm 0.4.4: identical output for a repeated seed, differing output across seeds, -1 equivalent to 2**64-1, boundary seeds (0, 2**64, -2**63, 2**70) all completing, the sampling chain still applied under a seed, and an unseeded request following a seeded one still varying across runs.

* feat(studio): honor frequency_penalty and logit_bias on the MLX and GGUF backends

Both fields were accepted by the chat-completions schema and then discarded by every local path, so a request setting either sampled exactly like one that did not. They now reach MLX text and vision, and llama-server through both the direct and passthrough request paths.

Where Studio samples -- the MLX backends -- semantics follow the presence penalty already implemented here: score the whole completion, exclude the prompt. mlx-lm's own presence and frequency processors window the last 20 tokens *including* the prompt, so using them would make the same request sample differently depending on the backend. Occurrence counts are scaled in float32 and cast once, matching llama-server; casting counts into a float16 logits dtype first rounds them above 2048 and overflows past 65504 repeats. llama-server is not made to match: Studio only forwards its native penalty fields, and it seeds penalty history with the prompt and windows it, so an identical request may still sample differently there.

`logit_bias` ids outside the model's logit width are dropped rather than rejected, because the vocabulary a client biases against is the tokenizer's, which can be narrower. On MLX that bound is load-bearing: MLX does no bounds checking, so indexing a stray id is undefined behavior. The schema no longer constrains ids to be non-negative either, since a negative id is just another id the processors discard.

`presence_penalty` now accepts OpenAI's documented [-2, 2] instead of [0, 2]. Widening a bound only admits requests that were previously rejected, so nothing that works today changes; a negative value boosts repetition, which every backend that reads it honors.

The transformers backend is deliberately untouched. It declares neither field and takes no `**kwargs`, so the worker forwards both only to a backend whose signature declares them, exactly as it already does for the per-request seed: a safetensors model on Apple Silicon is served by the MLX backend and honors them, while on a CUDA host the fields stay accepted and ignored as they were before.

Each tool-calling wrapper and payload builder re-declares its own argument list, which is how these fields came to be accepted and then dropped in the first place.

Validated on mlx-community/Qwen2.5-0.5B-Instruct-4bit under a fixed seed: a request setting neither field is byte-identical to the previous behavior, each knob changes the output, and out-of-range bias ids are dropped without crashing. Presence charges a token repeated three times once; frequency charges it three times; prompt tokens are untouched by both.

* fix(studio): report a truthful finish_reason on the MLX backends

Every MLX completion reported `finish_reason="stop"`, including one cut off at `max_tokens`. A client had no way to tell a finished answer from a truncated one, and a truncated tool call looked like a complete one. The transformers backend has the same defect and is deliberately left alone here; it is unreachable on Apple Silicon, where a safetensors model is served by MLX.

The reason is now derived where generation ends and travels on the existing generation-stats channel, so it reaches all four chat response sites: plain streaming and non-streaming, and both of the server-side tool loop's. mlx-lm reports the reason directly. mlx-vlm does not, and at the limit a token count cannot separate a stop token sampled as the final allowed token from ordinary exhaustion, so that case is decided by the last token's identity.

Which ids count as stopping is taken from what the runtime actually stops on: the tokenizer's stopping criteria first, then the model config that seeds them, then the tokenizer attribute. These disagree in practice — Kimi-VL's config lists two ids and its tokenizer reports a different one, so reading the tokenizer alone would report a real stop as truncation — and each source may be a bare int (Qwen2-VL) or a collection.

Truncation outranks the tool-call upgrade, matching the rule the GGUF path already documents: a call cut off at `max_tokens` keeps `"length"` so the client knows the arguments may be incomplete, while the healed call stays attached. All four sites share one resolver instead of repeating the policy inline.

An explicit Stop is not a completion, so the streams no longer synthesize a terminal chunk for one. Cancellation is latched once where generation exits rather than reread at the emit site: a Stop arriving while the finalization chunks are still being yielded must not suppress the terminal chunk of a completion that already finished, and the same latched value now decides the recorded request status. Non-streaming cancellation is unchanged either way. The plain handler drains its generator synchronously, so a cancel cannot be delivered to it in-process at all. The server-side tool loop drains through a worker thread and does observe one -- it records the request as cancelled -- but still answers with the partial text under `stop`, because this field's five values have no member for it. Both behave exactly as they did before this change; giving a cancelled reply a truthful terminal state needs a contract these values do not yet have.

Because the reason rides on the stats the nudge retry already saves and restores, a discarded retry cannot leave its terminal state behind.

The audio path is untouched: it generates greedily and reports its own terminal state.

Validated on mlx-community/Qwen2.5-0.5B-Instruct-4bit and mlx-community/Qwen2-VL-2B-Instruct-4bit: truncation at 10 and 12 tokens reports `length`, natural termination reports `stop`, and the vision runtime resolves its stopping ids from the integer-valued source that previously raised.

* feat(studio): honor stop sequences on the MLX backends

The request schema advertises `stop` and the llama-server path applies it, but the MLX backends had no such parameter, so the same request returned different text per backend with no error.

mlx-lm and mlx-vlm stop on token ids only, so a sequence spanning two tokens -- or spelled by a tokenization the vocabulary does not use -- never ends the turn there. Both generators yield cumulative snapshots, so the sequence is found in the decoded text instead, and the reply is cut before it. Text that could still grow into a sequence is held back for the same reason llama-server holds it: a client cannot unsee a fragment that the next token completes. A trailing replacement character is never matched, even once the turn has ended. A decode prints the same character for one the model wrote and for bytes that never finished arriving, and neither runtime resolves that: mlx-lm decodes what is left over with `errors="replace"` when generation stops, so its closing text carries the stand-in too, and mlx-vlm suppresses any text ending in the character at all. A sequence that can only match on such a character therefore does not match, while the character itself is delivered exactly as it is when no sequence was asked for. That cost falls on a caller that names one; the alternative, ending a turn on a character the model may never have written, would fall on every reply cut off mid-character. A turn that ends without the sequence arriving hands that fragment back as the ordinary text it turned out to be, so asking to stop on a sequence the model never writes returns exactly what the request would have returned without it.

The text path is delivered once when a request names sequences, rather than snapshot by snapshot. It re-decodes every id each step instead of appending, so a later snapshot can revise characters an earlier one showed -- joining two tokens into a word neither spelled alone. Its consumers, the streaming route and the tool loop both, turn snapshots into deltas by length, so handing them a revised snapshot splices two renderings into text no snapshot contained, which can spell the very sequence the request asked to stop on. Nothing can unsend that, so no sampled text is sent until the turn's text has settled -- the prompt's own `<think>` prefill still goes out first, since the prompt cannot be revised by a later token. A caller streaming such a request therefore sees the reply arrive at the end of the turn rather than as it is written. The alternative is to stream this path from an append-only decode instead, which would remove the revisions rather than wait them out; that changes how every plain reply on this path is delivered, including requests naming none of these parameters, so it belongs to its own change rather than to this one. Requests naming no sequences stream as before, as do the paths whose text only ever grows: the native reasoning channel, vision, and audio input.

Audio-input chat is a separate parameter chain that ignored `stop` for the same reason, and it forwards incremental deltas rather than the snapshots the text and vision paths diff, so its deltas are reassembled before the same matching runs over them. Its orchestrator command carries the sequences and the worker forwards them only to a backend whose signature declares the option, as the text branch already does, so a backend that never took one keeps ignoring it rather than raising TypeError.

A stop sequence applies to the completion the model produced, so each producer matches its own sampled text -- ahead of both rewrites this backend performs on the way out. One puts the open `<think>` a reasoning template left at the end of the prompt back in front of the reply, so the frontend renders the block; the other rewrites a model's native reasoning markers as `<think>`. Matching the delivered text instead would end a turn on markup the model never wrote, so `stop: ["<think>"]` would return nothing at all on a reasoning template, and would never find the native marker a caller did ask to stop on. The cut runs before the reasoning normalizer, so no sampled sequence reaches it; the markers the normalizer itself writes are this layer's own, and go unmatched for the same reason the prefill does.

The sequences are normalized once, above the branch that picks a backend, so the paths that return from inside it cannot disagree with the text paths about what the request asked for.

A turn cut this way reports `finish_reason: "stop"` rather than the reason the runtime would have named for the token it happened to be on, and a reasoning block the sequence landed inside is closed as any finished one is: the closing marker is this layer's to emit, not the model's, so leaving it out would hand the client markup it cannot render. Only a cancelled turn is still left unfinished, because there more text was coming.

* feat(studio): serve n > 1 on the non-GGUF chat path

The llama-server path drains `n` sequential generations for a non-streaming request; the non-GGUF path rejected `n > 1` outright, so the same request succeeded on one backend and 400'd on the other.

Plain text generation now serves extra choices the same way: one full generation per choice, stopping early if the turn is cancelled. The prompt is shared, so its tokens are counted once and only generated tokens accumulate, matching what the llama-server drain reports.

The rejection moves down onto the paths that answer something other than a sample -- an audio turn returns one waveform, a transcription one transcript, and the tool loop runs turns rather than samples -- and streaming keeps its existing refusal.

* feat(studio): declare response_format and refuse it where it cannot be served

`response_format` reached the request only through `extra="allow"`, so it was absent from the schema and invisible to clients reading it. Declaring it documents the contract and fixes its shape, so a value of the wrong kind entirely -- a bare string, say -- is turned away rather than reaching a backend as junk. What is inside the object stays the serving path's business. The router reads the value by attribute rather than off the chat model's field, so a request type that does not declare one is read the same way.

Guided decoding is the llama.cpp grammar engine's. A non-GGUF backend cannot constrain decoding, and previously dropped the field and answered with free text that silently violated the contract -- worse than refusing, because the client has no way to tell. It now returns 400 naming the parameter. Only `{"type": "text"}` is exempt: it is the default spelled out, constrains nothing, and is served exactly as a request naming no format. Every other object is a contract, including one whose type this build does not recognize and a text format carrying members it does not know -- guessing that an unfamiliar object constrains nothing is how a contract gets dropped silently again. The branch that answers speech refuses it as well, and says so where it stands: it returns before the routing below it, so nothing there can decide the contract's fate, and no grammar constrains a waveform. The GGUF paths that are not the passthrough refuse it too. Only an explicit request for Unsloth's tool loop reaches them carrying a contract, and neither that loop, which runs its own turns and forwards none, nor the ordinary generator below it, which has no such parameter, can keep one -- including when the loop's tool selection resolves empty and the request falls through to that generator.

The GGUF router asks the same question. Selecting the llama-server passthrough on the mere presence of the field sent an unconstrained request down a path whose only purpose is the grammar engine, and that path serves strictly less: a request naming `{"type": "text"}` was refused the n > 1 the ordinary path answers. An unrecognized type still goes there, where llama-server answers or rejects it as it did before.

* fix(studio): report token usage on the non-GGUF non-streaming replies

`ChatCompletion.usage` defaults to a zero-filled object, so a body built without it reports zeros rather than nothing -- a client cannot tell the difference between a real count and a missing one. All three non-streaming non-GGUF shapes -- the plain reply, the server-side tool loop, and the audio-input reply -- left it at that default while the counts they needed were already in hand, in the same stats the monitor reads from.

The plain reply now builds one usage object, so whatever the monitor is told is that same object and a client and the speed popover cannot disagree about the same turn. The monitor is still told nothing when no choice published stats, which is how it already reported a turn whose token count is unknown. For n > 1 the object carries the totals the llama-server drain reports -- the shared prompt counted once, generated tokens summed -- and drops the timings, which describe only the last choice and so describe neither.

The tool loop reports what the whole loop spent. Each of its turns published into the one holder, so the reply described only the last turn and hid the tokens that produced the tool call; the turns are summed instead, as the llama.cpp tool loop already counts them, with the prompt taken from the last turn since it already contains the tool results the earlier turns produced, and the reported rates recomputed from those totals rather than left describing the turn they arrived with. A turn is folded in when it reports, so a loop a cancel cuts short still reports the turns that finished: the interrupted turn's own report is discarded with the rest of its cancelled output, and folds in as nothing.

* Fix test loss, finish reason and usage gaps in PR #9262

Follow-ups found while reviewing the branch. All of them sit on top of the
existing commits rather than changing their direction.

test_vlm_seed_rides_on_the_sampler_not_a_seed_kwarg lost its body: the penalty
commit appended its parametrize block onto the docstring, so the one test
covering the seed design has been passing vacuously since. Restored, and rebuilt
so it runs off Apple Silicon too by recording what reaches the sampler builder
rather than calling into mlx.

The two numerical processor tests importorskip mlx, so they skip on the ubuntu
backend CI and mlx-ci.yml never collected studio/backend/tests at all. That left
the seed and penalty work with no executed coverage in any job. mlx-ci.yml now
runs test_mlx_inference_backend.py on the Mac runner.

_summed_tool_loop_stats had no test anywhere: replacing its body with
`return turn`, which is the exact bug it exists to prevent, passed the whole
backend suite. Added direct tests over the fold, and fixed two cases they found.
A final turn carrying usage but no timings dropped the loop's accumulated
timings, and one carrying timings but no usage dropped prompt_tokens entirely
and collapsed total_tokens to the summed completions. llama.cpp folds these
unconditionally and does not have the hole. completion_tokens_details now sums
with the completion it describes instead of being carried from the last turn
against a count it does not cover.

The pre-switch tool-policy gate still read response_format with bool(), while
the router reads it with _response_format_constrains_decoding. The two disagreed
in both directions: {"type": "text"} passed the pre-switch guard and was then
rejected on the GGUF local path after an auto-switch may have reloaded the
resident model, which is what that guard exists to prevent, and {} took the
local-confirm rejection with the wrong param name. Both use the router's
predicate now, so they cannot diverge.

A successful nudge retry reported only the retry's usage. Both attempts were
generated, and this branch is what promotes that number from the monitor into
the OpenAI usage field a caller bills against, so both are now folded, with the
winning attempt's prompt. The discarded-retry branch is folded the same way for
symmetry; its test is renamed and keeps asserting the delivered attempt's prompt.

The MLX audio-input producer never derived a finish_reason, so a turn that
exhausted max_tokens reported "stop". Pre-existing, but text and vision are
fixed here and leaving audio out makes the three siblings disagree.

_mlx_stop_token_ids returned () when its first source was present but empty,
never consulted a singular tokenizer.eos_token_id, and could raise on a numpy
scalar from inside a caller's finally. An empty stop set misreads every real
stop as truncation. Token id 0 is kept, since 0 is a valid id.

The VLM path left a reasoning block unclosed when a stop sequence and a cancel
landed on the same step; the text and audio paths already guarded for this.

Also corrected the n loop comment: the non-GGUF drain guards index 0 so a
cancelled turn still answers with one choice, which the llama-server drain it
claimed parity with does not do.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix nudge retry double-billing and report MLX cached prompt tokens

Both from the Codex round on 5cd980b6f.

The nudge retry saved the first attempt's stats with a read, not a take, so a
retry that raised before publishing its own report left stats_holder pointing at
that same object. The exception branch then folded it into itself and billed its
completion tokens twice: a 3-token first attempt was reported as 6. Popped now,
so a fresh report is distinguishable from the stale one and every branch folds
at most two distinct reports.

MLX puts its reused prefix inside usage.prompt_tokens but named the count only
under timings.cache_n, so the usage object this branch newly returns reported
cached_tokens 0 on every prompt-cache hit even though the backend had the number.
llama-server already forwards the real count, so the two backends disagreed on a
field a caller reads for cost attribution. _build_generation_stats now names it
where the rest of the usage lives.

That surfaced a latent split in the n loop: prompt_tokens was chosen
last-truthy-wins while prompt_tokens_details was chosen first-non-None-wins, two
rules over the same shared quantity. With MLX previously reporting no details it
was unreachable; naming cached_tokens makes it reachable, and the nudge does
rebuild a longer prompt for some choices, so a reply could claim more cached
tokens than prompt tokens. Both now come from the choice that supplied the count.
The same pairing is applied inside the tool-loop fold, where a turn that reports
no prompt count of its own inherits the loop's details along with it.

Three regression tests, each mutation-checked: reverting the pop doubles the
count, splitting the selection rules reproduces cached_tokens 1100 against
prompt_tokens 1000, and dropping the source field empties the details.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Read MLX EOS ids from every shape a model config arrives in

From the Codex round on a12a58f64.

A checkpoint's config is a dict on some models and an object on others, and it
sits under config or _config: exactly the spread _mlx_vlm_model_config already
walks to find a model_type, and the shape this file's own tests use. The EOS
lookup read it with getattr only, so the dict half returned nothing and the
search fell through to the tokenizer, which is the source the priority order
exists to outrank when the two disagree.

Reproduced: with a dict config listing [1, 2] and a tokenizer saying 9, the ids
resolve to (9,) and an EOS of 2 sampled on the last allowed token reports
finish_reason "length" instead of "stop", which is the reporting bug
02cc35642 exists to fix.

_mlx_config_field reads either representation across both attributes. The
tokenizer stays the fallback when no config carries the field. Parametrized
over all five shapes and mutation-checked: restoring the getattr-only read
fails four of the six cases.

* Tighten the comments this branch adds

Post-convergence comment pass over the comments added here, not the files they
sit in. Comments only, AST-verified against 507c1e76c with docstring edits
allowed: the one docstring touched is on an internal helper, not a route.

* Offset every fixed GGUF seed across choices, not just the non-negative ones

From the Codex round on 7cef80150.

_choice_seed exempted every negative seed from the per-choice offset on the
llama-server drain, on the premise that llama-server reads a negative seed as
"pick one at random". That is true of -1 only. The sampler seed is a uint32
(common/common.h: `uint32_t seed = LLAMA_DEFAULT_SEED`) and exactly one value,
LLAMA_DEFAULT_SEED = 0xFFFFFFFF (include/llama.h), selects a random draw; the
server README documents the sentinel as -1 specifically. -1 is the only request
seed that converts to it. -2 lands on 0xFFFFFFFE, an ordinary fixed seed, so a
GGUF request with seed -2 and n 3 sent the same fixed seed three times and
returned three copies of one run.

Exempting only -1 fixes that but introduces a second bug: the offset then walks
onto the sentinel, so seed -2 with n 3 serves -2, -1, 0 and the middle choice
samples at random where the caller asked for a fixed run. The offset is now
taken in the same uint32 domain llama-server reads, modulo 0xFFFFFFFF, so the
sentinel is unreachable by offsetting and consecutive choices stay distinct.

Non-negative seeds below 2^32 are unchanged, choice 0 is still the caller's own
seed, and the MLX side is untouched: it maps every seed onto its key domain, so
nothing is exempt there.

Regression test covers both failure modes; restoring `seed < 0` and dropping the
uint32 offset each fail it.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-08-24 02:55:03 -07:00

591 lines
29 KiB
YAML
Raw Permalink Blame History

# 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-26 (Apple Silicon standard runner, 3 vCPU / 7 GB
# -- FREE for public repositories per the GitHub Actions billing
# reference; larger -large/-xlarge variants 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. The Mac capability verdict /api/health publishes. Whether
# the sidebar's Train and Video rows are enabled, greyed out
# behind "Training needs MLX", or still spinning is decided by
# that one reply, and the decision reads a real `import
# mlx.core` and a real background reinstall. The unit tests
# drive it with a fake clock and a stand-in worker; only here
# is the host Apple Silicon and the MLX stack genuinely
# present or genuinely absent. mac_capability_verdict_smoke.py
# boots the server and judges every reply from the first one,
# because the reported bug was a verdict that was wrong for a
# while and right afterwards.
# 5. 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.
#
# Dispatch test files documented in tests/studio/README.md:
# - test_mlx_training_worker_behaviors.py AST contract checks on
# studio/backend/core/training/worker.py
# Run here because no other workflow claims it.
#
# test_hardware_dispatch_matrix.py and test_is_mlx_dispatch_gate.py are NOT run
# here. Both are fully spoofed (monkeypatched profiles plus a MetaPathFinder
# that blocks `import mlx.core`), so they assert identically on any host, and
# studio-backend-ci.yml runs both on ubuntu-latest. macOS concurrency is capped
# at 5 jobs account-wide; re-running host-independent tests on that pool is the
# most expensive way to learn nothing new.
#
# 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'
# The capability verdict is decided across these three: hardware/** measures it,
# mlx_repair.py says whether a self-heal can still overturn it, and main.py holds
# it back while one can. main.py is the studio's front door and is edited often,
# so listing it does cost macOS slots on PRs that never touch the verdict -- but a
# gate that cannot be triggered by the file the logic lives in is not a gate.
- 'studio/backend/main.py'
- 'studio/backend/utils/mlx_repair.py'
# test_hardware_dispatch_matrix.py and test_is_mlx_dispatch_gate.py are
# deliberately absent: this workflow no longer runs them (see the AST
# contract step below), so booting a macOS runner when only those files
# change would burn a slot from the 5-job account-wide macOS cap to run
# nothing related. studio-backend-ci.yml covers both on ubuntu-latest and
# already triggers on them.
- 'studio/backend/tests/test_mlx_inference_backend.py'
- 'tests/studio/test_mlx_training_worker_behaviors.py'
- 'tests/studio/run_real_mlx_smoke.py'
- 'tests/studio/mac_capability_verdict_smoke.py'
- 'tests/conftest.py'
# Every capability-verdict boot below shells out to these two, so an edit to
# either changes what this workflow actually runs.
- '.github/scripts/boot-studio-api-only.sh'
- '.github/scripts/wait-for-health.sh'
- '.github/workflows/mlx-ci.yml'
push:
branches: [main]
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'
- 'studio/backend/main.py'
- 'studio/backend/utils/mlx_repair.py'
- 'studio/backend/tests/test_mlx_inference_backend.py'
- 'tests/studio/test_mlx_training_worker_behaviors.py'
- 'tests/studio/run_real_mlx_smoke.py'
- 'tests/studio/mac_capability_verdict_smoke.py'
- 'tests/conftest.py'
- '.github/scripts/boot-studio-api-only.sh'
- '.github/scripts/wait-for-health.sh'
- '.github/workflows/mlx-ci.yml'
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ github.ref == 'refs/heads/main' && github.sha || '' }}
# Latest-only on a PR branch. On main this does less than it reads like: it stops
# a RUNNING main job being killed, but GitHub cancels any PENDING run in the group
# the moment a newer one is queued, so a merge burst still leaves only the tip.
# See studio-backend-ci.yml, which is grouped per commit on main for that reason.
cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
permissions:
contents: read
jobs:
dispatch:
name: dispatch
# macos-26, not macos-15, on a measurement rather than a preference. Both are
# free Apple Silicon standard runners, so this is the same class of machine.
#
# studio-mac-install-matrix runs both images from ONE matrix, so their queues
# can be compared on identical commits and identical triggers. Over 163 jobs
# per leg:
#
# exec med exec p90 queue p90
# macos-15 160s 713s 20667s
# macos-26 176s 597s 3866s
#
# The medians say the two are the same machine. The queue p90 says they are
# not the same pool: macos-15 is 5.3x worse in the tail, five and a half hours
# against one. That tail is what sets this repo's wall clock -- the census
# behind it found every commit's last finisher to be this job, minutes of work
# behind hours of waiting -- and the median hides it completely, because the
# median wait on both images is zero.
#
# The likely cause is the macos-14 retirement (brownouts 2026-10-05, removal
# 2026-11-02): everything migrated onto macos-15, including this job, and
# macos-26 was left comparatively empty. That also means this is a fact about
# today's pool and not a property of the image, so it is worth re-measuring
# rather than assuming. The comparison above is one query against the install
# matrix and can be re-run whenever the queue looks wrong again.
#
# Not macos-14 for the retirement above; all three are Apple Silicon, so the
# arm64 paths are still what gets tested.
runs-on: macos-26
# 40 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. The capability-verdict
# boots add roughly ten more: two settle in seconds, but the in-flight
# one has to wait out the whole torch warm before the self-heal is even
# scheduled, and then the length of the install it is holding a verdict
# against.
timeout-minutes: 40
steps:
# harden-runner v2.20.0 supports blocking mode on GitHub-hosted macOS
# runners. Keep this job in audit while its egress destinations are
# observed, then graduate it to a reviewed block-mode allowlist.
- name: Harden runner (audit)
uses: step-security/harden-runner@05e31511f85b41b11d1cf0ef85d0992719546e2c # v2.21.0
with:
egress-policy: audit
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
persist-credentials: false
- uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
with:
python-version: '3.12'
- name: Restore the pip cache
id: pip-cache
uses: ./.github/actions/pip-cache-restore
with:
key-files: |
pyproject.toml
studio/backend/requirements/*.txt
# 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. Unsloth'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 --extra-index-url https://pypi.org/simple \
'torch==2.10.0'
# github.com occasionally 500s on the git fetch; retry the
# zoo install so a single upstream blip does not fail CI.
for attempt in 1 2 3; do
if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then
break
fi
if [ "$attempt" -eq 3 ]; then
echo "::error::pip install unsloth_zoo failed after 3 attempts"
exit 1
fi
delay=$((5 * attempt))
echo "::warning::unsloth_zoo install failed (attempt $attempt/3), retrying in ${delay}s..."
sleep "$delay"
done
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')
"
# AST contract checks on worker.py. Pure `ast.parse`, no runtime and no
# mlx import, so it would run anywhere -- but this is the only workflow
# that runs it at all, so it stays here until something on Linux claims it.
#
# test_hardware_dispatch_matrix.py and test_is_mlx_dispatch_gate.py used
# to run here too. They are 100% spoofed monkeypatches (the fixture
# installs a MetaPathFinder that blocks `import mlx.core` for "no-mlx"
# profiles), so they assert the same thing on any host -- and
# studio-backend-ci.yml already runs both on ubuntu-latest in its
# "Hardware-spoof tests" step. Running them again on a macOS runner spent
# minutes from a pool capped at 5 concurrent jobs account-wide to
# re-derive a result Linux had already produced.
- name: MLX worker AST contract tests
env:
PYTHONPATH: ${{ github.workspace }}/studio
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python -m pytest -v --tb=short \
tests/studio/test_mlx_training_worker_behaviors.py
# The inference backend's numerical tests importorskip mlx, so they skip on
# studio-backend-ci.yml's ubuntu runner and only ever execute here. Run from
# studio/backend: its conftest puts that directory on sys.path.
- name: MLX inference backend tests
working-directory: studio/backend
env:
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python -m pytest -v --tb=short tests/test_mlx_inference_backend.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:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
mkdir -p mlx_workdir
# Authenticate llama.cpp's release-API lookup (anonymous 403s on rate-limit);
# read-only GITHUB_TOKEN scoped here only, never to steps that run binaries.
GH_TOKEN="${{ secrets.GITHUB_TOKEN }}" GITHUB_TOKEN="${{ secrets.GITHUB_TOKEN }}" \
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:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && 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:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && 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:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && 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
# ---------------------------------------------------------------
# The Mac capability verdict, on a real Apple Silicon host.
#
# /api/health is what decides whether the sidebar's Train and Video rows
# are enabled, greyed out behind "Training needs MLX. Run `unsloth studio
# update`", or still spinning. That decision reads a real `import
# mlx.core` and waits on a real background reinstall, and the unit tests
# for it necessarily fake both. Here the host is Apple Silicon and the
# MLX stack is genuinely present (first step) or genuinely unimportable
# (last two), so this is the only place the verdict is judged against the
# thing it is a verdict about.
#
# Three boots, one per step: the self-heal is once per process, so the
# opted-out and the in-flight cases cannot share a server.
# ---------------------------------------------------------------
- name: Provision a Studio venv for the capability-verdict boots
run: |
set -euo pipefail
# `unsloth studio` re-execs into $STUDIO_HOME/unsloth_studio and refuses to
# start without it ("Unsloth Studio not set up. Run install.sh first."), so a
# bare `pip install -e .` cannot boot a server. install.sh --local builds that
# venv, but it also reinstalls the whole stack for ten minutes, and
# studio-mac-ui-smoke.yml already pays that on every studio/** PR. Building the
# venv over this job's site-packages gets the same launch path -- no re-exec,
# the real run.py, the real uvicorn -- for the cost of a virtualenv.
STUDIO_VENV="$HOME/.unsloth/studio/unsloth_studio"
rm -rf "$STUDIO_VENV"
python -m venv --system-site-packages "$STUDIO_VENV"
"$STUDIO_VENV/bin/pip" install --no-deps -e .
"$STUDIO_VENV/bin/unsloth" --help > /dev/null
echo "studio venv ready at $STUDIO_VENV"
# PATH is set per step rather than through $GITHUB_PATH: the llama.cpp steps
# below deliberately run under the job's own interpreter, and a venv exported
# for the rest of the job would quietly move them.
- name: "Mac capability verdict: real MLX must never look chat-only"
env:
UNSLOTH_COMPILE_DISABLE: '1'
run: |
export PATH="$HOME/.unsloth/studio/unsloth_studio/bin:$PATH"
python tests/studio/mac_capability_verdict_smoke.py real-mlx \
--port 18901 \
--log logs/studio_verdict_real_mlx.log \
--workdir logs/verdict_real_mlx
- name: "Mac capability verdict: no MLX, self-heal opted out, must still settle"
env:
UNSLOTH_COMPILE_DISABLE: '1'
run: |
export PATH="$HOME/.unsloth/studio/unsloth_studio/bin:$PATH"
python tests/studio/mac_capability_verdict_smoke.py no-mlx-settles \
--port 18902 \
--log logs/studio_verdict_no_mlx.log \
--workdir logs/verdict_no_mlx
- name: "Mac capability verdict: no MLX, self-heal in flight, must not publish"
env:
UNSLOTH_COMPILE_DISABLE: '1'
run: |
export PATH="$HOME/.unsloth/studio/unsloth_studio/bin:$PATH"
python tests/studio/mac_capability_verdict_smoke.py no-mlx-repair \
--port 18903 \
--log logs/studio_verdict_repairing.log \
--workdir logs/verdict_repairing
- name: "Mac capability verdict: server logs"
if: always()
run: |
for f in logs/studio_verdict_real_mlx.log \
logs/studio_verdict_no_mlx.log \
logs/studio_verdict_repairing.log; do
echo "=== $f ==="
tail -60 "$f" 2>/dev/null || echo "(missing)"
echo
done
# Validates the macOS prebuilt path Unsloth's setup.sh uses (#5963): install the
# unslothai/llama.cpp fork's latest release, download a small public GGUF, and
# check llama-server /completion end to end. Split and placed last so the
# untrusted binary runs only in the final smoke step, after every HF_TOKEN step,
# leaving no token-bearing step or shared workspace for a tampered prebuilt to
# corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch.
- name: Unsloth prebuilt llama.cpp install + GGUF download (Mac M1)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# Download only -- no llama-quantize / llama-server launch in this step.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo unslothai/llama.cpp
mkdir -p /tmp/ggufs
bash .github/scripts/hf-download-with-retry.sh \
'unsloth/gemma-3-270m-it-GGUF' \
'gemma-3-270m-it-Q4_K_M.gguf' \
/tmp/ggufs
# Final step: runs the downloaded binaries with no secrets present, and clears
# the GitHub Actions command files so a tampered prebuilt cannot influence the job.
- name: Unsloth prebuilt llama.cpp GGUF inference smoke (Mac M1)
run: |
set -euo pipefail
unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
# Unsloth bundles only llama-server + llama-quantize (not llama-cli);
# inference goes through llama-server's HTTP /completion endpoint.
LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
[ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; }
[ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; }
echo "llama-server : $LLAMA_SERVER"
echo "llama-quantize: $LLAMA_QUANT"
"$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK"
PORT=18080
echo "=== starting llama-server on 127.0.0.1:$PORT ==="
"$LLAMA_SERVER" \
-m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \
--host 127.0.0.1 \
--port "$PORT" \
-c 256 \
-n 16 \
--no-warmup \
> /tmp/llama-server.log 2>&1 &
SERVER_PID=$!
trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT
# Wait for /health to come up
for i in $(seq 1 30); do
if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo " server up after ${i}s"
break
fi
sleep 1
done
if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo "::error::llama-server never became healthy"
tail -40 /tmp/llama-server.log
exit 1
fi
PROMPT="Hello, my name is"
echo "=== POST /completion ==="
RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \
-H 'Content-Type: application/json' \
-d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}")
echo "raw response (head): $(echo "$RESP" | head -c 600)"
CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))")
echo "completion content: $CONTENT"
if [ -z "$CONTENT" ]; then
echo "::error::llama-server /completion returned empty content"
tail -40 /tmp/llama-server.log
exit 1
fi
echo "OK: Unsloth prebuilt llama.cpp on Mac M1 + GGUF /completion works"
- name: Save the pip cache
if: always()
uses: ./.github/actions/pip-cache-save
with:
dir: ${{ steps.pip-cache.outputs.dir }}
key: ${{ steps.pip-cache.outputs.key }}
cache-hit: ${{ steps.pip-cache.outputs.cache-hit }}
Reference in a new issue View git blame Copy permalink