mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-30 01:54:02 +00:00
3 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
1dd2fc4583
|
tests: read checked-in files as UTF-8 instead of the platform default (#7438)
* tests: read checked-in files as UTF-8 instead of the platform default Path.read_text() with no encoding uses locale.getpreferredencoding(), which is UTF-8 on the Linux runners and cp1252 on a stock Windows install. Nine module-level reads of checked-in source files were relying on that default. studio/backend/routes/inference.py carries the DeepSeek tool-call token regexes, so it holds U+FF5C and U+2581. Under cp1252 that read raised UnicodeDecodeError on byte 0x81 at position 97806, and because the reads run at import time it took test_cancel_atomicity.py and test_cancel_id_wiring.py out at collection, not as failures. Green on CI, permanently broken for a Windows contributor running the suite locally. Adds a guard: at module scope there is no tmp_path fixture, so a bare read_text()/write_text()/open() there is always touching a checked-in file. That makes the rule mechanical enough to enforce with no allowlist, while staying quiet about temp-dir I/O inside test bodies where the platform default is harmless. The repo already spells this correctly in 464 other places; this only stops the stragglers coming back. * tests: cover import-time helper reads and keep the guard py3.9-safe Follows up on the Codex review: - add `from __future__ import annotations`, since `str | None` in `_offender` is evaluated at import on Python 3.9 and pyproject declares requires-python ">=3.9,<3.15". - widen the guard from module scope to import time. Class bodies and the bodies of module-level helpers called from an executing statement run during collection too, so `CODE = _extract_mixed_precision_code()` was the same hazard as an inline read. `if __name__ == "__main__":` blocks are skipped: pytest never executes them. - scan studio/backend/tests/ as well as tests/. Both trees are collected on Windows by separate CI jobs, and the offender that started this, test_tool_xml_strip.py reading routes/inference.py, lives there. Widening it surfaced seven more import-time reads of checked-in sources; all now name utf-8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden the import-time encoding guard for PR #7438 Close the detector gaps raised in review, all of which I reproduced against the actual AST before changing anything. False negatives (the guard let a real hazard through): - _is_main_guard ignored the comparison operator, so if __name__ != "__main__" counted as script-only even though its body runs at import. - The else arm of a main guard was discarded with the rest of the If node. - Decorators and argument defaults on a module-level def were skipped with the body, though both are evaluated when the def executes. - Path.open() in text mode was invisible; only builtin open() was matched. - encoding = None and encoding = "locale" both re-select the platform default, but the keyword merely being present counted as pinned. False positives (the guard would have blocked a compliant contributor): - A non-literal mode fell through to the "r" default, so open(p, mode) was flagged even when mode is "rb", where adding encoding= is a ValueError and there is no edit that satisfies the rule. - Same for open(*args) and a **kwargs splat, which hide the mode and can hide an encoding. - Lambda bodies and comprehension elements were walked even though neither runs at definition. Verified: still reports the same 22 offenders on unpatched main, green on this branch and on the tree merged with latest main (557 files), and an adversarial corpus of 33 cases now scores zero false positives and zero false negatives. Also corrected two docstring claims: neither collecting job runs on Windows, and the read is governed by locale.getencoding(). * Walk eager comprehensions and treat io.open as the builtin Two regressions from the previous commit, both reproduced against the AST before changing anything. Lumping list, set and dict comprehensions in with generator expressions was wrong. Only a genexp is lazy; the other three run their element expression, their filters and their nested iterators immediately, so CONTENTS = [p.read_text() for p in PATHS] at module scope is an import-time read the guard was silently missing. Comprehensions are now walked in full and only the genexp keeps the outermost-iterable-only treatment. io was also in the not-a-path-opener list, but io.open is the builtin, with the same mode position and the same platform default. io.open(CHECKED_IN_FILE) is exactly the hazard this guard exists for, so it is matched now, with binary modes and a pinned encoding still exempt. tarfile.open and fitz.open stay exempt since neither has an encoding to name. Verified: 13 targeted cases covering all five eager comprehension forms and io.open in text, binary and pinned shapes all classify correctly; still 22 offenders on unpatched main; green on this branch and on the tree merged with latest main. * Close three more walker gaps in the import-time guard All three reproduced against the AST first. A generator expression handed straight to a call is consumed there, so DATA = "".join(p.read_text() for p in paths) runs its element at import. Only an unconsumed genexp bound to a name stays lazy, so the walker now follows the consumed ones in full and keeps the outermost-iterable-only treatment for the rest. if "__main__" == __name__ is an equivalent and accepted spelling of the main guard, but requiring __name__ on the left meant its body was treated as import-time code. That is a false positive on a block pytest never runs, so both operand orders are recognised now. The helper table was built from module-level defs only, so a def in a class body invoked while the class is constructed was never followed, contradicting the walker's stated coverage of class bodies. Helpers are now collected from the module body and from class bodies at any nesting. Verified: 15 targeted cases including all three fixes and the earlier ones still classify correctly; still 22 offenders on unpatched main; green on this branch and on the tree merged with latest main. * Handle positional read_text encodings, lazy generators and nested helpers * Guard reads reached from test bodies, unbound Path calls and __file__ paths * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Follow derived paths, skip lazy generator helpers, cover compressed openers * Guard the CLI tests, helper parameters and unbound Path arguments * Discover test roots and follow literal, in-place and tuple-derived paths * Identify module openers by import, unwrap starred paths, pin subprocess snippets * Resolve import origins, seed helper locals, follow named generators and parametrize * Scope imports lexically, list tracked test files, bind unpacked names * Resolve aliased openers, keyword-only params, destructured targets, next() * Pin the encoding on subprocess snippets, workflow lint and CLI output for PR #7438 * Harden the CLI encoding guard against detached streams for PR #7438 * Tighten the encoding guard's path and scope analysis for PR #7438 * Resolve path provenance more precisely and keep POSIX stream encodings for PR #7438 * Resolve qualified path classes and scope conditional imports for PR #7438 * Scope CLI stream setup to the entry point and align two encoding pairs for PR #7438 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <danielhanchen@gmail.com> |
||
|
|
3ce187da02
|
Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent. |
||
|
|
d65149795b
|
feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export) (#5265)
* Add Apple Silicon MLX routing
Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports
Extract original GPU init to _gpu_init.py (unchanged)
MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code
GPU path unchanged: from ._gpu_init import *
* Add Apple Silicon MLX routing
- Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports
- Extract original GPU init to _gpu_init.py (unchanged)
- MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code
- GPU path unchanged: from ._gpu_init import *
* mlx with studio
* mlx with studio
* updating temporary install.sh
* updating temporary install.sh
* adding t_v5 path
* adding t_v5 path
* fixing vision training
* fixing vision training
* adding chat
* adding chat
* minor
* minor
* Adding export and fixing training issues, inference with lora adaptors
* Adding export and fixing training issues, inference with lora adaptors
* fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM
* fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM
* Merge mlx-apple-silicon into main
* update install.sh to point to main branch
* update install.sh to point to main branch
* fix: export returns 3 values (success, message, output_path) matching upstream worker
* fix: export returns 3 values (success, message, output_path) matching upstream worker
* fix(mlx): show training-process peak memory in Studio UI, not system-wide
Studio UI was showing ~95 GB during MLX training because get_gpu_utilization
read "In use system memory" from IORegistry's AGXAccelerator — system-wide
GPU memory across all processes (training + backend + browser + Display).
Now the trainer's mx.get_peak_memory value is forwarded through the
progress event and surfaced via /api/train/hardware while training is
active. Falls back to the system-wide reading when training is not running.
* fix(mlx): show training-process peak memory in Studio UI, not system-wide
Studio UI was showing ~95 GB during MLX training because get_gpu_utilization
read "In use system memory" from IORegistry's AGXAccelerator — system-wide
GPU memory across all processes (training + backend + browser + Display).
Now the trainer's mx.get_peak_memory() value is forwarded through the
progress event and surfaced via /api/train/hardware while training is
active. Falls back to the system-wide reading when training is not running.
* fix(mlx): make is_bfloat16_supported detect M1/M2 (no native bf16)
M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70%
slower prefill compared to native fp16. M3+ have native bf16 (macOS
Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware
detection via mx.device_info.
* fix(mlx): make is_bfloat16_supported() detect M1/M2 (no native bf16)
M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70%
slower prefill compared to native fp16. M3+ have native bf16 (macOS
Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware
detection via mx.device_info().
* feat(mlx): wire training_type="Full Finetuning" through MLX worker
Compute use_lora from the UI's training_type before loading the model,
pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and
let the existing 'if use_lora' branch skip get_peft_model. Matches the
GPU worker's flow.
* feat(mlx): wire training_type="Full Finetuning" through MLX worker
Compute use_lora from the UI's training_type before loading the model,
pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and
let the existing 'if use_lora' branch skip get_peft_model. Matches the
GPU worker's flow.
* fix(mlx): pass save_method='merged_16bit' from Studio's export page
Previously the MLX path called save_pretrained_merged with no
save_method, which fell through to a no-op that didn't actually fuse
LoRA into the base. Now Studio's "Merged Model" export properly
fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU
behavior for the same UI option.
* fix(mlx): pass save_method='merged_16bit' from Studio's export page
Previously the MLX path called save_pretrained_merged() with no
save_method, which fell through to a no-op that didn't actually fuse
LoRA into the base. Now Studio's "Merged Model" export properly
fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU
behavior for the same UI option.
* fix(studio): pass private to MLX push, return 3-tuples consistently
MLX push_to_hub branch now forwards private=private (matches GPU)
Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model
needed') were tripping the route's 3-tuple unpack. Added a None
output_path so the unpack always succeeds.
* fix(studio): pass private to MLX push, return 3-tuples consistently
- MLX push_to_hub branch now forwards private=private (matches GPU)
- Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model
needed') were tripping the route's 3-tuple unpack. Added a None
output_path so the unpack always succeeds.
* studio wirings
* studio wirings
* Merge pull request #5 from Manan17/feat/quant_config
studio wirings
* fix(mlx): wire train_on_completions for VLM via per-template lookup
Mirror the GPU worker: stop excluding VLMs and stop hardcoding
template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and
fetch the per-template instruction/response markers from
TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables
train_on_completions for vision+image and audio cases, so backend
just trusts the flag.
* fix(mlx): wire train_on_completions for VLM via per-template lookup
Mirror the GPU worker: stop excluding VLMs and stop hardcoding
template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and
fetch the per-template instruction/response markers from
TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables
train_on_completions for vision+image and audio cases, so backend
just trusts the flag.
* wire in lora rslora, init lora weights, random_state
* wire in lora rslora, init lora weights, random_state
* loftq studio error message fix
* loftq studio error message fix
* handle unknown optim and lr scheduler
* handle unknown optim and lr scheduler
* Merge pull request #6 from Manan17/update/peftkwargs
Update/peftkwargs
* feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel
Studio's four UI checkboxes now actually flow through to MLX get_peft_model
(which was just updated in unsloth-zoo to honor them). Also drops the
incorrect train_projector wiring that tied projector LoRA to the
attn/mlp flags — those are language-side toggles, not projector toggles.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel
Studio's four UI checkboxes now actually flow through to MLX get_peft_model
(which was just updated in unsloth-zoo to honor them). Also drops the
incorrect train_projector wiring that tied projector LoRA to the
attn/mlp flags — those are language-side toggles, not projector toggles.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp
UI guardrail. The four checkboxes (vision/language/attention/MLP) carry
"scope × module-type" semantics that aren't obvious — picking just
"Attention modules" + "MLP modules" without "Language layers" naturally
reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune
attn/mlp modules in *no* tower" → empty target_modules → zero
trainable params → crash inside value_and_grad.
If user selected attn or mlp module types but no layer scope, default
to language scope. Power users can still explicitly choose
language=False, vision=True if they want vision-only fine-tuning of
attn/mlp.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp
UI guardrail. The four checkboxes (vision/language/attention/MLP) carry
"scope × module-type" semantics that aren't obvious — picking just
"Attention modules" + "MLP modules" without "Language layers" naturally
reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune
attn/mlp modules in *no* tower" → empty target_modules → zero
trainable params → crash inside value_and_grad.
If user selected attn or mlp module types but no layer scope, default
to language scope. Power users can still explicitly choose
language=False, vision=True if they want vision-only fine-tuning of
attn/mlp.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm
Inference UI sliders for top_k and repetition_penalty had no effect on
MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing
bug: mlx_vlm.generate_step expects temperature= (long form), but we
were passing temp= which silently fell into **kwargs — every VLM chat
was effectively greedy regardless of the temperature slider.
Text path (_generate_text):
make_sampler now receives top_k in addition to temp/top_p
make_logits_processors built and forwarded when repetition_penalty is
non-trivial (skip when 0.0/1.0 to avoid pointless overhead)
VLM path (_generate_vlm):
Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate
forwards them to generate_step's sampler/logits_processor builders)
Rename temp= → temperature= so it's actually consumed
Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and
Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10,
rep_pen=1.5} now produces a distinct output, proving the parameters
reach the sampler.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm
Inference UI sliders for top_k and repetition_penalty had no effect on
MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing
bug: mlx_vlm.generate_step expects temperature= (long form), but we
were passing temp= which silently fell into **kwargs — every VLM chat
was effectively greedy regardless of the temperature slider.
Text path (_generate_text):
- make_sampler now receives top_k in addition to temp/top_p
- make_logits_processors built and forwarded when repetition_penalty is
non-trivial (skip when 0.0/1.0 to avoid pointless overhead)
VLM path (_generate_vlm):
- Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate
forwards them to generate_step's sampler/logits_processor builders)
- Rename temp= → temperature= so it's actually consumed
Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and
Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10,
rep_pen=1.5} now produces a distinct output, proving the parameters
reach the sampler.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push
export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit"
(was hardcoded merged_16bit, ignoring the UI choice).
Both export_merged_model and export_base_model now pass save_directory=
to push_to_hub_merged so it reuses the just-written local folder
instead of re-saving under a relative "username/model" directory.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push
- export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit"
(was hardcoded merged_16bit, ignoring the UI choice).
- Both export_merged_model and export_base_model now pass save_directory=
to push_to_hub_merged so it reuses the just-written local folder
instead of re-saving under a relative "username/model" directory.
Co-Authored-By: Manan17 <shahmanan170602@gmail.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* restore install
* restore install
* fix(mlx): restore FastVisionModel as a distinct class
unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel`
right after defining `class FastVisionModel(FastLanguageModel)` with a
`for_training` static method. The alias erased the class binding, so
the documented `FastVisionModel.for_training(model)` call from upstream
Unsloth's VLM notebooks raised `AttributeError` on MLX.
Remove the offending alias. `FastVisionModel` is now a real subclass of
`FastLanguageModel` again — inherits `from_pretrained` /
`get_peft_model` / `for_inference`, exposes `for_training` as a no-op
pass-through (no-op because MLX doesn't have a train/eval mode flag;
the call exists purely for GPU/MLX notebook parity).
Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via
FastVisionModel.from_pretrained → get_peft_model → for_training →
MLXTrainer.train runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs,
peak 5.89 GB).
Studio's path (FastLanguageModel.from_pretrained for any repo,
auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8.
* fix(mlx): restore FastVisionModel as a distinct class
unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel`
right after defining `class FastVisionModel(FastLanguageModel)` with a
`for_training` static method. The alias erased the class binding, so
the documented `FastVisionModel.for_training(model)` call from upstream
Unsloth's VLM notebooks raised `AttributeError` on MLX.
Remove the offending alias. `FastVisionModel` is now a real subclass of
`FastLanguageModel` again — inherits `from_pretrained` /
`get_peft_model` / `for_inference`, exposes `for_training` as a no-op
pass-through (no-op because MLX doesn't have a train/eval mode flag;
the call exists purely for GPU/MLX notebook parity).
Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via
FastVisionModel.from_pretrained → get_peft_model → for_training →
MLXTrainer.train() runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs,
peak 5.89 GB).
Studio's path (FastLanguageModel.from_pretrained for any repo,
auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8.
* Studio: harden MLX training and export, restore GPU init guards
Studio export
Restore Tuple[bool, str, Optional[str]] contract on export_merged_model,
export_base_model, export_gguf, and export_lora_adapter, populating
output_path on successful local saves so routes/worker/CLI/frontend
details.output_path is non-empty again.
Lift the GPU save_method assignment out of the local-save branch so
Hub-only merged exports (save_directory='', push_to_hub=True) no longer
hit UnboundLocalError on the push branch.
For MLX merged and base hub-only export, stage to a tempfile.TemporaryDirectory
before push_to_hub_merged instead of passing save_directory=''.
Source _IS_MLX from unsloth instead of recomputing the platform check
(single source of truth, also enforces mlx-package availability).
Studio MLX training/inference
Pass token=hf_token into FastMLXModel.from_pretrained for gated/private
models, matching the inference path.
Strip hf_token and wandb_token from wandb.init(config=...) so secrets
do not leak into the W&B run config.
Replace load_from_disk(local_datasets[0]) with the existing
UnslothTrainer._resolve_local_files / _loader_for_files helpers so
uploaded JSON/JSONL/CSV/Parquet files train through the normal datasets
loader (load_from_disk still used for HF save_to_disk directories).
Make the dataset slice helper inclusive at the end and treat 0 as a real
index instead of "unset", matching the GPU and embedding paths.
Add a status_message -> message alias inside _send so the existing parent
pump (training.py) renders MLX status updates instead of blanks.
Forward min_p through generate_chat_response into _generate_text /
_generate_vlm and into make_sampler / vlm_kwargs so the sampling control
is no longer a no-op on MLX.
Wrap unsloth_zoo.mlx_loader / mlx_trainer imports with a clearer
ImportError pointing users at install.sh for Apple Silicon.
Exit the MLX stop-polling thread on EOFError/OSError instead of
busy-looping when the queue/pipe is permanently closed (one-line
why-safe rationale inline).
Studio frontend
ParamsSection subscribes to platform deviceType via the Zustand hook so
the gradient checkpointing dropdown re-renders after the async device
fetch completes.
Studio hardware
get_gpu_utilization MLX branch now reads _read_apple_gpu_stats once and
derives VRAM totals from psutil, removing the second ioreg subprocess
per utilization poll.
Unsloth core
Restore the os.geteuid == 0 guard around the CUDA ldconfig recovery
that was lost when GPU initialization moved into _gpu_init.py, plus the
non-root manual-fix warning branch. Non-root CUDA users no longer shell
out to ldconfig at import time.
Load dataprep/raw_text via importlib so the MLX import path no longer
pulls torch in through dataprep/__init__.py -> synthetic.py.
FastVisionModel.from_pretrained overrides the inherited delegator only
to inject text_only=False; this is an extension, not a duplication, and
is needed so VLM checkpoint loads keep the vision tower.
Wrap the MLX-branch unsloth_zoo import with a clearer ImportError.
* Studio: regression tests for MLX training/export and GPU init ldconfig guard
tests/python/test_gpu_init_ldconfig_guard.py asserts the geteuid root
check still wraps the ldconfig recovery and the non-root branch warns
bnb users; AST + source-text inspection so the test runs without torch.
tests/studio/test_export_output_path_contract.py covers the
Tuple[bool, str, Optional[str]] return contract on every export method,
the output_path assignment after successful local save, the Hub-only
GPU save_method binding fix, the MLX hub-only TemporaryDirectory
staging, and the single-source `_IS_MLX` import from unsloth.
tests/studio/test_mlx_training_worker_behaviors.py covers token
forwarding to FastMLXModel.from_pretrained, wandb config secret
stripping, file-aware local dataset loading, status_message ->
message aliasing, inclusive slice semantics, EOFError/OSError stop
thread exit, and the friendly mlx_loader / mlx_trainer ImportError.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(mlx): cap inference memory + release wired on unload + tame worker pre-pin
Three memory-hardening fixes for Studio's MLX path:
1. Inference applies the same Metal caps as the trainer.
load_model previously only called set_wired_limit(100% of recommended)
with no upper memory_limit, leaving large VLM checkpoints unbounded
during the loader allocation. Add _configure_memory_limits() that sets
memory_limit to 85% of recommended and wired_limit to min(recommended,
memory_limit) — matching MLXTrainer's defaults so behavior is the same
whether the user trains or just runs inference.
2. unload_model releases pinned memory back to the OS — but only when
the cache is empty. Without this, pinned wired bytes stayed allocated
to MLX after the model was gone, starving other apps. The release is
guarded on `not self.models` so unloading one of several cached
models doesn't un-pin weights still in use.
3. Worker pre-cap is conservative instead of aggressive.
The previous pre-pin set_wired_limit(100% of recommended) competed
with MLXTrainer's later more conservative cap. Replace with the same
85%-memory / min(rec, memory) pair that the trainer applies later
(idempotent re-apply). Bounds the model load + LoRA setup window
without over-pinning.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* tests/studio: regression tests for the _IS_MLX dispatch gate
Two gates drive every MLX-vs-CUDA dispatch decision in Studio:
1. unsloth._IS_MLX in unsloth/__init__.py — evaluated once at import
time, read by Studio worker code to choose the GPU vs MLX trainer
and inference paths. Defined as
Darwin AND arm64 AND find_spec("mlx") is not None.
2. utils.hardware.detect_hardware() — runtime probe with priority
CUDA > XPU > MLX > CPU. The MLX branch is reached only when both
CUDA and XPU are unavailable and the host is Apple Silicon and
mlx is importable.
Neither gate had a direct test. Adds tests/studio/test_is_mlx_dispatch_gate.py
with six tests:
test_is_mlx_gate_uses_three_required_predicates
AST-walks unsloth/__init__.py and asserts the _IS_MLX assignment
is a BoolOp(And) of platform.system()=="Darwin",
platform.machine()=="arm64", and find_spec("mlx") is not None.
Catches accidental rewrites that drop a predicate.
test_is_mlx_gate_true_on_apple_silicon_with_mlx_present
Spoofs platform to Darwin/arm64, injects a fake mlx module so
find_spec returns a real ModuleSpec, re-evaluates the gate
expression. Verifies it flips True under the exact conditions
Studio expects.
test_is_mlx_gate_false_when_mlx_missing
Spoofs Apple Silicon but with mlx absent. Verifies the gate stays
False (so a Mac without mlx installed does not pretend to have
MLX support).
test_is_mlx_gate_false_on_non_apple_silicon
Canary on the actual Linux+CUDA / AMD / Intel test host: the gate
must remain False regardless of whether mlx happens to be
importable. Protects existing GPU users from accidental MLX
hijack when MLX support evolves.
test_detect_hardware_picks_mlx_when_only_apple_silicon_available
Forces torch.cuda and torch.xpu off, spoofs Apple Silicon, injects
fake mlx and mlx.core. detect_hardware() must return DeviceType.MLX.
test_detect_hardware_picks_cuda_on_real_host
Canary: on a real CUDA host detect_hardware() must return
DeviceType.CUDA. Protects against the MLX branch shadowing CUDA
dispatch on NVIDIA / AMD ROCm hosts.
Uses the same monkeypatch.setitem(sys.modules, ...) fake-mlx pattern as
the existing test_mlx_inference_backend.py — no new test infrastructure,
no real mlx install required.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add AGPL-3.0 SPDX header to Studio MLX regression tests
Four Studio MLX test files shipped without an SPDX-License-Identifier:
studio/backend/tests/test_mlx_training_worker_config.py
tests/studio/test_mlx_training_worker_behaviors.py
tests/studio/test_export_output_path_contract.py
tests/studio/test_is_mlx_dispatch_gate.py
They sit in or alongside studio/backend/, which is governed by
studio/LICENSE.AGPL-3.0, and exercise AGPL Studio code. Add the same
"# SPDX-License-Identifier: AGPL-3.0-only" header that's already on
test_mlx_inference_backend.py so the license declaration matches
the code under test rather than defaulting to the repo-root
Apache-2.0.
* Wrap MLX submodule imports with friendly install hint
The _IS_MLX block at the top of unsloth/__init__.py already catches the
missing-package case with a friendly install hint, but the follow-up
"from unsloth_zoo.mlx_trainer import ..." and "from unsloth_zoo.mlx_loader import ..."
lines run unguarded. An Apple Silicon user who has unsloth-zoo installed
but on an older version (e.g. the current PyPI release, before the MLX
modules ship) sees a raw ImportError on the submodule rather than the
hint that points at install.sh.
Wrap the two submodule imports in the same try/except shape so the
friendly install message fires whether the package is missing entirely
or just predates the MLX submodules. No-op once both packages release
together; smooths the transitional window where unsloth/main has merged
but unsloth-zoo on PyPI has not.
---------
Co-authored-by: DoubleMathew <mmathew23@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
|