From a932294627176b742ff425817f320c60c28d284d Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Thu, 14 May 2026 07:24:20 -0500 Subject: [PATCH 01/10] MLX training support for Studio on Apple Silicon (#5340) * mlx fixes * Fix studio integration, local dataset files, chat templates without the torch gpu imports * pass grad norm in mlx worker * fix(studio): pass MLX grad clipping settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * mlx: update grad value * fix(mlx): address ci and clipping review * fix backward compatibility and CI tests * unsloth local is mlx function * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * dont reference runtime * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio mlx: hardcode value clipping, drop max_grad_value from frontend Simplifies the MLX grad-clipping plumbing now that we are standardising on elementwise value clipping at [-5, 5] for the compiled MLX path and norm clipping disabled. The MLX worker no longer reads max_grad_norm / max_grad_value from the request; both are pinned in one place. Frontend stops sending the field at all, and the TypeScript request type drops it to match. Non-MLX (CUDA/AMD/Intel) is untouched and continues to pick up HF TrainingArguments' default max_grad_norm = 1.0. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../backend/core/inference/mlx_inference.py | 4 +- studio/backend/core/training/training.py | 2 + studio/backend/core/training/worker.py | 87 +++++++++++++++++-- studio/backend/models/training.py | 13 +++ studio/backend/routes/training.py | 2 + .../tests/test_mlx_inference_backend.py | 9 +- .../tests/test_training_raw_support.py | 44 ++++++++++ .../backend/utils/datasets/chat_templates.py | 25 +++++- .../src/features/training/api/mappers.ts | 1 + .../src/features/training/types/api.ts | 1 + tests/studio/run_real_mlx_smoke.py | 6 +- tests/studio/test_is_mlx_dispatch_gate.py | 66 +++++++++----- .../test_mlx_training_worker_behaviors.py | 8 +- unsloth/__init__.py | 37 +++++--- unsloth/chat_templates.py | 14 ++- unsloth/device_type.py | 23 ++++- unsloth/kernels/utils.py | 9 ++ 17 files changed, 295 insertions(+), 56 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index ee7c9bbd51..e7bce2d33e 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -116,11 +116,11 @@ class MLXInferenceBackend: ) try: - from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as e: raise ImportError( "Unsloth: MLX inference requires unsloth-zoo with the MLX modules " - "(unsloth_zoo.mlx_loader). Reinstall via install.sh on Apple Silicon." + "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e model, tokenizer_or_processor = FastMLXModel.from_pretrained( diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 549d733252..556545e680 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -214,6 +214,8 @@ class TrainingBackend: "max_steps": kwargs.get("max_steps", 0), "save_steps": kwargs.get("save_steps", 0), "weight_decay": kwargs.get("weight_decay", 0.001), + "max_grad_norm": kwargs.get("max_grad_norm", 0.0), + "max_grad_value": kwargs.get("max_grad_value"), "random_seed": kwargs.get("random_seed", 3407), "packing": kwargs.get("packing", False), "optim": kwargs.get("optim", "adamw_8bit"), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ef5cafb175..8cd1d1f9c8 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -417,6 +417,55 @@ def _normalize_mlx_studio_scheduler(value): return raw +def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: + """Resolve Studio local dataset uploads without importing the GPU trainer.""" + from utils.paths import resolve_dataset_path + + all_files: list[str] = [] + for dataset_file in file_paths or []: + file_path = ( + dataset_file + if os.path.isabs(dataset_file) + else str(resolve_dataset_path(dataset_file)) + ) + file_path_obj = Path(file_path) + + if file_path_obj.is_dir(): + parquet_dir = ( + file_path_obj / "parquet-files" + if (file_path_obj / "parquet-files").exists() + else file_path_obj + ) + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if parquet_files: + all_files.extend(str(p) for p in parquet_files) + continue + + candidates: list[Path] = [] + for ext in (".json", ".jsonl", ".csv", ".parquet"): + candidates.extend(sorted(file_path_obj.glob(f"*{ext}"))) + if candidates: + all_files.extend(str(c) for c in candidates) + continue + + raise ValueError(f"No supported data files in directory: {file_path_obj}") + + all_files.append(str(file_path_obj)) + + return all_files + + +def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: + first_ext = Path(files[0]).suffix.lower() + if first_ext in (".json", ".jsonl"): + return "json" + if first_ext == ".csv": + return "csv" + if first_ext == ".parquet": + return "parquet" + raise ValueError(f"Unsupported dataset format: {files[0]}") + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -442,8 +491,8 @@ def _run_mlx_training(event_queue, stop_queue, config): import mlx.core as mx try: - from unsloth_zoo.mlx_loader import FastMLXModel - from unsloth_zoo.mlx_trainer import ( + from unsloth_zoo.mlx.loader import FastMLXModel + from unsloth_zoo.mlx.trainer import ( MLXTrainer, MLXTrainingConfig, train_on_responses_only, @@ -451,7 +500,7 @@ def _run_mlx_training(event_queue, stop_queue, config): except ImportError as e: raise ImportError( "Unsloth: MLX training requires unsloth-zoo with the MLX modules " - "(unsloth_zoo.mlx_loader / unsloth_zoo.mlx_trainer). Reinstall via " + "(unsloth_zoo.mlx.loader / unsloth_zoo.mlx.trainer). Reinstall via " "install.sh on Apple Silicon." ) from e from datasets import load_dataset @@ -572,7 +621,6 @@ def _run_mlx_training(event_queue, stop_queue, config): return ds def _load_local(file_paths): - from core.training.trainer import UnslothTrainer from datasets import load_from_disk if len(file_paths) == 1: @@ -581,10 +629,10 @@ def _run_mlx_training(event_queue, stop_queue, config): (p / "dataset_info.json").exists() or (p / "state.json").exists() ): return load_from_disk(str(p)) - all_files = UnslothTrainer._resolve_local_files(file_paths) + all_files = _resolve_mlx_local_dataset_files(file_paths) if not all_files: raise ValueError("No local dataset files found") - loader = UnslothTrainer._loader_for_files(all_files) + loader = _mlx_local_dataset_loader_for_files(all_files) return load_dataset(loader, data_files = all_files, split = "train") if hf_dataset: @@ -718,6 +766,10 @@ def _run_mlx_training(event_queue, stop_queue, config): else: eval_steps_val = int(eval_steps_val) + # MLX: value-clip grads to [-5, 5]; norm clipping disabled for compile-friendliness. + max_grad_norm = 0.0 + max_grad_value = 5.0 # TODO: expose MLX grad-clip in Studio UI for power users + trainer = MLXTrainer( model = model, tokenizer = tokenizer, @@ -732,6 +784,8 @@ def _run_mlx_training(event_queue, stop_queue, config): lr_scheduler_type = lr_scheduler_type, optim = optim_name, weight_decay = float(config.get("weight_decay", 0.001) or 0.001), + max_grad_norm = max_grad_norm, + max_grad_value = max_grad_value, logging_steps = 1, max_seq_length = max_seq_length, seed = config.get("random_seed", 3407), @@ -820,7 +874,17 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 9. Real-time progress callback ── _send("status", status_message = f"Training {model_name}...") - def _on_step(step, total, loss, lr, tok_s, peak_gb, elapsed, num_tokens): + def _on_step( + step, + total, + loss, + lr, + tok_s, + peak_gb, + elapsed, + num_tokens, + grad_norm = None, + ): eta = (elapsed / step * (total - step)) if step > 0 else 0 _send( "progress", @@ -831,7 +895,7 @@ def _run_mlx_training(event_queue, stop_queue, config): total_steps = total, elapsed_seconds = elapsed, eta_seconds = max(0, eta), - grad_norm = None, + grad_norm = grad_norm, num_tokens = num_tokens, eval_loss = None, status_message = None, @@ -846,6 +910,11 @@ def _run_mlx_training(event_queue, stop_queue, config): "train/tokens_per_sec": tok_s, "train/peak_gb": peak_gb, "train/num_tokens": num_tokens, + **( + {"train/grad_norm": grad_norm} + if grad_norm is not None + else {} + ), }, step = step, ) @@ -857,6 +926,8 @@ def _run_mlx_training(event_queue, stop_queue, config): tb_writer.add_scalar("train/learning_rate", lr, step) tb_writer.add_scalar("train/tokens_per_sec", tok_s, step) tb_writer.add_scalar("train/peak_gb", peak_gb, step) + if grad_norm is not None: + tb_writer.add_scalar("train/grad_norm", grad_norm, step) except Exception: pass diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 31f1d575d7..34ae9cfe72 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -262,6 +262,19 @@ class TrainingStartRequest(BaseModel): max_steps: Optional[int] = Field(None, description = "Maximum training steps") save_steps: int = Field(100, description = "Steps between checkpoints") weight_decay: float = Field(0.001, description = "Weight decay") + max_grad_norm: float = Field( + 0.0, + ge = 0, + description = "Global gradient norm clipping threshold. Set 0 to disable.", + ) + max_grad_value: Optional[float] = Field( + None, + ge = 0, + description = ( + "Elementwise gradient value clipping threshold. Set 0 to disable. " + "If omitted, MLX defaults to 1 unless max_grad_norm is set." + ), + ) random_seed: int = Field(42, description = "Random seed") packing: bool = Field(False, description = "Enable sequence packing") optim: str = Field("adamw_8bit", description = "Optimizer") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 19202f3883..6875a13206 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -215,6 +215,8 @@ async def start_training( "max_steps": request.max_steps, "save_steps": request.save_steps, "weight_decay": request.weight_decay, + "max_grad_norm": request.max_grad_norm, + "max_grad_value": request.max_grad_value, "random_seed": request.random_seed, "packing": request.packing, "optim": request.optim, diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 868e537372..ce447bdd1f 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -56,11 +56,14 @@ def _install_fake_fast_mlx(monkeypatch, calls): return _DummyModel(), _DummyTokenizer() unsloth_zoo_pkg = types.ModuleType("unsloth_zoo") - mlx_loader = types.ModuleType("unsloth_zoo.mlx_loader") + mlx_pkg = types.ModuleType("unsloth_zoo.mlx") + mlx_loader = types.ModuleType("unsloth_zoo.mlx.loader") mlx_loader.FastMLXModel = _FastMLXModel - unsloth_zoo_pkg.mlx_loader = mlx_loader + unsloth_zoo_pkg.mlx = mlx_pkg + mlx_pkg.loader = mlx_loader monkeypatch.setitem(sys.modules, "unsloth_zoo", unsloth_zoo_pkg) - monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx_loader", mlx_loader) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader) def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index 876ee34686..c4aaaf3298 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/studio/backend/tests/test_training_raw_support.py @@ -70,6 +70,50 @@ class TestTrainingRawSupport(unittest.TestCase): self.assertTrue(config["load_in_4bit"]) self.assertEqual(config["embedding_learning_rate"], 1e-5) + def test_training_backend_forwards_grad_clipping_controls(self): + backend = TrainingBackend() + + class DummyProcess: + pid = 12345 + + def start(self): + return None + + class DummyThread: + def start(self): + return None + + dummy_queue = object() + + with ( + patch( + "core.training.training.prepare_gpu_selection", + return_value = ([0], {"selection_mode": "auto"}), + ), + patch( + "core.training.training._CTX.Queue", + side_effect = [dummy_queue, dummy_queue], + ), + patch( + "core.training.training._CTX.Process", return_value = DummyProcess() + ) as mock_process, + patch( + "core.training.training.threading.Thread", + return_value = DummyThread(), + ), + ): + backend.start_training( + job_id = "test-grad-clip", + model_name = "unsloth/test", + training_type = "LoRA/QLoRA", + max_grad_norm = 0.7, + max_grad_value = 0.0, + ) + + config = mock_process.call_args.kwargs["kwargs"]["config"] + self.assertEqual(config["max_grad_norm"], 0.7) + self.assertEqual(config["max_grad_value"], 0.0) + def test_training_route_forwards_embedding_learning_rate(self): training_route = _load_route_module( "training_route_module_raw_support", diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 35fbaba8f0..cfdd811853 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -28,6 +28,23 @@ DEFAULT_ALPACA_TEMPLATE = """Below is an instruction that describes a task, pair {}""" +def _is_mlx_runtime() -> bool: + try: + from unsloth_zoo.mlx import is_mlx_available + except ImportError: + return False + return is_mlx_available() + + +def _chat_template_kwargs() -> dict: + if not _is_mlx_runtime(): + return {} + return { + "patch_saving": False, + "use_zoo_tokenizer_patch": True, + } + + def get_tokenizer_chat_template(tokenizer, model_name): """ Gets appropriate chat template for tokenizer based on model. @@ -60,6 +77,7 @@ def get_tokenizer_chat_template(tokenizer, model_name): tokenizer = get_chat_template( tokenizer, chat_template = matched_template, + **_chat_template_kwargs(), ) except Exception as e: logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}") @@ -79,6 +97,7 @@ def get_tokenizer_chat_template(tokenizer, model_name): tokenizer = get_chat_template( tokenizer, chat_template = "chatml", + **_chat_template_kwargs(), ) except Exception as e: logger.info(f"⚠️ Failed to apply default ChatML template: {e}") @@ -255,7 +274,11 @@ def apply_chat_template_to_dataset( if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template): try: from unsloth.chat_templates import get_chat_template - tokenizer = get_chat_template(tokenizer, chat_template = "alpaca") + tokenizer = get_chat_template( + tokenizer, + chat_template = "alpaca", + **_chat_template_kwargs(), + ) logger.info(f"📝 Set alpaca chat template on tokenizer for model saving") except Exception as e: logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}") diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 5e68ccd72c..5d81f0df9c 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -83,6 +83,7 @@ export function buildTrainingStartPayload( save_steps: config.saveSteps, eval_steps: config.evalSteps, weight_decay: config.weightDecay, + max_grad_norm: 0.0, random_seed: config.randomSeed, packing: isEmbedding ? false : config.packing, optim: config.optimizerType, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index fb8a2f899e..0cb881e634 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -31,6 +31,7 @@ export interface TrainingStartRequest { save_steps: number; eval_steps: number; weight_decay: number; + max_grad_norm: number; random_seed: number; packing: boolean; optim: string; diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 168e9ad329..f0c90dd9c6 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -211,8 +211,8 @@ def cmd_train(args) -> int: workdir.mkdir(parents = True, exist_ok = True) import mlx.core as mx - from unsloth_zoo.mlx_loader import FastMLXModel - from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.loader import FastMLXModel + from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig hf_token = os.environ.get("HF_TOKEN") or None @@ -440,7 +440,7 @@ def cmd_reload(args) -> int: return _reload_gguf(save_dir, metrics) import mlx.core as mx - from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx.loader import FastMLXModel from mlx_lm import generate hf_token = os.environ.get("HF_TOKEN") or None diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py index fc07a497e7..043f2796aa 100644 --- a/tests/studio/test_is_mlx_dispatch_gate.py +++ b/tests/studio/test_is_mlx_dispatch_gate.py @@ -7,8 +7,9 @@ Two gates drive every dispatch decision in Studio's MLX path: 1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated once at import time and read by Studio worker code to choose between - the GPU and MLX trainer / inference / export paths. Defined as - ``Darwin AND arm64 AND find_spec("mlx") is not None``. + the GPU and MLX trainer / inference / export paths. It delegates to + the shared zoo MLX runtime gate, with a local import barrier while the + paired unsloth-zoo runtime rollout is in flight. 2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is @@ -18,8 +19,8 @@ Two gates drive every dispatch decision in Studio's MLX path: These gates are the canaries for "MLX support accidentally hijacks CUDA/AMD/Intel users". The tests here: - * verify the source-level structure of the ``_IS_MLX`` expression so an - accidental rewrite (e.g. dropping the ``arm64`` check) is caught, + * verify the source-level structure of the ``_IS_MLX`` helper so an + accidental rewrite importing zoo before the local MLX precheck is caught, * exercise the runtime gate logic under a spoofed Darwin+arm64 platform with a fake ``mlx`` module in ``sys.modules`` to confirm both gates flip True together, @@ -64,20 +65,36 @@ def test_is_mlx_gate_uses_three_required_predicates(): target = node.value break assert target is not None, "_IS_MLX assignment not found in unsloth/__init__.py" - assert isinstance(target, ast.BoolOp) and isinstance( - target.op, ast.And - ), "_IS_MLX must be a BoolOp(And) of platform + mlx checks" - + assert isinstance(target, ast.Call), "_IS_MLX must call the shared MLX helper" expr_src = ast.unparse(target) assert ( - "platform.system()" in expr_src and "Darwin" in expr_src - ), "_IS_MLX must check platform.system() == 'Darwin'" + expr_src == "_is_mlx_available()" + ), "_IS_MLX must delegate to the shared MLX runtime gate" + + helper = None + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "_is_mlx_available": + helper = node + break + assert helper is not None, "_is_mlx_available helper not found" + + helper_src = ast.unparse(helper) assert ( - "platform.machine()" in expr_src and "arm64" in expr_src - ), "_IS_MLX must check platform.machine() == 'arm64'" + "platform.system()" in helper_src + and "'Darwin'" in helper_src + and "platform.machine()" in helper_src + and "'arm64'" in helper_src + and "find_spec" in helper_src + and "'mlx'" in helper_src + and "from unsloth_zoo.mlx import is_mlx_available" in helper_src + ), "_IS_MLX helper must precheck local MLX predicates before importing zoo" assert ( - "find_spec" in expr_src and "'mlx'" in expr_src - ), "_IS_MLX must check importlib.util.find_spec('mlx')" + "from unsloth_zoo.mlx import is_mlx_available" in helper_src + and "return is_mlx_available()" in helper_src + ), "_IS_MLX helper must delegate final detection to the shared zoo MLX runtime gate" + assert helper_src.index("UNSLOTH_FORCE_GPU_PATH") < helper_src.index( + "from unsloth_zoo.mlx import is_mlx_available" + ), "_IS_MLX helper must run the local MLX precheck before importing zoo" # --------------------------------------------------------------------------- @@ -87,13 +104,14 @@ def test_is_mlx_gate_uses_three_required_predicates(): # --------------------------------------------------------------------------- -def _evaluate_is_mlx_gate(platform_module, importlib_util): - """Re-evaluate the _IS_MLX expression using injected dependencies. +def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module): + """Re-evaluate the local _is_mlx_available precheck using injected dependencies. - Mirrors the assignment in unsloth/__init__.py exactly. + Mirrors only the cheap import barrier before unsloth imports unsloth_zoo. """ return ( - platform_module.system() == "Darwin" + os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1" + and platform_module.system() == "Darwin" and platform_module.machine() == "arm64" and importlib_util.find_spec("mlx") is not None ) @@ -112,7 +130,9 @@ def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch): monkeypatch.setattr(platform, "system", lambda: "Darwin") monkeypatch.setattr(platform, "machine", lambda: "arm64") - assert _evaluate_is_mlx_gate(platform, importlib.util) is True + import os + + assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is True def test_is_mlx_gate_false_when_mlx_missing(monkeypatch): @@ -133,7 +153,9 @@ def test_is_mlx_gate_false_when_mlx_missing(monkeypatch): monkeypatch.setattr(importlib.util, "find_spec", _no_mlx) - assert _evaluate_is_mlx_gate(platform, importlib.util) is False + import os + + assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False def test_is_mlx_gate_false_on_non_apple_silicon(): @@ -147,7 +169,9 @@ def test_is_mlx_gate_false_on_non_apple_silicon(): pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.") - assert _evaluate_is_mlx_gate(platform, importlib.util) is False + import os + + assert _evaluate_is_mlx_precheck(platform, importlib.util, os) is False # --------------------------------------------------------------------------- diff --git a/tests/studio/test_mlx_training_worker_behaviors.py b/tests/studio/test_mlx_training_worker_behaviors.py index 6c067ea00b..78b229d6e9 100644 --- a/tests/studio/test_mlx_training_worker_behaviors.py +++ b/tests/studio/test_mlx_training_worker_behaviors.py @@ -46,8 +46,8 @@ def test_wandb_init_strips_secret_keys(): def test_local_dataset_loader_uses_load_dataset_path(): src = WORKER.read_text() - assert "_resolve_local_files" in src - assert "_loader_for_files" in src + assert "_resolve_mlx_local_dataset_files" in src + assert "_mlx_local_dataset_loader_for_files" in src assert "data_files = all_files" in src or "data_files=all_files" in src @@ -84,7 +84,7 @@ def test_poll_stop_returns_on_broken_pipe(): def test_unsloth_zoo_mlx_imports_have_friendly_error(): src = WORKER.read_text() - assert "from unsloth_zoo.mlx_loader import FastMLXModel" in src - assert "from unsloth_zoo.mlx_trainer import" in src + assert "from unsloth_zoo.mlx.loader import FastMLXModel" in src + assert "from unsloth_zoo.mlx.trainer import" in src assert "raise ImportError" in src assert "install.sh" in src diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 9b620a5c76..75fb169e83 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -12,16 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os, platform, importlib.util +import os, importlib.util, platform os.environ["UNSLOTH_IS_PRESENT"] = "1" + +def _is_mlx_available(): + # Transitional import barrier: while the paired unsloth-zoo MLX runtime + # rollout is in flight, keep non-Apple-Silicon imports from touching + # unsloth_zoo here. After both PRs are released together and + # unsloth_zoo.mlx is guaranteed to be import-safe on GPU hosts, + # this helper can collapse back to the centralized zoo runtime call below. + if ( + os.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") == "1" + or platform.system() != "Darwin" + or platform.machine() != "arm64" + or importlib.util.find_spec("mlx") is None + ): + return False + try: + from unsloth_zoo.mlx import is_mlx_available + except ImportError: + return False + return is_mlx_available() + + # Detect Apple Silicon + MLX before any torch/numpy imports -_IS_MLX = ( - platform.system() == "Darwin" - and platform.machine() == "arm64" - and importlib.util.find_spec("mlx") is not None -) +_IS_MLX = _is_mlx_available() if _IS_MLX: try: @@ -31,18 +48,18 @@ if _IS_MLX: "Unsloth: MLX support requires `unsloth-zoo` with MLX modules. " "Reinstall with `pip install unsloth-zoo` or rerun install.sh." ) from _e - # The mlx_trainer / mlx_loader submodules ship with unsloth-zoo's MLX + # The mlx.trainer / mlx.loader submodules ship with unsloth-zoo's MLX # support. An older installed unsloth-zoo (e.g. from PyPI before the # MLX release lands) will satisfy `import unsloth_zoo` but be missing # these submodules. Surface the same friendly install hint instead of # a raw ImportError on the submodule path. try: - from unsloth_zoo.mlx_trainer import MLXTrainer, MLXTrainingConfig - from unsloth_zoo.mlx_loader import FastMLXModel + from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( "Unsloth: MLX support requires an unsloth-zoo build that includes " - "`unsloth_zoo.mlx_trainer` and `unsloth_zoo.mlx_loader`. Upgrade with " + "`unsloth_zoo.mlx.trainer` and `unsloth_zoo.mlx.loader`. Upgrade with " "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 8376dc7e39..1c94e10f9f 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -30,11 +30,9 @@ __all__ = [ from transformers import StoppingCriteria, StoppingCriteriaList from torch import LongTensor, FloatTensor from transformers.models.llama.modeling_llama import logger -from .save import patch_saving_functions import os import shutil from .tokenizer_utils import * -from .models._utils import patch_tokenizer import re from .ollama_template_mappers import OLLAMA_TEMPLATES from unsloth_zoo.dataset_utils import ( @@ -1844,6 +1842,8 @@ def get_chat_template( mapping = {"role" : "role", "content" : "content", "user" : "user", "assistant" : "assistant"}, map_eos_token = True, system_message = None, + patch_saving = True, + use_zoo_tokenizer_patch = False, ): assert(type(map_eos_token) is bool) old_tokenizer = tokenizer @@ -2026,6 +2026,12 @@ def get_chat_template( .replace("'user'", "'" + mapping["user"] + "'")\ .replace("'assistant'", "'" + mapping["assistant"] + "'") + if use_zoo_tokenizer_patch: + # Studio MLX avoids the model-utils tokenizer wrapper because that + # import path pulls in Torch/GPU-specific modules before MLX training. + from unsloth_zoo.tokenizer_utils import patch_tokenizer + else: + from .models._utils import patch_tokenizer _, tokenizer = patch_tokenizer(model = None, tokenizer = tokenizer) tokenizer.padding_side = old_padding_side @@ -2059,7 +2065,9 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - tokenizer = patch_saving_functions(tokenizer) + if patch_saving: + from .save import patch_saving_functions + tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile diff --git a/unsloth/device_type.py b/unsloth/device_type.py index a42d2b9fab..9bad9be0e4 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -20,21 +20,40 @@ __all__ = [ "DEVICE_COUNT", "ALLOW_PREQUANTIZED_MODELS", "ALLOW_BITSANDBYTES", + "is_mlx_available", ] -import torch import functools import inspect +import os from unsloth_zoo.utils import Version +def is_mlx_available(): + try: + from unsloth_zoo.mlx import is_mlx_available as _is_mlx_available + except ImportError: + return False + return _is_mlx_available() + + +_IS_MLX = is_mlx_available() + +if not _IS_MLX: + import torch + + @functools.cache def is_hip(): + if _IS_MLX: + return False return bool(getattr(getattr(torch, "version", None), "hip", None)) @functools.cache def get_device_type(): + if _IS_MLX: + return "mlx" if hasattr(torch, "cuda") and torch.cuda.is_available(): if is_hip(): return "hip" @@ -64,6 +83,8 @@ DEVICE_TYPE: str = get_device_type() DEVICE_TYPE_TORCH = DEVICE_TYPE if DEVICE_TYPE_TORCH == "hip": DEVICE_TYPE_TORCH = "cuda" +elif DEVICE_TYPE_TORCH == "mlx": + DEVICE_TYPE_TORCH = "mps" @functools.cache diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index dd5a9cbf0e..f77baea281 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -160,6 +160,10 @@ else: # INTEL GPU Specific Logic if DEVICE_TYPE == "xpu": _gpu_getCurrentRawStream = torch._C._xpu_getCurrentRawStream +elif DEVICE_TYPE == "mlx": + + def _gpu_getCurrentRawStream(_index = 0): + return 0 # NVIDIA GPU Default Logic elif hasattr(torch._C, "_cuda_getCurrentRawStream"): _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream @@ -206,6 +210,11 @@ if DEVICE_TYPE == "xpu": XPU_STREAMS = () WEIGHT_BUFFERS = [] ABSMAX_BUFFERS = [] +elif DEVICE_TYPE == "mlx": + CUDA_STREAMS = () + XPU_STREAMS = () + WEIGHT_BUFFERS = [] + ABSMAX_BUFFERS = [] else: # NVIDIA GPU Default Logic if DEVICE_COUNT > 0: From 4192fe6ebea2a89fde4daf410577ff30bab54bdd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 05:43:58 -0700 Subject: [PATCH 02/10] studio: drop unused max_grad_value schema + route plumbing (#5424) * studio: drop unused max_grad_value schema + route plumbing The MLX worker hardcodes max_grad_value to 5.0 after PR #5340. The schema field, frontend payload type, route forwarder, and start_training kwarg threading were all left in place as a transitional buffer for old clients. The field is now genuinely unused everywhere except inside the MLX worker, so the schema, route forwarder, and config-build entries can go. Pydantic still tolerates older clients that send max_grad_value because TrainingStartRequest's model_config defaults to extra=ignore. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- scripts/verify_comment_only_diff.py | 8 +++++--- studio/backend/core/training/training.py | 1 - studio/backend/models/training.py | 8 -------- studio/backend/routes/training.py | 1 - studio/backend/tests/test_training_raw_support.py | 2 -- 5 files changed, 5 insertions(+), 15 deletions(-) diff --git a/scripts/verify_comment_only_diff.py b/scripts/verify_comment_only_diff.py index 068d3244df..90eafb7f8f 100644 --- a/scripts/verify_comment_only_diff.py +++ b/scripts/verify_comment_only_diff.py @@ -35,6 +35,7 @@ Example: git diff --name-only origin/main..HEAD \\ | xargs python scripts/verify_comment_only_diff.py --base origin/main """ + from __future__ import annotations import argparse @@ -49,7 +50,9 @@ import yaml def _git_show(rev: str, path: str) -> str: return subprocess.check_output( - ["git", "show", f"{rev}:{path}"], text = True, stderr = subprocess.DEVNULL, + ["git", "show", f"{rev}:{path}"], + text = True, + stderr = subprocess.DEVNULL, ) @@ -145,8 +148,7 @@ def _walk_yaml_diff(b: Any, a: Any, prefix: str = "") -> None: elif isinstance(b, list): if len(b) != len(a): print( - f" list len at {prefix or '/'}: " - f"{len(b)} -> {len(a)}", + f" list len at {prefix or '/'}: " f"{len(b)} -> {len(a)}", ) for i, (bi, ai) in enumerate(zip(b, a)): _walk_yaml_diff(bi, ai, f"{prefix}[{i}]") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 556545e680..e4abb64b8b 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -215,7 +215,6 @@ class TrainingBackend: "save_steps": kwargs.get("save_steps", 0), "weight_decay": kwargs.get("weight_decay", 0.001), "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": kwargs.get("max_grad_value"), "random_seed": kwargs.get("random_seed", 3407), "packing": kwargs.get("packing", False), "optim": kwargs.get("optim", "adamw_8bit"), diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 34ae9cfe72..7c53b0fee5 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -267,14 +267,6 @@ class TrainingStartRequest(BaseModel): ge = 0, description = "Global gradient norm clipping threshold. Set 0 to disable.", ) - max_grad_value: Optional[float] = Field( - None, - ge = 0, - description = ( - "Elementwise gradient value clipping threshold. Set 0 to disable. " - "If omitted, MLX defaults to 1 unless max_grad_norm is set." - ), - ) random_seed: int = Field(42, description = "Random seed") packing: bool = Field(False, description = "Enable sequence packing") optim: str = Field("adamw_8bit", description = "Optimizer") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6875a13206..6e2413b3e9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -216,7 +216,6 @@ async def start_training( "save_steps": request.save_steps, "weight_decay": request.weight_decay, "max_grad_norm": request.max_grad_norm, - "max_grad_value": request.max_grad_value, "random_seed": request.random_seed, "packing": request.packing, "optim": request.optim, diff --git a/studio/backend/tests/test_training_raw_support.py b/studio/backend/tests/test_training_raw_support.py index c4aaaf3298..384247a191 100644 --- a/studio/backend/tests/test_training_raw_support.py +++ b/studio/backend/tests/test_training_raw_support.py @@ -107,12 +107,10 @@ class TestTrainingRawSupport(unittest.TestCase): model_name = "unsloth/test", training_type = "LoRA/QLoRA", max_grad_norm = 0.7, - max_grad_value = 0.0, ) config = mock_process.call_args.kwargs["kwargs"]["config"] self.assertEqual(config["max_grad_norm"], 0.7) - self.assertEqual(config["max_grad_value"], 0.0) def test_training_route_forwards_embedding_learning_rate(self): training_route = _load_route_module( From 000ca89301ba8fb9a58ab718b885d567cfbb4ca1 Mon Sep 17 00:00:00 2001 From: "U. I. I. Derbashi" <32719218+uderbashi@users.noreply.github.com> Date: Thu, 14 May 2026 15:48:28 +0200 Subject: [PATCH 03/10] Studio: Passing batch size for eval (#5168) * add eval batch size * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index a3f063694f..62f1e23e60 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3208,6 +3208,9 @@ class UnslothTrainer: if eval_steps_val > 0: config_args["eval_strategy"] = "steps" config_args["eval_steps"] = eval_steps_val + config_args["per_device_eval_batch_size"] = config_args[ + "per_device_train_batch_size" + ] logger.info( f"✅ Evaluation enabled: eval_steps={eval_steps_val} (fraction of total steps)\n" ) From 79adfd9c71ac57fbbc96b0798af83d8cd7949c35 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Thu, 14 May 2026 18:13:50 +0400 Subject: [PATCH 04/10] studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420) * studio: skip flash-attn install on Blackwell GPUs (sm_100+) Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120, or sm_121, and the older-arch wheels fail to load on Blackwell. Add a shared has_blackwell_gpu() helper and gate both the install-time (install_python_stack._ensure_flash_attn) and runtime (worker._ensure_flash_attn_for_long_context) paths on it. Detection uses nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows. * test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests prefers_prebuilt_wheel and falls_back_to_pypi exercise the install paths that the Blackwell guard now short-circuits. Make them explicit about non-Blackwell so they pass on real Blackwell hosts. * studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH - Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a single process avoid redundant nvidia-smi spawns. Tests clear the cache via setup_method/teardown_method. - In _ensure_flash_attn, run the NO_TORCH short-circuit before the Blackwell check so GGUF-only users (who never install torch anyway) do not see a Blackwell warning. Blackwell check still runs above the IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see the explicit reason rather than a silent OS skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: add has_blackwell_gpu to mlx worker test wheel_utils stub test_mlx_training_worker_config loads worker.py against a hand-rolled utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol list so worker's import line resolves. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/worker.py | 7 + .../tests/test_mlx_training_worker_config.py | 1 + .../tests/test_training_worker_flash_attn.py | 25 +++ studio/backend/utils/wheel_utils.py | 44 ++++ studio/install_python_stack.py | 14 +- .../test_flash_attn_install_python_stack.py | 195 ++++++++++++++++++ 6 files changed, 284 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 8cd1d1f9c8..6b3b3b6609 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -30,6 +30,7 @@ from utils.hardware import apply_gpu_ids from utils.wheel_utils import ( direct_wheel_url, flash_attn_wheel_url, + has_blackwell_gpu, install_wheel, probe_torch_wheel_env, url_exists, @@ -313,6 +314,12 @@ def _should_try_runtime_flash_attn_install(max_seq_length: int) -> bool: def _ensure_flash_attn_for_long_context(event_queue: Any, max_seq_length: int) -> None: if not _should_try_runtime_flash_attn_install(max_seq_length): return + if has_blackwell_gpu(): + _send_status( + event_queue, + "Skipping flash-attn install: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel", + ) + return installed = _install_package_wheel_first( event_queue = event_queue, diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 5900af4e3d..98c7bdaa55 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,6 +37,7 @@ def _load_worker_module(): for name in ( "direct_wheel_url", "flash_attn_wheel_url", + "has_blackwell_gpu", "install_wheel", "probe_torch_wheel_env", "url_exists", diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 41a7c87df1..0737bdc82f 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -37,6 +37,7 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -65,6 +66,7 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -112,6 +114,29 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() +def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): + statuses: list[str] = [] + install_mock = mock.Mock() + + monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) + monkeypatch.setattr( + worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True + ) + monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) + monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) + monkeypatch.setattr( + worker, + "_send_status", + lambda queue, message: statuses.append(message), + ) + + worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) + + install_mock.assert_not_called() + assert len(statuses) == 1 + assert "Blackwell" in statuses[0] + + def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 3ed9bda827..5c42e890d1 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -3,6 +3,7 @@ from __future__ import annotations +import functools import json import logging import platform @@ -22,6 +23,49 @@ FLASH_ATTN_RELEASE_BASE_URL = ( ) +@functools.lru_cache(maxsize = 1) +def has_blackwell_gpu() -> bool: + """Return True if any visible NVIDIA GPU has compute capability >= 10.0 + (Blackwell: sm_100, sm_120, sm_121, ...). + + Dao-AILab does not publish prebuilt flash-attention wheels for these + architectures, and the older-arch wheels fail to load on Blackwell, so + callers use this gate to skip the flash-attn install/upgrade path. + + Result is cached for the process lifetime since GPU hardware does not + change. Tests that mock subprocess/nvidia-smi must call + ``has_blackwell_gpu.cache_clear()`` before each invocation. + """ + exe = shutil.which("nvidia-smi") + if not exe: + return False + try: + result = subprocess.run( + [exe, "--query-gpu=compute_cap", "--format=csv,noheader"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + env = child_env_without_native_path_secret(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + if result.returncode != 0: + return False + for line in result.stdout.splitlines(): + cap = line.strip() + if not cap: + continue + major_part = cap.split(".", 1)[0] + try: + major = int(major_part) + except ValueError: + continue + if major >= 10: + return True + return False + + def linux_wheel_platform_tag() -> str | None: machine = platform.machine().lower() if sys.platform.startswith("linux"): diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 3fd1e6af66..ab234ad566 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -28,6 +28,7 @@ if str(_BACKEND_DIR) not in sys.path: from backend.utils.wheel_utils import ( flash_attn_package_version, flash_attn_wheel_url, + has_blackwell_gpu, install_wheel, probe_torch_wheel_env, url_exists, @@ -628,10 +629,19 @@ def _flash_attn_install_disabled() -> bool: def _ensure_flash_attn() -> None: - if NO_TORCH or IS_WINDOWS or IS_MACOS: - return if _flash_attn_install_disabled(): return + if NO_TORCH: + return + if has_blackwell_gpu(): + _step( + "warning", + "Skipping flash-attn: Blackwell GPU detected (sm_100+); no compatible prebuilt wheel", + _cyan, + ) + return + if IS_WINDOWS or IS_MACOS: + return if ( subprocess.run( [sys.executable, "-c", "import flash_attn"], diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 9881f2258e..49f4350a7b 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -10,8 +10,133 @@ from unittest import mock STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio" sys.path.insert(0, str(STUDIO_DIR)) +sys.path.insert(0, str(STUDIO_DIR / "backend")) import install_python_stack as ips +from backend.utils import wheel_utils + + +def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess: + return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "") + + +class TestHasBlackwellGpu: + def setup_method(self): + wheel_utils.has_blackwell_gpu.cache_clear() + + def teardown_method(self): + wheel_utils.has_blackwell_gpu.cache_clear() + + def test_returns_false_when_nvidia_smi_missing(self): + with mock.patch.object(wheel_utils.shutil, "which", return_value = None): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_true_for_sm_100(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_true_for_sm_120(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_true_for_sm_121(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_false_for_sm_90(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_false_for_sm_89(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n") + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_mixed_gpus_with_one_blackwell_returns_true(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + return_value = _smi_result("8.0\n10.0\n"), + ), + ): + assert wheel_utils.has_blackwell_gpu() is True + + def test_returns_false_when_nvidia_smi_fails(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + return_value = _smi_result("", returncode = 1), + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_false_on_subprocess_timeout(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10), + ), + ): + assert wheel_utils.has_blackwell_gpu() is False + + def test_returns_false_on_malformed_output(self): + with ( + mock.patch.object( + wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi" + ), + mock.patch.object( + wheel_utils.subprocess, + "run", + return_value = _smi_result("not-a-number\n\n"), + ), + ): + assert wheel_utils.has_blackwell_gpu() is False class TestFlashAttnWheelSelection: @@ -234,6 +359,76 @@ class TestEnsureFlashAttn: mock_probe.assert_not_called() mock_install_wheel.assert_not_called() + def test_blackwell_gpu_skips_install_with_warning(self): + step_messages: list[tuple[str, str]] = [] + + def fake_step(label: str, value: str, color_fn = None): + step_messages.append((label, value)) + + with ( + mock.patch.object(ips, "NO_TORCH", False), + mock.patch.object(ips, "IS_WINDOWS", False), + mock.patch.object(ips, "IS_MACOS", False), + mock.patch.object(ips, "has_blackwell_gpu", return_value = True), + mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, + mock.patch.object(ips, "install_wheel") as mock_install_wheel, + mock.patch.object(ips, "_step", side_effect = fake_step), + mock.patch("subprocess.run", return_value = self._import_check()), + ): + ips._ensure_flash_attn() + + mock_probe.assert_not_called() + mock_install_wheel.assert_not_called() + assert any( + label == "warning" and "Blackwell" in msg for label, msg in step_messages + ) + + def test_blackwell_gpu_on_windows_emits_blackwell_warning(self): + step_messages: list[tuple[str, str]] = [] + + def fake_step(label: str, value: str, color_fn = None): + step_messages.append((label, value)) + + with ( + mock.patch.object(ips, "NO_TORCH", False), + mock.patch.object(ips, "IS_WINDOWS", True), + mock.patch.object(ips, "IS_MACOS", False), + mock.patch.object(ips, "has_blackwell_gpu", return_value = True), + mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, + mock.patch.object(ips, "install_wheel") as mock_install_wheel, + mock.patch.object(ips, "_step", side_effect = fake_step), + mock.patch("subprocess.run", return_value = self._import_check()), + ): + ips._ensure_flash_attn() + + mock_probe.assert_not_called() + mock_install_wheel.assert_not_called() + assert any( + label == "warning" and "Blackwell" in msg for label, msg in step_messages + ) + + def test_non_blackwell_windows_does_not_emit_blackwell_warning(self): + step_messages: list[tuple[str, str]] = [] + + def fake_step(label: str, value: str, color_fn = None): + step_messages.append((label, value)) + + with ( + mock.patch.object(ips, "NO_TORCH", False), + mock.patch.object(ips, "IS_WINDOWS", True), + mock.patch.object(ips, "IS_MACOS", False), + mock.patch.object(ips, "has_blackwell_gpu", return_value = False), + mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, + mock.patch.object(ips, "install_wheel") as mock_install_wheel, + mock.patch.object(ips, "_step", side_effect = fake_step), + mock.patch("subprocess.run", return_value = self._import_check()), + ): + ips._ensure_flash_attn() + + mock_probe.assert_not_called() + mock_install_wheel.assert_not_called() + assert not any("Blackwell" in msg for _, msg in step_messages) + class TestInstallPythonStackFlashAttnIntegration: def _run_install(self, *, no_torch: bool, is_macos: bool, is_windows: bool) -> int: From ba833b4ade401115618b8d2c2347ee58e2feafcf Mon Sep 17 00:00:00 2001 From: Tenith Hasintha <102300949+Tenith01@users.noreply.github.com> Date: Thu, 14 May 2026 19:45:27 +0530 Subject: [PATCH 05/10] Fix: Add missing utf-8 encoding to text-mode file operations (#5356) Fixes #2795 by explicitly adding encoding='utf-8' to open() calls. This prevents UnicodeDecodeError on Windows with non-UTF-8 system locales when processing files containing UTF-8 characters. Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- unsloth/kernels/moe/autotune_cache.py | 4 ++-- unsloth/kernels/moe/benchmark/utils.py | 2 +- unsloth/models/sentence_transformer.py | 8 ++++---- unsloth/save.py | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/unsloth/kernels/moe/autotune_cache.py b/unsloth/kernels/moe/autotune_cache.py index f23d9688ea..eac6b02b08 100644 --- a/unsloth/kernels/moe/autotune_cache.py +++ b/unsloth/kernels/moe/autotune_cache.py @@ -71,7 +71,7 @@ def load_cached_config(cache_key: str) -> Optional[Dict[str, Any]]: return None try: - with open(cache_file, "r") as f: + with open(cache_file, "r", encoding = "utf-8") as f: cached_data = json.load(f) # Verify cache is still valid (same device, etc.) @@ -118,7 +118,7 @@ def save_cached_config( } try: - with open(cache_file, "w") as f: + with open(cache_file, "w", encoding = "utf-8") as f: json.dump(cache_data, f, indent = 2) logger.info(f"Saved MoE kernel config cache: {cache_key}") except Exception as e: diff --git a/unsloth/kernels/moe/benchmark/utils.py b/unsloth/kernels/moe/benchmark/utils.py index 21905d8df1..56396687e4 100644 --- a/unsloth/kernels/moe/benchmark/utils.py +++ b/unsloth/kernels/moe/benchmark/utils.py @@ -183,7 +183,7 @@ def save_autotune_results(autotune_cache, mode, ref_time, fused_time, results_di filename = "_".join(key) save_path = f"{save_dir}/{filename}.json" print(f"Saving autotune results to {save_path}") - with open(save_path, "w") as f: + with open(save_path, "w", encoding = "utf-8") as f: result = { **config.all_kwargs(), "ref_time": ref_time, diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 52ff268314..c53e3a7a81 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -70,7 +70,7 @@ def _save_pretrained_torchao( modules_path = os.path.join(save_directory, "modules.json") if os.path.exists(modules_path): try: - with open(modules_path, "r") as f: + with open(modules_path, "r", encoding = "utf-8") as f: modules = json.load(f) for m in modules: if m.get("type", "").endswith("Transformer"): @@ -177,7 +177,7 @@ def _save_pretrained_gguf( modules_path = os.path.join(save_directory, "modules.json") if os.path.exists(modules_path): try: - with open(modules_path, "r") as f: + with open(modules_path, "r", encoding = "utf-8") as f: modules = json.load(f) for m in modules: if m.get("type", "").endswith("Transformer"): @@ -542,7 +542,7 @@ class FastSentenceTransformer(FastModel): model_name, "modules.json", token = token ) - with open(modules_json_path, "r") as f: + with open(modules_json_path, "r", encoding = "utf-8") as f: modules_config = json.load(f) pooling_config_path = None @@ -566,7 +566,7 @@ class FastSentenceTransformer(FastModel): break if pooling_config_path: - with open(pooling_config_path, "r") as f: + with open(pooling_config_path, "r", encoding = "utf-8") as f: pooling_config = json.load(f) # from here: # https://github.com/huggingface/sentence-transformers/blob/main/sentence_transformers/models/Pooling.py#L43 diff --git a/unsloth/save.py b/unsloth/save.py index 472eb14933..d67a1a0550 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -2641,7 +2641,7 @@ This model was finetuned and converted to GGUF format using [Unsloth](https://gi ) readme_path = os.path.join(actual_save_directory, "README.md") - with open(readme_path, "w") as f: + with open(readme_path, "w", encoding = "utf-8") as f: f.write(readme_content) api.upload_file( From 739ebeea8241d125c1bb8fa401207e62dcacb326 Mon Sep 17 00:00:00 2001 From: Tenith Hasintha <102300949+Tenith01@users.noreply.github.com> Date: Thu, 14 May 2026 20:19:47 +0530 Subject: [PATCH 06/10] Fix/issue 3667 vicuna template (#5357) * Fix: Add missing utf-8 encoding to text-mode file operations Fixes #2795 by explicitly adding encoding='utf-8' to open() calls. This prevents UnicodeDecodeError on Windows with non-UTF-8 system locales when processing files containing UTF-8 characters. * fix(chat_templates): escape apostrophe in vicuna default system message The unescaped apostrophe in 'user's' broke Jinja2 template parsing when _change_system_message() substituted the default message into the template string via re.sub. Escape it with \' to match the existing vicuna_old pattern. Fixes #3667 --------- Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- unsloth/chat_templates.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 1c94e10f9f..e8a34cbc60 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -211,7 +211,7 @@ vicuna_ollama = _ollama_template("vicuna") vicuna_eos_token = "eos_token" CHAT_TEMPLATES["vicuna"] = (vicuna_template, vicuna_eos_token, False, vicuna_ollama,) -DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions." +DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user\\'s questions." # =========================================== Vicuna Old # https://github.com/lm-sys/FastChat/blob/main/docs/vicuna_weights_version.md#prompt-template From ab21dc25b41a1d7fd142d6beeea572ed4295ae7d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 19:56:21 -0700 Subject: [PATCH 07/10] tests: public-api surface drift detector (companion to test_import_fixes_drift.py) (#5428) * tests: ship public-api surface drift detector + wire into Core matrix Companion to tests/test_import_fixes_drift.py (PR #5414): that file catches drift in THIRD-PARTY libs (transformers / trl / triton / peft / vllm / torchcodec / xformers); this file catches drift in unsloth's OWN public-surface API -- the top-9 classmethods + symbols that unslothai/notebooks calls at ~2000 cumulative sites. Closes the gap where a refactor on this repo (e.g. renaming FastLanguageModel.from_pretrained -> .load) would pass unsloth CI green and surface only on the next unslothai/notebooks CI run, or worse, on a user's Colab crash report. Coverage (call-site counts measured against unslothai/notebooks main): test_fast_language_model_class_present test_fast_language_model_from_pretrained_kwargs 506 sites test_fast_language_model_get_peft_model_kwargs 304 sites test_fast_language_model_for_inference_callable 370 sites test_fast_vision_model_class_and_methods (4 methods) test_fast_vision_model_get_peft_model_vision_kwargs (4 kwargs) test_fast_model_class_and_methods (2 methods) test_fast_model_from_pretrained_kwargs 103 sites test_is_bf16_supported_or_alias_callable 48 + 8 sites Each test asserts the healthy public shape via inspect.signature; on regression fires pytest.fail("DRIFT DETECTED: ...") -- never pytest.skip -- so the Core matrix cell goes red. Mirrors the same skeleton used by tests/test_import_fixes_drift.py. Wired as a new step in consolidated-tests-ci.yml right after the import_fixes drift step, inside every Core matrix cell. Local verification on transformers 4.57.6 + unsloth main: pytest tests/test_public_api_surface.py -v -> 9 passed in 0.02s * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 11 + tests/test_public_api_surface.py | 216 ++++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 tests/test_public_api_surface.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index abceb91567..218d467f04 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -304,6 +304,17 @@ jobs: run: | python -m pytest -v --tb=short tests/test_import_fixes_drift.py + - name: public-api surface drift detectors (9 tests, HARD GATE) + # Companion to test_import_fixes_drift.py: that file catches + # third-party drift; this one catches drift in unsloth's OWN + # public surface (FastLanguageModel / FastVisionModel / + # FastModel + their classmethods + is_bf16_supported). A + # rename here would silently break the unslothai/notebooks tree + # one PR cycle later -- this gate catches it BEFORE the + # breakage reaches users. + run: | + python -m pytest -v --tb=short tests/test_public_api_surface.py + - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) # 16 tests across 5 files. They live inside tests/saving/ and # tests/utils/, both of which Repo tests (CPU) excludes via --ignore diff --git a/tests/test_public_api_surface.py b/tests/test_public_api_surface.py new file mode 100644 index 0000000000..50cc504749 --- /dev/null +++ b/tests/test_public_api_surface.py @@ -0,0 +1,216 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Public-API surface drift detectors for unsloth itself. + +Companion to tests/test_import_fixes_drift.py: that file catches drift +in THIRD-PARTY libraries (transformers / trl / triton / peft / etc.) +that unsloth's import_fixes patches around. This file catches drift in +unsloth's OWN public-surface API -- the top-10 symbols and classmethods +that the unslothai/notebooks tree (and therefore every user on Colab) +calls. If a refactor on this repo renames FastLanguageModel.from_pretrained +or drops one of the documented kwargs, the test fires DRIFT DETECTED +here BEFORE the breakage reaches users. + +Call-site counts measured against unslothai/notebooks @ main: + FastLanguageModel.from_pretrained 506 + FastLanguageModel.for_inference 370 + FastLanguageModel.get_peft_model 304 + FastVisionModel.for_inference 183 + FastVisionModel.from_pretrained 176 + FastVisionModel.get_peft_model 99 + FastVisionModel.for_training 60 + FastModel.from_pretrained 103 + FastModel.get_peft_model 67 + +Mirrors the unsloth-zoo / unsloth drift-detector skeleton: +``pytest.importorskip("unsloth")`` to gate, assert the healthy upstream +shape, ``pytest.fail("DRIFT DETECTED: ...")`` (never ``pytest.skip``) on +regression so the matrix cell goes red. +""" + +from __future__ import annotations + +import inspect + +import pytest + + +def _signature_param_names(callable_obj) -> set[str]: + try: + sig = inspect.signature(callable_obj) + except (TypeError, ValueError): + return set() + return set(sig.parameters) + + +def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]: + """True if every name in ``kwargs`` is either a named parameter on + ``callable_obj`` OR the callable's signature has a ``**kwargs`` + catch-all. Returns (ok, missing_set).""" + try: + sig = inspect.signature(callable_obj) + except (TypeError, ValueError): + return True, set() + params = sig.parameters + has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + if has_var_kw: + return True, set() + missing = kwargs - set(params) + return (not missing), missing + + +# =========================================================================== +# FastLanguageModel: the headline class. 506 from_pretrained + 370 +# for_inference + 304 get_peft_model call sites across the notebooks. +# =========================================================================== + + +def test_fast_language_model_class_present(): + unsloth = pytest.importorskip("unsloth") + if not hasattr(unsloth, "FastLanguageModel"): + pytest.fail( + "DRIFT DETECTED: unsloth.FastLanguageModel is missing; every " + "LoRA notebook fails at the first import cell." + ) + + +def test_fast_language_model_from_pretrained_kwargs(): + """from_pretrained must accept the canonical kwargs the notebooks pass.""" + unsloth = pytest.importorskip("unsloth") + required = {"model_name", "max_seq_length", "dtype", "load_in_4bit"} + ok, missing = _accepts(unsloth.FastLanguageModel.from_pretrained, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastLanguageModel.from_pretrained dropped " + f"kwargs {sorted(missing)}; 506 notebook call sites would " + f"crash with TypeError." + ) + + +def test_fast_language_model_get_peft_model_kwargs(): + unsloth = pytest.importorskip("unsloth") + required = { + "r", + "lora_alpha", + "lora_dropout", + "target_modules", + "bias", + "use_gradient_checkpointing", + "random_state", + } + ok, missing = _accepts(unsloth.FastLanguageModel.get_peft_model, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastLanguageModel.get_peft_model dropped " + f"kwargs {sorted(missing)}; 304 notebook call sites would crash." + ) + + +def test_fast_language_model_for_inference_callable(): + unsloth = pytest.importorskip("unsloth") + if not callable(getattr(unsloth.FastLanguageModel, "for_inference", None)): + pytest.fail( + "DRIFT DETECTED: FastLanguageModel.for_inference is missing; " + "370 inference-cell call sites would crash." + ) + + +# =========================================================================== +# FastVisionModel: 183 + 176 + 99 + 60 call sites across vision notebooks. +# =========================================================================== + + +def test_fast_vision_model_class_and_methods(): + unsloth = pytest.importorskip("unsloth") + if not hasattr(unsloth, "FastVisionModel"): + pytest.fail( + "DRIFT DETECTED: unsloth.FastVisionModel is missing; every " + "vision fine-tuning notebook fails at import." + ) + cls = unsloth.FastVisionModel + missing = [ + m + for m in ("from_pretrained", "get_peft_model", "for_inference", "for_training") + if not callable(getattr(cls, m, None)) + ] + if missing: + pytest.fail(f"DRIFT DETECTED: FastVisionModel is missing methods {missing}.") + + +def test_fast_vision_model_get_peft_model_vision_kwargs(): + """Vision-specific kwargs the notebooks pass on the vision LoRA path.""" + unsloth = pytest.importorskip("unsloth") + required = { + "finetune_vision_layers", + "finetune_language_layers", + "finetune_attention_modules", + "finetune_mlp_modules", + } + ok, missing = _accepts(unsloth.FastVisionModel.get_peft_model, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastVisionModel.get_peft_model dropped " + f"vision kwargs {sorted(missing)}." + ) + + +# =========================================================================== +# FastModel: the modern unified entry point. 103 + 67 call sites. +# =========================================================================== + + +def test_fast_model_class_and_methods(): + unsloth = pytest.importorskip("unsloth") + if not hasattr(unsloth, "FastModel"): + pytest.fail( + "DRIFT DETECTED: unsloth.FastModel is missing; the modern " + "unified entry point used by 100+ notebooks would crash." + ) + missing = [ + m + for m in ("from_pretrained", "get_peft_model") + if not callable(getattr(unsloth.FastModel, m, None)) + ] + if missing: + pytest.fail(f"DRIFT DETECTED: FastModel is missing methods {missing}.") + + +def test_fast_model_from_pretrained_kwargs(): + unsloth = pytest.importorskip("unsloth") + required = {"model_name", "max_seq_length", "dtype", "load_in_4bit"} + ok, missing = _accepts(unsloth.FastModel.from_pretrained, required) + if not ok: + pytest.fail( + f"DRIFT DETECTED: FastModel.from_pretrained dropped kwargs " + f"{sorted(missing)}; 103 notebook call sites would crash." + ) + + +# =========================================================================== +# Bf16 helper alias (renamed once already; keep both accepted). +# =========================================================================== + + +def test_is_bf16_supported_or_alias_callable(): + """48 notebook import sites for is_bf16_supported plus 8 for the + legacy is_bfloat16_supported alias. Either must remain importable.""" + unsloth = pytest.importorskip("unsloth") + has_new = callable(getattr(unsloth, "is_bf16_supported", None)) + has_old = callable(getattr(unsloth, "is_bfloat16_supported", None)) + if not (has_new or has_old): + pytest.fail( + "DRIFT DETECTED: neither unsloth.is_bf16_supported nor " + "unsloth.is_bfloat16_supported is callable; dtype probing " + "in 50+ notebooks fails." + ) From cb15a7a5b66bec7ae8d4beea45bbda008cadfb9f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 20:27:14 -0700 Subject: [PATCH 08/10] add UNSLOTH_ALLOW_CPU=1 path for CPU-only CI (#5429) Lets `import unsloth.trainer` succeed on hosts without a CUDA/XPU/HIP accelerator (typical of zoo's source-inspection test matrix). The env var is read exactly once per process via @functools.cache on `get_device_type()`, so production hosts pay no runtime cost. Three edits beyond the device_type fallback: * `_gpu_init.py:212/247` -- the bf16 + libcuda/bnb setup blocks call `torch.cuda.get_device_capability()` and `libcuda_dirs()`/`bnb.functional.lib.*` unconditionally when DEVICE_TYPE == "cuda". Guard with `and torch.cuda.is_available()` so the new CPU-CI sentinel doesn't fault those. * `_gpu_init.py:353` -- gate `_patch_trl_trainer()` (the `_backwards_compatible_trainer.__init__` wrapper). Under UNSLOTH_ALLOW_CPU we want pristine upstream TRL classes for downstream `inspect.getsource(SFTTrainer)` drift detectors. * `models/_utils.py:1196` -- same `and torch.cuda.is_available()` guard for `get_device_capability()` at import time. * `models/rl.py:PatchFastRL` -- early-return under UNSLOTH_ALLOW_CPU=1 so the heavier `patch_trl_rl_trainers()` (which replaces `trl.SFTTrainer` with the compiled `UnslothSFTTrainer` class) doesn't fire either. Without this gate the drift detectors that do `inspect.getsource(SFTTrainer)` see the wrapper source and spurious fail. Local sanity: `UNSLOTH_ALLOW_CPU=1 python -c "import unsloth.trainer"` succeeds on a CPU-only venv, `trl.SFTTrainer.__init__.__qualname__` stays `SFTTrainer.__init__` (not `UnslothSFTTrainer.__init__`), and `inspect.getsource(SFTTrainer)` still contains `self._signature_columns`. Without the env var on a CUDA host, TRL is still patched normally (verified `UnslothSFTTrainer.__init__`). --- unsloth/_gpu_init.py | 19 +++++++++++++++---- unsloth/device_type.py | 6 ++++++ unsloth/models/_utils.py | 2 +- unsloth/models/rl.py | 5 +++++ 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index aa94c4a568..a30111b529 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -209,7 +209,7 @@ del fix_peft_transformers_weight_conversion_import del patch_peft_weight_converter_compatibility # Torch 2.4 has including_emulation -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): major_version, minor_version = torch.cuda.get_device_capability() SUPPORTS_BFLOAT16 = major_version >= 8 @@ -233,12 +233,18 @@ elif DEVICE_TYPE == "xpu": # torch.xpu.is_bf16_supported() does not have including_emulation # set SUPPORTS_BFLOAT16 as torch.xpu.is_bf16_supported() SUPPORTS_BFLOAT16 = torch.xpu.is_bf16_supported() +else: + # CPU-only CI under UNSLOTH_ALLOW_CPU=1. We can't probe device + # capability, so assume no bf16 -- training won't run on this host + # anyway, this branch only exists to let `import unsloth.trainer` + # succeed for source-inspection tests. + SUPPORTS_BFLOAT16 = False # For Gradio HF Spaces? # if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ: import triton -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): libcuda_dirs = lambda: None if Version(triton.__version__) >= Version("3.0.0"): try: @@ -349,5 +355,10 @@ from unsloth_zoo.rl_environments import ( launch_openenv, ) -# Patch TRL trainers for backwards compatibility -_patch_trl_trainer() +# Patch TRL trainers for backwards compatibility. +# Skipped under UNSLOTH_ALLOW_CPU=1 (CPU-only CI) because rebinding +# trl.SFTTrainer.__init__ to a generic wrapper changes +# inspect.getsource(SFTTrainer.__init__) and corrupts downstream +# drift detectors that anchor on the pristine upstream source. +if os.environ.get("UNSLOTH_ALLOW_CPU", "0") != "1": + _patch_trl_trainer() diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 9bad9be0e4..6a82e42e8c 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -63,6 +63,10 @@ def get_device_type(): # Check torch.accelerator if hasattr(torch, "accelerator"): if not torch.accelerator.is_available(): + # Test-only CPU fallback. The env var is read exactly once per + # process because get_device_type is @functools.cache'd. + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return "cuda" raise NotImplementedError( "Unsloth cannot find any torch accelerator? You need a GPU." ) @@ -73,6 +77,8 @@ def get_device_type(): f"But `torch.accelerator.current_accelerator()` works with it being = `{accelerator}`\n" f"Please reinstall torch - it's most likely broken :(" ) + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return "cuda" raise NotImplementedError( "Unsloth currently only works on NVIDIA, AMD and Intel GPUs." ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index ccd75aa000..df498e89fb 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1193,7 +1193,7 @@ SUPPORTS_BFLOAT16 = False HAS_FLASH_ATTENTION = False HAS_FLASH_ATTENTION_SOFTCAPPING = False -if DEVICE_TYPE == "cuda": +if DEVICE_TYPE == "cuda" and torch.cuda.is_available(): major_version, minor_version = torch.cuda.get_device_capability() torch.cuda.get_device_capability = functools.cache(torch.cuda.get_device_capability) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index ee9bdda26a..31a498eada 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -2270,6 +2270,11 @@ def patch_trl_vllm_generation(): def PatchFastRL(algorithm = None, FastLanguageModel = None): if FastLanguageModel is not None: PatchRL(FastLanguageModel) + # Under UNSLOTH_ALLOW_CPU=1 (CPU-only CI), skip TRL trainer rewriting so + # downstream `inspect.getsource(trl.SFTTrainer)` drift detectors see the + # pristine upstream class, not the compiled Unsloth* wrappers. + if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1": + return # Install the disable_gradient_checkpointing noop BEFORE # patch_trl_rl_trainers. patch_trl_rl_trainers imports extra trl.* trainer # submodules while generating the compiled cache; any new trl.* modules From 63c6750532bb1c8657b09be87bac31f4f4468802 Mon Sep 17 00:00:00 2001 From: Tai An Date: Thu, 14 May 2026 20:31:20 -0700 Subject: [PATCH 09/10] fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347) (#5350) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347) When a flat local GGUF directory holds several unrelated models with their own mmproj siblings, detect_mmproj_file() returned the first projector it walked into. For the layout reported in #5347 (Qwen weights + a Gemma mmproj in the same dir) that meant llama-server was launched with --mmproj pointing at the Gemma projector, which fails to load and surfaces as a confusing crash. Disambiguation rules: - Drop candidates whose family token (qwen/gemma/llama/mistral/phi/...) disagrees with the model's family. Candidates with no recognised family token (e.g. the HF-convention 'mmproj-F16.gguf') are kept. - Among same-family candidates, prefer the one whose stem shares the longest prefix with the model (Qwen3.5-9B mmproj beats Qwen3.5-35B mmproj for a Qwen3.5-9B model). - If every candidate is dropped, return None — better than attaching a wrong projector and getting a server-launch failure. Tests cover the cross-family block, multi-candidate prefix tie-break, HF-convention 'mmproj-F16.gguf', unrecognised families, and the existing search_root walk. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/mmproj: word-bounded family match, expanded token list, launcher guard Tighten the family-token detector to match only on word boundaries so substring collisions stop tagging false families: phi no longer matches sapphire, yi no longer matches yip, mimo no longer matches mimosa, and mistral does not bleed into ministral/magistral/devstral. Pick the token whose first occurrence is leftmost in the filename rather than the first hit in tuple order, so merge models disambiguate predictably (llama-phi tags llama; phi-llama tags phi). Expand _MODEL_FAMILY_TOKENS with the families an audit of the unsloth HF org turned up that the previous list missed: devstral, ministral, magistral (Mistral-derivative naming), nemotron, kimi, nanonets, cosmos, mimo, apriel, lfm. Without these, a flat local GGUF directory containing one of these weights plus an unrelated renamed projector still hit the original #5347 failure. Add mmproj_matches_model_family() and call it at the llama-server launch site in core/inference/llama_cpp.py. detect_mmproj_file already drops cross-family candidates at discovery time, but mmproj_path can also reach the launcher via config injection or future overrides; this guard keeps those paths from silently loading a known-wrong projector. Tests: 12 new cases covering substring rejection, leftmost-position selection, new family tokens, a new flat-dir Nemotron + Gemma rejection case, and the launcher-level guard. All 21 detect_mmproj_file tests and the existing 106 llama_cpp tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/mmproj: pair via GGUF general.* metadata, not just filenames Real Unsloth vision GGUFs carry rich identity metadata that has been ignored by the discovery path. Every projector under the unsloth org has general.type='mmproj' plus general.base_model.0.repo_url pointing at the same upstream HF repo as its weight, and the equivalent basename, base_model.0.name, and base_model.0.organization fields. A flat-dir mismatch is therefore decidable from the headers alone, no matter how the user has renamed the files. Add utils/models/gguf_metadata.py with read_gguf_general_metadata(): a fast (~30 ms) header walk that pulls only the general.* string fields and skips everything else, cached by (resolved path, mtime_ns, size). Mirrors the parser shape already used by LlamaCppBackend._read_gguf_metadata so the format handling is consistent. is_mmproj_by_metadata() returns True/False/None from general.type, and pairing_score() returns 100 for an exact base_model URL match, 80 for basename plus organization match, 60 for basename only, -1 for definitive metadata disagreement, and 0 when neither side has enough metadata to decide. Rewire detect_mmproj_file() to a two-stage selector: 1. Detect projectors via metadata (general.type) when present, else fall back to the filename substring heuristic. This recovers headerless projectors AND projectors whose name does not contain 'mmproj' but whose header advertises one. 2. Score each candidate against the weight via pairing_score. Drop candidates with score -1 (definitive metadata disagreement). For candidates with score 0 (no usable metadata) fall back to the existing filename family-token check, dropping recognised-family mismatches. Pick the survivor with the highest (score, longest_prefix, -len(stem)) tuple, so a metadata URL match always wins over a filename-prefix match. Tests: 16 new cases. tests/test_gguf_metadata.py covers the parser (missing file, non-GGUF, string extraction, walking past arrays and uint32s, cache invalidation by mtime/size) and the score helpers. tests/test_detect_mmproj_file.py adds end-to-end cases that synthesise real on-disk GGUF headers: URL match wins over a longer-prefix sibling, URL mismatch returns None even when filenames match, a projector named 'vision-projector.gguf' is still discovered via general.type, and a 100-score header match outranks a near-perfect filename prefix on a headerless candidate. All 75 tests across detect_mmproj_file, gguf_metadata, llama_cpp load progress, cached gguf routes, trained model scan, and vision cache pass. * studio/mmproj: shorten comments and docstrings across the #5347 changes Trim verbose explanations to one-line statements of intent. The behaviour is unchanged: 161 tests across detect_mmproj_file, gguf_metadata, llama_cpp_load_progress (+ matrix), llama_server_args, llama_cpp_cache_aware_disk_check, trained_model_scan, and vision_cache all pass. * studio/mmproj: shorten remaining detect_mmproj_file body comments Trim the docstring and the dir-walking block comments inside detect_mmproj_file to one-liners. Behaviour unchanged; 44 mmproj + gguf_metadata + llama_cpp_load_progress tests pass. * studio/mmproj: cap gguf_metadata cache below ceiling on every insert The eviction branch popped exactly one entry when len >= max, so the cache size could only converge to the cap when entries were added slowly enough for natural growth. After a sandbox sim that reduced the cap mid-run, len stayed above the cap because each insert popped one and added one. Switch to a while loop so we evict until len is strictly below the cap before inserting. Steady-state behaviour at the default 4096 ceiling is unchanged. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 18 +- .../backend/tests/test_detect_mmproj_file.py | 326 ++++++++++++++++++ studio/backend/tests/test_gguf_metadata.py | 216 ++++++++++++ studio/backend/utils/models/gguf_metadata.py | 233 +++++++++++++ studio/backend/utils/models/model_config.py | 198 ++++++++--- 5 files changed, 949 insertions(+), 42 deletions(-) create mode 100644 studio/backend/tests/test_detect_mmproj_file.py create mode 100644 studio/backend/tests/test_gguf_metadata.py create mode 100644 studio/backend/utils/models/gguf_metadata.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 35933e6685..7ef687035c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2367,8 +2367,20 @@ class LlamaCppBackend: if not Path(mmproj_path).is_file(): logger.warning(f"mmproj file not found: {mmproj_path}") else: - cmd.extend(["--mmproj", mmproj_path]) - logger.info(f"Using mmproj for vision: {mmproj_path}") + # #5347 guard for paths that bypass detect_mmproj_file. + from utils.models.model_config import ( + mmproj_matches_model_family, + ) + + if not mmproj_matches_model_family(model_path, mmproj_path): + logger.warning( + f"Skipping mmproj with mismatched family: " + f"model={Path(model_path).name}, " + f"mmproj={Path(mmproj_path).name}" + ) + else: + cmd.extend(["--mmproj", mmproj_path]) + logger.info(f"Using mmproj for vision: {mmproj_path}") # Option C: add --api-key for direct client access when enabled import os as _os @@ -3747,7 +3759,7 @@ class LlamaCppBackend: except json.JSONDecodeError: logger.debug( - f"Skipping malformed SSE line: " f"{line[:100]}" + f"Skipping malformed SSE line: {line[:100]}" ) if _stream_done: break # exit outer for diff --git a/studio/backend/tests/test_detect_mmproj_file.py b/studio/backend/tests/test_detect_mmproj_file.py new file mode 100644 index 0000000000..cdb73448be --- /dev/null +++ b/studio/backend/tests/test_detect_mmproj_file.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for :func:`utils.models.model_config.detect_mmproj_file` (#5347).""" + +from __future__ import annotations + +from pathlib import Path + +import struct + +from utils.models.model_config import ( + _detect_family_token, + detect_mmproj_file, + mmproj_matches_model_family, +) + + +_GGUF_MAGIC = 0x46554747 + + +def _gguf_with_general(path: Path, fields: dict) -> Path: + """Write a minimal GGUF with only ``general.*`` string KVs.""" + body = b"" + for k, v in fields.items(): + kb = k.encode("utf-8") + vb = v.encode("utf-8") + body += struct.pack(" Path: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(b"") + return path + + +def test_returns_none_when_no_mmproj(tmp_path: Path): + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + assert detect_mmproj_file(str(model)) is None + + +def test_single_matching_family_mmproj_picked(tmp_path: Path): + """Single same-family projector: returned (historical behaviour).""" + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + mmproj = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(mmproj.resolve()) + + +def test_hf_style_unprefixed_mmproj_still_works(tmp_path: Path): + """HF convention: weight + ``mmproj-F16.gguf`` sibling.""" + model = _touch(tmp_path / "model.gguf") + mmproj = _touch(tmp_path / "mmproj-F16.gguf") + assert detect_mmproj_file(str(model)) == str(mmproj.resolve()) + + +def test_blocks_single_cross_family_projector(tmp_path: Path): + """#5347 core: Qwen weight + lone Gemma mmproj returns None.""" + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + assert detect_mmproj_file(str(model)) is None + + +def test_picks_matching_family_among_mixed_candidates(tmp_path: Path): + """Mixed Qwen + Gemma projectors: pick Qwen, drop Gemma.""" + model = _touch(tmp_path / "Qwen3.5-9B-Q4_K_M.gguf") + qwen_mm = _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + assert detect_mmproj_file(str(model)) == str(qwen_mm.resolve()) + + +def test_prefers_longest_prefix_within_same_family(tmp_path: Path): + """Same family, different sizes: longest shared stem prefix wins.""" + model = _touch(tmp_path / "Qwen3.5-35B-A3B-UD-Q4_K_L.gguf") + _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + big_mm = _touch(tmp_path / "Qwen3.5-35B-A3B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(big_mm.resolve()) + + +def test_unrecognised_family_does_not_break_detection(tmp_path: Path): + """Unknown model family must not return None on a sole candidate.""" + model = _touch(tmp_path / "MyCustomBrand-7B-Q4_K_M.gguf") + mmproj = _touch(tmp_path / "MyCustomBrand-7B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(mmproj.resolve()) + + +def test_directory_path_returns_first_candidate(tmp_path: Path): + """Directory path: no model stem to compare; legacy first-candidate.""" + _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + result = detect_mmproj_file(str(tmp_path)) + assert result is not None + assert "mmproj" in Path(result).name.lower() + + +def test_search_root_walk_still_works(tmp_path: Path): + """Snapshot layout: weight in quant subdir, mmproj at snapshot root.""" + snapshot = tmp_path / "snapshot" + weight = _touch(snapshot / "BF16" / "Qwen3.5-9B-BF16.gguf") + mmproj = _touch(snapshot / "Qwen3.5-9B-BF16-mmproj.gguf") + result = detect_mmproj_file(str(weight), search_root = str(snapshot)) + assert result == str(mmproj.resolve()) + + +# -- Family token detection: word-bounded matching ---------------------- + + +def test_family_token_phi_does_not_match_sapphire(): + """``phi`` substring inside ``sapphire`` must not tag Phi.""" + assert _detect_family_token("sapphire-7b-q4_k_m.gguf") is None + + +def test_family_token_yi_does_not_match_tinyish_names(): + """``yi`` must not cross letter boundaries (``yip``).""" + assert _detect_family_token("yip-7b.gguf") is None + assert _detect_family_token("yi-vl-6b.gguf") == "yi" + + +def test_family_token_mimo_does_not_match_mimosa(): + """``mimo`` must not tag ``mimosa``.""" + assert _detect_family_token("mimosa-rosa-7b.gguf") is None + assert _detect_family_token("MiMo-VL-7B-RL-BF16.gguf") == "mimo" + + +def test_family_token_mistral_does_not_match_ministral(): + """Pin Mistral-derivative tagging.""" + assert _detect_family_token("Ministral-3-8B-Instruct-2512-BF16.gguf") == "ministral" + assert _detect_family_token("Mistral-7B-Instruct-v0.3.gguf") == "mistral" + assert _detect_family_token("Magistral-Small-2506-BF16.gguf") == "magistral" + assert ( + _detect_family_token("Devstral-Small-2-24B-Instruct-2512-BF16.gguf") + == "devstral" + ) + + +def test_family_token_picks_leftmost_when_multiple_present(): + """Leftmost family token wins, not tuple order.""" + assert _detect_family_token("llama-phi-merge.gguf") == "llama" + assert _detect_family_token("phi-llama-merge.gguf") == "phi" + assert _detect_family_token("llama3-3b-instruct.gguf") == "llama" + + +def test_family_token_new_families_recognised(): + """Catalogue-audit additions tag correctly.""" + assert _detect_family_token("NVIDIA-Nemotron-3-Nano-Omni-30B.gguf") == "nemotron" + assert _detect_family_token("Kimi-K2.6-BF16.gguf") == "kimi" + assert _detect_family_token("Nanonets-OCR-s-BF16.gguf") == "nanonets" + assert _detect_family_token("Cosmos-Reason1-7B-BF16.gguf") == "cosmos" + assert _detect_family_token("Apriel-1.5-15b-Thinker-BF16.gguf") == "apriel" + assert _detect_family_token("LFM2.5-VL-1.6B-BF16.gguf") == "lfm" + + +# -- Cross-family rejection with the expanded token list ---------------- + + +def test_blocks_cross_family_for_new_token_pair(tmp_path: Path): + """Nemotron weight + lone Gemma projector returns None.""" + model = _touch( + tmp_path / "NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-MXFP4_MOE.gguf" + ) + _touch(tmp_path / "gemma-4-26B-A4B-it.mmproj-q8_0.gguf") + assert detect_mmproj_file(str(model)) is None + + +def test_picks_devstral_mmproj_in_mixed_dir(tmp_path: Path): + """Devstral weight + Devstral mmproj + a Qwen mmproj: pick Devstral.""" + model = _touch(tmp_path / "Devstral-Small-2-24B-Instruct-2512-BF16.gguf") + dev_mm = _touch(tmp_path / "Devstral-Small-2-mmproj-bf16.gguf") + _touch(tmp_path / "Qwen3.5-9B-BF16-mmproj.gguf") + assert detect_mmproj_file(str(model)) == str(dev_mm.resolve()) + + +# -- Launcher-level family guard ---------------------------------------- + + +def test_mmproj_family_guard_blocks_cross_family(): + assert ( + mmproj_matches_model_family( + "/models/Qwen3.5-9B-Q4_K_M.gguf", + "/models/gemma-4-26B-A4B-it.mmproj-q8_0.gguf", + ) + is False + ) + + +def test_mmproj_family_guard_allows_same_family(): + assert ( + mmproj_matches_model_family( + "/models/Qwen3.5-9B-Q4_K_M.gguf", + "/models/Qwen3.5-9B-BF16-mmproj.gguf", + ) + is True + ) + + +def test_mmproj_family_guard_allows_generic_hf_mmproj(): + """No family token on the projector: wildcard.""" + assert ( + mmproj_matches_model_family( + "/models/Qwen3.5-9B-Q4_K_M.gguf", + "/models/mmproj-F16.gguf", + ) + is True + ) + + +def test_mmproj_family_guard_allows_unrecognised_model_family(): + """No family token on the model: wildcard.""" + assert ( + mmproj_matches_model_family( + "/models/Apriel-1.5-15b-Thinker-BF16.gguf", + "/models/mmproj-F16.gguf", + ) + is True + ) + + +# -- Metadata-primary pairing in detect_mmproj_file --------------------- + + +def test_metadata_url_match_picked_over_filename_lookalike(tmp_path: Path): + """URL match beats a longer-prefix sibling.""" + weight = _gguf_with_general( + tmp_path / "Qwen3.5-9B-Q4_K_M.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + # Closer filename prefix, wrong upstream. + _gguf_with_general( + tmp_path / "Qwen3.5-9B-mmproj-bf16.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-1.5B", + }, + ) + # Matching upstream. + correct = _gguf_with_general( + tmp_path / "mmproj-BF16.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + assert detect_mmproj_file(str(weight)) == str(correct.resolve()) + + +def test_metadata_url_mismatch_dropped(tmp_path: Path): + """Filenames match family but metadata disagrees: returns None.""" + weight = _gguf_with_general( + tmp_path / "qwen-9b.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + _gguf_with_general( + tmp_path / "qwen-9b-mmproj.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.base_model.0.repo_url": "https://huggingface.co/google/gemma-3-9B", + }, + ) + assert detect_mmproj_file(str(weight)) is None + + +def test_metadata_identifies_mmproj_without_filename_hint(tmp_path: Path): + """Projector named ``vision-projector.gguf`` discovered via header.""" + weight = _gguf_with_general( + tmp_path / "Qwen3.5-9B.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + projector = _gguf_with_general( + tmp_path / "vision-projector.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + assert detect_mmproj_file(str(weight)) == str(projector.resolve()) + + +def test_metadata_score_outranks_filename_prefix(tmp_path: Path): + """Score 100 (URL match) beats score 0 (long filename prefix).""" + weight = _gguf_with_general( + tmp_path / "Qwen3.5-9B-Q4_K_M.gguf", + { + "general.architecture": "qwen2vl", + "general.type": "model", + "general.basename": "Qwen3.5", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + # Headerless: long shared stem, score 0. + _touch(tmp_path / "Qwen3.5-9B-Q4_K_M-mmproj.gguf") + # Headered: generic name, score 100. + correct = _gguf_with_general( + tmp_path / "mmproj-BF16.gguf", + { + "general.architecture": "clip", + "general.type": "mmproj", + "general.base_model.0.repo_url": "https://huggingface.co/Qwen/Qwen3.5-9B", + }, + ) + assert detect_mmproj_file(str(weight)) == str(correct.resolve()) diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py new file mode 100644 index 0000000000..cf1a17347f --- /dev/null +++ b/studio/backend/tests/test_gguf_metadata.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for :mod:`utils.models.gguf_metadata`. Synthesise small GGUF +headers in tmp dirs so we never depend on real model files.""" + +from __future__ import annotations + +import struct +from pathlib import Path +from typing import Iterable, Mapping + +from utils.models.gguf_metadata import ( + is_mmproj_by_metadata, + pairing_score, + read_gguf_general_metadata, +) + + +_GGUF_MAGIC = 0x46554747 +_VTYPE_STRING = 8 +_VTYPE_UINT32 = 4 +_VTYPE_ARRAY = 9 + + +def _enc_string(s: str) -> bytes: + b = s.encode("utf-8") + return struct.pack(" bytes: + return _enc_string(key) + struct.pack(" bytes: + return ( + _enc_string(key) + struct.pack(" bytes: + vals = list(values) + out = _enc_string(key) + struct.pack(" Path: + """Minimal GGUF: header + KV body, no tensors.""" + extra_uint32 = extra_uint32 or {} + extra_string_arrays = extra_string_arrays or {} + kv_count = len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + body = b"" + for k, v in general_strings.items(): + body += _enc_kv_string(k, v) + for k, v in extra_uint32.items(): + body += _enc_kv_uint32(k, v) + for k, v in extra_string_arrays.items(): + body += _enc_kv_string_array(k, v) + header = struct.pack( + " Optional[_CacheKey]: + try: + st = os.stat(path) + except OSError: + return None + try: + resolved = str(Path(path).resolve()) + except OSError: + resolved = str(path) + return (resolved, st.st_mtime_ns, st.st_size) + + +def read_gguf_general_metadata(path: str) -> Optional[Dict[str, str]]: + """Return ``general.*`` strings from a GGUF header, or ``None`` if + the file is missing, unreadable, or not a GGUF. ``{}`` means the + file is valid but carries none of the wanted keys.""" + key = _cache_key(path) + if key is None: + return None + with _CACHE_LOCK: + if key in _METADATA_CACHE: + return _METADATA_CACHE[key] + result = _parse_gguf_header(path) + with _CACHE_LOCK: + # Arbitrary eviction; header reads are cheap so true LRU is overkill. + while len(_METADATA_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _METADATA_CACHE.pop(next(iter(_METADATA_CACHE))) + except StopIteration: + break + _METADATA_CACHE[key] = result + return result + + +def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: + out: Dict[str, str] = {} + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: # 1 MB sanity bound + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: # 4 MB sanity bound + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + out[key] = sbytes.decode("utf-8", "replace") + else: + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"read_gguf_general_metadata: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"read_gguf_general_metadata: parse failure on {path}: {e}") + return None + return out + + +# Strings (8) and arrays (9) are handled inline. +_FIXED_VTYPE_SIZES: Dict[int, int] = { + 0: 1, # uint8 + 1: 1, # int8 + 2: 2, # uint16 + 3: 2, # int16 + 4: 4, # uint32 + 5: 4, # int32 + 6: 4, # float32 + 7: 1, # bool + 10: 8, # uint64 + 11: 8, # int64 + 12: 8, # float64 +} + + +def _skip_gguf_value(f, vtype: int) -> bool: + """Advance past one GGUF value. False on truncation or unknown type.""" + if vtype == 8: # STRING + slen_bytes = f.read(8) + if len(slen_bytes) < 8: + return False + slen = struct.unpack(" 1 << 30: # 1 GB sanity bound + return False + return len(f.read(slen)) == slen + if vtype == 9: # ARRAY + head = f.read(12) + if len(head) < 12: + return False + atype, alen = struct.unpack(" 1 << 30: + return False + if atype == 8: + for _ in range(alen): + slen_bytes = f.read(8) + if len(slen_bytes) < 8: + return False + slen = struct.unpack(" 1 << 30: + return False + if len(f.read(slen)) != slen: + return False + return True + sz = _FIXED_VTYPE_SIZES.get(atype) + if sz is None: + return False + total = sz * alen + return len(f.read(total)) == total + sz = _FIXED_VTYPE_SIZES.get(vtype) + if sz is None: + return False + return len(f.read(sz)) == sz + + +def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]: + """True/False from ``general.type``; None means fall back to filename.""" + if not meta: + return None + t = meta.get("general.type") + if t is None: + return None + return t.lower() == "mmproj" + + +def pairing_score( + weight_meta: Optional[Dict[str, str]], + mmproj_meta: Optional[Dict[str, str]], +) -> int: + """Pairing confidence: 100 = base_model URL match, 80 = basename + org, + 60 = basename, -1 = definitive mismatch, 0 = decide from filename.""" + if not weight_meta or not mmproj_meta: + return 0 + + w_url = weight_meta.get("general.base_model.0.repo_url") + p_url = mmproj_meta.get("general.base_model.0.repo_url") + if w_url and p_url: + return 100 if w_url.strip().rstrip("/") == p_url.strip().rstrip("/") else -1 + + w_base = weight_meta.get("general.basename") + p_base = mmproj_meta.get("general.basename") + w_org = weight_meta.get("general.base_model.0.organization") or weight_meta.get( + "general.organization" + ) + p_org = mmproj_meta.get("general.base_model.0.organization") or mmproj_meta.get( + "general.organization" + ) + if w_base and p_base and w_org and p_org: + if w_base.lower() == p_base.lower() and w_org.lower() == p_org.lower(): + return 80 + return -1 + + if w_base and p_base: + return 60 if w_base.lower() == p_base.lower() else -1 + + return 0 diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index ebf85c5320..bf7f7a009b 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -19,6 +19,11 @@ from utils.paths import ( resolve_export_dir, ) from utils.utils import without_hf_auth +from utils.models.gguf_metadata import ( + is_mmproj_by_metadata, + pairing_score, + read_gguf_general_metadata, +) import structlog from loggers import get_logger import os @@ -801,12 +806,15 @@ _AUDIO_TOKEN_PATTERNS = { "whisper": lambda tokens: "<|startoftranscript|>" in tokens, "audio_vlm": lambda tokens: "" in tokens, "bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens), - "dac": lambda tokens: "<|audio_start|>" in tokens - and "<|audio_end|>" in tokens - and "<|text_start|>" in tokens - and "<|text_end|>" in tokens, - "snac": lambda tokens: sum(1 for t in tokens if t.startswith(" 10000, + "dac": lambda tokens: ( + "<|audio_start|>" in tokens + and "<|audio_end|>" in tokens + and "<|text_start|>" in tokens + and "<|text_end|>" in tokens + ), + "snac": lambda tokens: ( + sum(1 for t in tokens if t.startswith(" 10000 + ), } @@ -913,6 +921,85 @@ def _is_mmproj(filename: str) -> bool: return "mmproj" in filename.lower() +# Family tokens for #5347's filename fallback. Lowercase. Order does not +# matter (see ``_detect_family_token``). +_MODEL_FAMILY_TOKENS: tuple[str, ...] = ( + "qwen", + "gemma", + "llama", + "mistral", + "ministral", + "magistral", + "devstral", + "phi", + "deepseek", + "internvl", + "minicpm", + "llava", + "glm", + "yi", + "command-r", + "molmo", + "pixtral", + "smolvlm", + "moondream", + "granite", + "ovis", + "nemotron", + "kimi", + "nanonets", + "cosmos", + "mimo", + "apriel", + "lfm", +) + + +# Word-bounded match: any letter on either side disqualifies. Stops +# ``phi`` matching ``sapphire``, ``yi`` matching ``tiny``, etc. +_FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {} + + +def _family_token_re(token: str) -> "_re.Pattern[str]": + pat = _FAMILY_TOKEN_RE_CACHE.get(token) + if pat is None: + pat = _re.compile(rf"(?:^|[^a-z])({_re.escape(token)})(?:[^a-z]|$)") + _FAMILY_TOKEN_RE_CACHE[token] = pat + return pat + + +def _detect_family_token(filename: str) -> Optional[str]: + """Leftmost-position match; ties prefer the longer token.""" + name = filename.lower() + best: Optional[tuple[int, int, str]] = None # (start, -len, token) + for token in _MODEL_FAMILY_TOKENS: + m = _family_token_re(token).search(name) + if m is None: + continue + key = (m.start(1), -len(token), token) + if best is None or key < best: + best = key + return None if best is None else best[2] + + +def mmproj_matches_model_family(model_path: str, mmproj_path: str) -> bool: + """Defense-in-depth guard for the launcher: True unless both filenames + carry recognised family tokens that disagree.""" + model_fam = _detect_family_token(Path(model_path).name) + mmproj_fam = _detect_family_token(Path(mmproj_path).name) + if model_fam is None or mmproj_fam is None: + return True + return model_fam == mmproj_fam + + +def _shared_prefix_len(a: str, b: str) -> int: + n = min(len(a), len(b)) + for i in range(n): + if a[i] != b[i]: + return i + return n + + def _is_gguf_filename(filename: str) -> bool: return filename.lower().endswith(".gguf") @@ -927,33 +1014,18 @@ def _iter_gguf_files(directory: Path, recursive: bool = False): def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]: - """ - Find the mmproj (vision projection) GGUF file for a given model. + """Find the mmproj GGUF for a model. - Args: - path: Directory to search — or a .gguf file (uses its parent dir - as the starting point). - search_root: Optional outer directory that should also be scanned - (and any directory between it and ``path``). This handles - local layouts where the model weights live in a quant-named - subdir (``snapshot/BF16/foo.gguf``) but the mmproj sits at - the snapshot root (``snapshot/mmproj-BF16.gguf``). When - ``None``, only the immediate parent dir is scanned, matching - the historical behavior. - - Returns: - Full path to the mmproj .gguf file, or None if not found. - """ + ``path``: directory or a .gguf file. ``search_root``: optional ancestor + to also walk (snapshot layouts where the weight is in ``snapshot/BF16/`` + but the projector sits at ``snapshot/``). Returns the projector path or + ``None``.""" p = Path(path) start_dir = p.parent if p.is_file() else p if not start_dir.is_dir(): return None - # Build the list of dirs to scan: immediate dir first, then walk up - # to (and including) ``search_root`` if it is an ancestor. We walk - # incrementally rather than recursing into ``search_root`` so we - # don't accidentally pick up an mmproj from a sibling subdir - # belonging to a different model variant. + # Walk incrementally so a sibling subdir's mmproj cannot leak in. seen: set[Path] = set() scan_order: list[Path] = [] @@ -969,12 +1041,7 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional _add(start_dir) - # When ``path`` is a symlink (e.g. Ollama's ``.studio_links/...gguf`` - # -> ``blobs/sha256-...``), the symlink's parent directory rarely - # contains the mmproj sibling; the real mmproj file lives next to - # the symlink target. Add the target's parent to the scan so vision - # GGUFs that are surfaced via symlinks are still recognised as - # vision models. + # Ollama's .studio_links/foo.gguf -> blobs/sha256-...: also scan target dir. try: if p.is_symlink() and p.is_file(): target_parent = p.resolve().parent @@ -986,14 +1053,12 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional try: root_resolved = Path(search_root).resolve() start_resolved = start_dir.resolve() - # Only walk if start_dir is inside (or equal to) search_root. if root_resolved == start_resolved or ( start_resolved.is_relative_to(root_resolved) if hasattr(start_resolved, "is_relative_to") else str(start_resolved).startswith(str(root_resolved) + "/") ): cur = start_resolved - # Walk up from start_dir to (and including) root_resolved. while cur != root_resolved and cur.parent != cur: cur = cur.parent _add(cur) @@ -1002,11 +1067,66 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional except OSError: pass + candidates: list[Path] = [] + seen_resolved: set[Path] = set() for d in scan_order: for f in _iter_gguf_files(d): - if _is_mmproj(f.name): - return str(f.resolve()) - return None + try: + resolved = f.resolve() + except OSError: + continue + if resolved in seen_resolved: + continue + # Prefer ``general.type=='mmproj'``; fall back to filename. + meta = read_gguf_general_metadata(str(resolved)) + by_meta = is_mmproj_by_metadata(meta) + if by_meta is True or (by_meta is None and _is_mmproj(f.name)): + seen_resolved.add(resolved) + candidates.append(resolved) + + if not candidates: + return None + + # Directory path: no model name to compare against; legacy behaviour. + if not p.is_file(): + return str(candidates[0]) + + # Stage 1: GGUF metadata. Stage 2: filename family token (#5347). + model_stem = p.stem.lower() + model_family = _detect_family_token(p.name) + weight_meta = read_gguf_general_metadata(str(p)) + + scored: list[tuple[int, Path]] = [] + for c in candidates: + cand_meta = read_gguf_general_metadata(str(c)) + meta_score = pairing_score(weight_meta, cand_meta) + if meta_score == -1: + logger.info(f"detect_mmproj_file: dropped {c.name} (metadata mismatch)") + continue + if meta_score == 0 and model_family is not None: + # Unrecognised candidate family is a wildcard (``mmproj-F16.gguf``). + cand_family = _detect_family_token(c.name) + if cand_family is not None and cand_family != model_family: + logger.info( + f"detect_mmproj_file: dropped {c.name} " + f"(filename family {cand_family!r} vs model {model_family!r})" + ) + continue + scored.append((meta_score, c)) + + if not scored: + return None + + # Score first, then longest shared prefix, then shorter stem. + best = max( + scored, + key = lambda sc: ( + sc[0], + _shared_prefix_len(model_stem, sc[1].stem.lower()), + -len(sc[1].stem), + ), + ) + return str(best[1]) def detect_gguf_model(path: str) -> Optional[str]: @@ -1360,7 +1480,7 @@ def detect_gguf_model_remote( if attempt < 2: time.sleep(2**attempt) logger.warning( - f"Could not check GGUF files for '{repo_id}' after 3 attempts: " f"{last_err}" + f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}" ) return None From bbd0ba0c259353fc5147a3255af1af34185b9ed2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 14 May 2026 21:57:04 -0700 Subject: [PATCH 10/10] studio/mmproj: skip unwanted GGUF values via seek instead of read (#5431) The previous _skip_gguf_value walked past discarded values with f.read(n), which allocates and immediately drops a Python bytes object. For weight GGUFs that carry tokenizer.ggml.tokens (~150K unicode strings) this wasted ~10 MB of allocation per cold call. Switch the discard path to f.seek(n, 1). The kernel never has to copy the bytes into userspace and Python never allocates. Truncation is now detected on the next read attempt rather than inline (an out-of-range seek on a regular file is legal and the next read returns short). Measured on real downloaded GGUFs (Qwen3.5-4B IQ2_XXS 1.52 GB, bartowski Qwen3.5-4B IQ2_M 1.70 GB, Qwen3.5-4B-MTP IQ2_M 1.94 GB): before: 142 ms cold per weight, ~11 MB read after: 90 ms cold per weight, ~4 MB read Mmproj reads are unaffected (no tokenizer to skip). Cached re-reads remain ~50 microseconds. All 161 in-tree backend tests + 85 isolated sandbox tests pass. --- studio/backend/utils/models/gguf_metadata.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index 9aab88329d..5629bac58b 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -152,7 +152,9 @@ _FIXED_VTYPE_SIZES: Dict[int, int] = { def _skip_gguf_value(f, vtype: int) -> bool: - """Advance past one GGUF value. False on truncation or unknown type.""" + """Advance past one GGUF value. ``f.seek(.., 1)`` past EOF is legal + on a regular file so truncation is detected on the next read; we + only return False for unknown types or sanity-bound overflow.""" if vtype == 8: # STRING slen_bytes = f.read(8) if len(slen_bytes) < 8: @@ -160,7 +162,8 @@ def _skip_gguf_value(f, vtype: int) -> bool: slen = struct.unpack(" 1 << 30: # 1 GB sanity bound return False - return len(f.read(slen)) == slen + f.seek(slen, 1) + return True if vtype == 9: # ARRAY head = f.read(12) if len(head) < 12: @@ -176,18 +179,18 @@ def _skip_gguf_value(f, vtype: int) -> bool: slen = struct.unpack(" 1 << 30: return False - if len(f.read(slen)) != slen: - return False + f.seek(slen, 1) return True sz = _FIXED_VTYPE_SIZES.get(atype) if sz is None: return False - total = sz * alen - return len(f.read(total)) == total + f.seek(sz * alen, 1) + return True sz = _FIXED_VTYPE_SIZES.get(vtype) if sz is None: return False - return len(f.read(sz)) == sz + f.seek(sz, 1) + return True def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]: