From 157cecb25c3c7277a6da33c81a409373adb2b4b1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 4 Jun 2026 00:07:58 -0700 Subject: [PATCH 1/7] Port KTO logps truncation guard to TRL 1.x _compute_logps refactor (#5996) * Port KTO logps truncation guard to TRL 1.x _compute_logps refactor * [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> --- .../test_trl_grpo_pinned_symbols.py | 20 ++++--- unsloth/models/rl_replacements.py | 57 +++++++++++++++++++ 2 files changed, 69 insertions(+), 8 deletions(-) diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index 4c7dcc4234..8f1435ba5f 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -551,12 +551,12 @@ def test_trl_grpo_source_inference_mode_unwrap(tag: str): @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_kto_get_batch_logps_signature(tag: str): - """TRL 0.27+ moved KTOTrainer to trl.experimental.kto and the - canonical kto_trainer.py shrank to a thin re-export wrapper. The - real `get_batch_logps` lives at trl/experimental/kto/kto_trainer.py. - Unsloth's MRO walk in models/rl.py:592-708 already follows - trl.experimental.* parents, so either path is fine — we just - require the symbol to exist SOMEWHERE.""" + """KTO log-prob computation must stay patchable. Through TRL 1.x the + target was KTOTrainer.get_batch_logps; TRL 1.x dropped it and moved the + math into _compute_logps / compute_ref_log_probs calling + selective_log_softmax. unsloth/models/rl_replacements.py patches BOTH + shapes (kto_trainer_get_batch_logps + kto_trainer_align_completion_logps), + so we require EITHER form to exist wherever KTOTrainer lives.""" candidates = [ "trl/trainer/kto_trainer.py", "trl/experimental/kto/kto_trainer.py", @@ -566,11 +566,15 @@ def test_trl_kto_get_batch_logps_signature(tag: str): src = fetch_text("huggingface/trl", tag, path) if src is None: continue + # Legacy: explicit get_batch_logps method. if has_def(src, "get_batch_logps", "func"): return + # TRL 1.x: refactored into _compute_logps + selective_log_softmax. + if has_def(src, "_compute_logps", "func") and "selective_log_softmax" in src: + return pytest.fail( - f"{tag}: KTOTrainer.get_batch_logps not found in any of {candidates}; " - f"unsloth/models/rl_replacements.py:1675 rewrite silently skipped" + f"{tag}: KTO log-prob computation not found in any of {candidates}; " + f"unsloth/models/rl_replacements.py KTO rewrite silently skipped" ) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 31d54675c9..c1d92a31c4 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1962,6 +1962,63 @@ def kto_trainer_get_batch_logps(function_name, function): RL_FUNCTIONS["kto_trainer"].append(kto_trainer_get_batch_logps) +# TRL 1.x dropped KTOTrainer.get_batch_logps and moved the log-prob math into +# _compute_logps / compute_ref_log_probs / _compute_kl_logps, which call +# selective_log_softmax on completion-only tokens. Same truncation hazard as +# above, so clamp logits/ids/mask to the shorter seq length (no-op when equal). +_KTO_COMPLETION_RE = re.compile( + r"(?P[ \t]*)shift_logits = completion_logits\[:, :-1, :\]\.contiguous\(\)\n" + r"(?P=ws)per_token_logps = selective_log_softmax\(\s*shift_logits,\s*" + r"(?P\w+)\[\"completion_input_ids\"\]\[:, 1:\]\.contiguous\(\)\s*\)\n" + r"(?P=ws)per_token_logps\[(?P=var)\[\"completion_mask\"\]\[:, 1:\] == 0\] = 0\.0" +) +_KTO_KL_RE = re.compile( + r"(?P[ \t]*)shift_KL_logits = KL_logits\[:, :-1, :\]\.contiguous\(\)\n" + r"(?P=ws)KL_per_token_logps = selective_log_softmax\(\s*shift_KL_logits,\s*" + r"(?P\w+)\[\"KL_completion_input_ids\"\]\[:, 1:\]\.contiguous\(\)\s*\)\n" + r"(?P=ws)KL_per_token_logps\[(?P=var)\[\"KL_completion_mask\"\]\[:, 1:\] == 0\] = 0\.0" +) + + +def _kto_completion_repl(m): + ws, var = m.group("ws"), m.group("var") + return ( + f"{ws}shift_logits = completion_logits[:, :-1, :].contiguous()\n" + f"{ws}# Unsloth: clamp logits/ids/mask to shorter seq len (model may truncate input_ids)\n" + f'{ws}_uns_ids = {var}["completion_input_ids"][:, 1:].contiguous()\n' + f"{ws}_uns_n = min(shift_logits.shape[1], _uns_ids.shape[1])\n" + f"{ws}per_token_logps = selective_log_softmax(shift_logits[:, :_uns_n], _uns_ids[:, :_uns_n])\n" + f'{ws}per_token_logps[{var}["completion_mask"][:, 1:][:, :_uns_n] == 0] = 0.0' + ) + + +def _kto_kl_repl(m): + ws, var = m.group("ws"), m.group("var") + return ( + f"{ws}shift_KL_logits = KL_logits[:, :-1, :].contiguous()\n" + f"{ws}# Unsloth: clamp logits/ids/mask to shorter seq len (model may truncate input_ids)\n" + f'{ws}_uns_kl_ids = {var}["KL_completion_input_ids"][:, 1:].contiguous()\n' + f"{ws}_uns_kl_n = min(shift_KL_logits.shape[1], _uns_kl_ids.shape[1])\n" + f"{ws}KL_per_token_logps = selective_log_softmax(shift_KL_logits[:, :_uns_kl_n], _uns_kl_ids[:, :_uns_kl_n])\n" + f'{ws}KL_per_token_logps[{var}["KL_completion_mask"][:, 1:][:, :_uns_kl_n] == 0] = 0.0' + ) + + +def kto_trainer_align_completion_logps(function_name, function): + if function_name not in ( + "_compute_logps", + "compute_ref_log_probs", + "_compute_kl_logps", + ): + return function + function = _KTO_COMPLETION_RE.sub(_kto_completion_repl, function) + function = _KTO_KL_RE.sub(_kto_kl_repl, function) + return function + + +RL_FUNCTIONS["kto_trainer"].append(kto_trainer_align_completion_logps) + + # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356 # TRL warns if batch size is not a multiple of num_generations -> fix this. def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source): From 4eac5272476372b37febdbfc541076ab4355f72f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 4 Jun 2026 00:08:19 -0700 Subject: [PATCH 2/7] CI: mark deepseek_ocr2 as known-broken compile timeout (#5995) --- .github/workflows/consolidated-tests-ci.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index f8eccc7d18..94ec26ef31 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -990,6 +990,7 @@ jobs: # First seen on transformers >=5,<6; each represents a slow # or recursive source-rewriter path the zoo can address. "beit": "TimeoutError: compile exceeds per-model budget", + "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", "sam": "TimeoutError: compile exceeds per-model budget", "sam_hq": "TimeoutError: compile exceeds per-model budget", } From b1ee492982791e02a55979866d1e3b2e90c8af9a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 4 Jun 2026 07:17:55 +0000 Subject: [PATCH 3/7] Revert "CI: mark deepseek_ocr2 as known-broken compile timeout (#5995)" This reverts commit 4eac5272476372b37febdbfc541076ab4355f72f. --- .github/workflows/consolidated-tests-ci.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 94ec26ef31..f8eccc7d18 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -990,7 +990,6 @@ jobs: # First seen on transformers >=5,<6; each represents a slow # or recursive source-rewriter path the zoo can address. "beit": "TimeoutError: compile exceeds per-model budget", - "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", "sam": "TimeoutError: compile exceeds per-model budget", "sam_hq": "TimeoutError: compile exceeds per-model budget", } From 636455a7d6dad04a941ca58a93e2f27e01abc11b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 4 Jun 2026 07:17:55 +0000 Subject: [PATCH 4/7] Revert "Port KTO logps truncation guard to TRL 1.x _compute_logps refactor (#5996)" This reverts commit 157cecb25c3c7277a6da33c81a409373adb2b4b1. --- .../test_trl_grpo_pinned_symbols.py | 20 +++---- unsloth/models/rl_replacements.py | 57 ------------------- 2 files changed, 8 insertions(+), 69 deletions(-) diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index 8f1435ba5f..4c7dcc4234 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -551,12 +551,12 @@ def test_trl_grpo_source_inference_mode_unwrap(tag: str): @pytest.mark.parametrize("tag", TRL_TAGS) def test_trl_kto_get_batch_logps_signature(tag: str): - """KTO log-prob computation must stay patchable. Through TRL 1.x the - target was KTOTrainer.get_batch_logps; TRL 1.x dropped it and moved the - math into _compute_logps / compute_ref_log_probs calling - selective_log_softmax. unsloth/models/rl_replacements.py patches BOTH - shapes (kto_trainer_get_batch_logps + kto_trainer_align_completion_logps), - so we require EITHER form to exist wherever KTOTrainer lives.""" + """TRL 0.27+ moved KTOTrainer to trl.experimental.kto and the + canonical kto_trainer.py shrank to a thin re-export wrapper. The + real `get_batch_logps` lives at trl/experimental/kto/kto_trainer.py. + Unsloth's MRO walk in models/rl.py:592-708 already follows + trl.experimental.* parents, so either path is fine — we just + require the symbol to exist SOMEWHERE.""" candidates = [ "trl/trainer/kto_trainer.py", "trl/experimental/kto/kto_trainer.py", @@ -566,15 +566,11 @@ def test_trl_kto_get_batch_logps_signature(tag: str): src = fetch_text("huggingface/trl", tag, path) if src is None: continue - # Legacy: explicit get_batch_logps method. if has_def(src, "get_batch_logps", "func"): return - # TRL 1.x: refactored into _compute_logps + selective_log_softmax. - if has_def(src, "_compute_logps", "func") and "selective_log_softmax" in src: - return pytest.fail( - f"{tag}: KTO log-prob computation not found in any of {candidates}; " - f"unsloth/models/rl_replacements.py KTO rewrite silently skipped" + f"{tag}: KTOTrainer.get_batch_logps not found in any of {candidates}; " + f"unsloth/models/rl_replacements.py:1675 rewrite silently skipped" ) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index c1d92a31c4..31d54675c9 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1962,63 +1962,6 @@ def kto_trainer_get_batch_logps(function_name, function): RL_FUNCTIONS["kto_trainer"].append(kto_trainer_get_batch_logps) -# TRL 1.x dropped KTOTrainer.get_batch_logps and moved the log-prob math into -# _compute_logps / compute_ref_log_probs / _compute_kl_logps, which call -# selective_log_softmax on completion-only tokens. Same truncation hazard as -# above, so clamp logits/ids/mask to the shorter seq length (no-op when equal). -_KTO_COMPLETION_RE = re.compile( - r"(?P[ \t]*)shift_logits = completion_logits\[:, :-1, :\]\.contiguous\(\)\n" - r"(?P=ws)per_token_logps = selective_log_softmax\(\s*shift_logits,\s*" - r"(?P\w+)\[\"completion_input_ids\"\]\[:, 1:\]\.contiguous\(\)\s*\)\n" - r"(?P=ws)per_token_logps\[(?P=var)\[\"completion_mask\"\]\[:, 1:\] == 0\] = 0\.0" -) -_KTO_KL_RE = re.compile( - r"(?P[ \t]*)shift_KL_logits = KL_logits\[:, :-1, :\]\.contiguous\(\)\n" - r"(?P=ws)KL_per_token_logps = selective_log_softmax\(\s*shift_KL_logits,\s*" - r"(?P\w+)\[\"KL_completion_input_ids\"\]\[:, 1:\]\.contiguous\(\)\s*\)\n" - r"(?P=ws)KL_per_token_logps\[(?P=var)\[\"KL_completion_mask\"\]\[:, 1:\] == 0\] = 0\.0" -) - - -def _kto_completion_repl(m): - ws, var = m.group("ws"), m.group("var") - return ( - f"{ws}shift_logits = completion_logits[:, :-1, :].contiguous()\n" - f"{ws}# Unsloth: clamp logits/ids/mask to shorter seq len (model may truncate input_ids)\n" - f'{ws}_uns_ids = {var}["completion_input_ids"][:, 1:].contiguous()\n' - f"{ws}_uns_n = min(shift_logits.shape[1], _uns_ids.shape[1])\n" - f"{ws}per_token_logps = selective_log_softmax(shift_logits[:, :_uns_n], _uns_ids[:, :_uns_n])\n" - f'{ws}per_token_logps[{var}["completion_mask"][:, 1:][:, :_uns_n] == 0] = 0.0' - ) - - -def _kto_kl_repl(m): - ws, var = m.group("ws"), m.group("var") - return ( - f"{ws}shift_KL_logits = KL_logits[:, :-1, :].contiguous()\n" - f"{ws}# Unsloth: clamp logits/ids/mask to shorter seq len (model may truncate input_ids)\n" - f'{ws}_uns_kl_ids = {var}["KL_completion_input_ids"][:, 1:].contiguous()\n' - f"{ws}_uns_kl_n = min(shift_KL_logits.shape[1], _uns_kl_ids.shape[1])\n" - f"{ws}KL_per_token_logps = selective_log_softmax(shift_KL_logits[:, :_uns_kl_n], _uns_kl_ids[:, :_uns_kl_n])\n" - f'{ws}KL_per_token_logps[{var}["KL_completion_mask"][:, 1:][:, :_uns_kl_n] == 0] = 0.0' - ) - - -def kto_trainer_align_completion_logps(function_name, function): - if function_name not in ( - "_compute_logps", - "compute_ref_log_probs", - "_compute_kl_logps", - ): - return function - function = _KTO_COMPLETION_RE.sub(_kto_completion_repl, function) - function = _KTO_KL_RE.sub(_kto_kl_repl, function) - return function - - -RL_FUNCTIONS["kto_trainer"].append(kto_trainer_align_completion_logps) - - # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L356 # TRL warns if batch size is not a multiple of num_generations -> fix this. def grpo_trainer_fix_batch_size(RLTrainer_source, RLConfig_source): From 63dc27f76e00128807b400f91bbe96e37427e64f Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Thu, 4 Jun 2026 15:38:45 +0800 Subject: [PATCH 5/7] fix(studio): disable mlx gc for none (#5991) --- studio/backend/core/training/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index a825321597..e4ef95a375 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1362,7 +1362,7 @@ def _run_mlx_training(event_queue, stop_queue, config): gc_setting = config.get("gradient_checkpointing", "mlx") if isinstance(gc_setting, str): use_grad_checkpoint = ( - gc_setting if gc_setting.lower() not in ("false", "") else False + gc_setting if gc_setting.lower() not in ("false", "none", "") else False ) else: use_grad_checkpoint = gc_setting From 0425a3c0a1d69e0a99b1436249d012e9703198fd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 4 Jun 2026 00:39:29 -0700 Subject: [PATCH 6/7] Normalize shell scripts to LF in .gitattributes (#5997) Shell scripts are stored as LF in git, but without an eol rule a Windows clone with core.autocrlf=true checks them out as CRLF. The trailing \r then breaks them when run in WSL/Linux -- e.g. `set -e` becomes `set -e\r` and dash/sh aborts with "set: Illegal option -". This bites developers who clone on Windows and run the repo's *.sh directly in WSL, increasingly common with the AMD Strix Halo ROCm-on-WSL support. Add `*.sh text eol=lf` so every shell script always checks out with LF regardless of the contributor's platform or core.autocrlf setting. All tracked *.sh use Unix shebangs; none need CRLF. PowerShell/batch scripts are left untouched -- they tolerate LF and are unaffected by this bug. Verified with `git ls-files --eol`: every *.sh now resolves to i/lf w/lf attr/text eol=lf. Co-authored-by: Claude Opus 4.8 --- .gitattributes | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.gitattributes b/.gitattributes index 264fdfd02f..75fba5d6ab 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,2 +1,7 @@ # Normalize Python files to LF line endings *.py text eol=lf + +# Always check out shell scripts with LF endings. Without this rule a Windows +# clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks +# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -"). +*.sh text eol=lf From 4c06c1dcc771f3d4f2aecc8d3ab77f2e01aa9107 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 4 Jun 2026 00:56:53 -0700 Subject: [PATCH 7/7] Studio: enable audio input for Gemma 4 GGUFs; default chat model to Qwen3.5-4B-MTP (#6000) * Studio: enable audio input for Gemma 4 GGUF models Audio file upload was disabled for Gemma 4 vision+audio GGUFs (e.g. gemma-4-12b-it-GGUF) even though their mmproj carries an audio encoder (clip.has_audio_encoder, gemma4ua). Two causes: - Audio-input detection only matched Gemma 3n's ; Gemma 4 uses <|audio|>, so audio_vlm was never detected. - The GGUF load/status responses hardcoded has_audio_input=False, so the flag was dropped even when audio_vlm was detected (affected Gemma 3n GGUFs too). Changes: - Recognize <|audio|> alongside in the llama-server token probe and the tokenizer-config pattern. - Read clip.has_audio_encoder from the mmproj as an independent, model-agnostic signal (read_mmproj_audio_capability). - Emit the computed has_audio_input on the GGUF load/status responses. - Tests for the new pattern and the mmproj reader. * Studio: default chat model and dataset helper to Qwen3.5-4B-MTP Switch the auto-loaded chat default and the dataset-analysis helper GGUF from gemma-4-E2B-it to unsloth/Qwen3.5-4B-MTP-GGUF (UD-Q4_K_XL). * [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> --- studio/backend/core/inference/llama_cpp.py | 37 ++++++++- studio/backend/routes/inference.py | 6 +- .../tests/test_audio_token_detection.py | 46 +++++++++++ studio/backend/tests/test_gguf_metadata.py | 67 +++++++++++++++- studio/backend/utils/datasets/llm_assist.py | 2 +- studio/backend/utils/models/gguf_metadata.py | 78 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 3 +- .../src/features/chat/api/chat-adapter.ts | 16 ++-- 8 files changed, 240 insertions(+), 15 deletions(-) create mode 100644 studio/backend/tests/test_audio_token_detection.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7bcf02dc35..0f23549138 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -697,6 +697,9 @@ class LlamaCppBackend: self._is_audio: bool = False self._audio_type: Optional[str] = None self._audio_probed: bool = False + # Audio INPUT capability (distinct from _is_audio, which is TTS output). + self._has_audio_input: bool = False + self._mmproj_has_audio: bool = False # clip.has_audio_encoder, set at load # Monotonic timestamp set in _kill_process; read by load_model # to decide whether to wait for the VRAM reclaim to finish. self._last_kill_monotonic: float = 0.0 @@ -2782,6 +2785,12 @@ class LlamaCppBackend: if not self._healthy: return False self._audio_type = detected + # Re-derive after a retried probe (_mmproj_has_audio persists). + from utils.models.model_config import is_audio_input_type + + self._has_audio_input = bool( + is_audio_input_type(self._audio_type) + ) or bool(self._mmproj_has_audio) if not self._healthy: return False return True @@ -3095,6 +3104,21 @@ class LlamaCppBackend: "image input will be disabled for this session" ) + # Audio input straight from the mmproj (clip.has_audio_encoder), + # independent of token names. + self._mmproj_has_audio = False + if launch_mmproj_path: + try: + from utils.models.gguf_metadata import ( + read_mmproj_audio_capability, + ) + + self._mmproj_has_audio = bool( + read_mmproj_audio_capability(launch_mmproj_path) + ) + except Exception as e: + logger.debug(f"mmproj audio-capability read failed: {e}") + cmd = [ binary, "-m", @@ -3527,6 +3551,7 @@ class LlamaCppBackend: self._is_audio = False self._audio_type = None self._audio_probed = False + self._has_audio_input = False try: detected = self._detect_audio_type_strict() self._audio_probed = True @@ -3558,6 +3583,13 @@ class LlamaCppBackend: return False self._audio_type = detected + # Audio input = token probe (audio_vlm/whisper) OR mmproj audio encoder. + from utils.models.model_config import is_audio_input_type + + self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool( + self._mmproj_has_audio + ) + if not self._healthy: return False return True @@ -3901,6 +3933,8 @@ class LlamaCppBackend: self._is_audio = False self._audio_type = None self._audio_probed = False + self._has_audio_input = False + self._mmproj_has_audio = False self._port = None self._healthy = False self._context_length = None @@ -5591,7 +5625,8 @@ class LlamaCppBackend: return "csm" if len(_tok("<|startoftranscript|>")) == 1: return "whisper" - if len(_tok("")) == 1: + # Gemma 3n: ; Gemma 4: <|audio|> (not csm's <|AUDIO|>). + if len(_tok("")) == 1 or len(_tok("<|audio|>")) == 1: return "audio_vlm" if ( len(_tok("<|bicodec_semantic_0|>")) == 1 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2964a0801f..15d73405cb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -794,7 +794,7 @@ async def load_model( is_gguf = True, is_audio = _gguf_is_audio, audio_type = _gguf_audio, - has_audio_input = False, + has_audio_input = getattr(llama_backend, "_has_audio_input", False), inference = inference_config, requires_trust_remote_code = bool( inference_config.get("trust_remote_code", False) @@ -1044,7 +1044,7 @@ async def load_model( is_gguf = True, is_audio = _gguf_is_audio, audio_type = _gguf_audio, - has_audio_input = False, + has_audio_input = llama_backend._has_audio_input, inference = inference_config, requires_trust_remote_code = bool( inference_config.get("trust_remote_code", False) @@ -1531,7 +1531,7 @@ async def get_status( gguf_variant = llama_backend.hf_variant, is_audio = getattr(llama_backend, "_is_audio", False), audio_type = _audio_type, - has_audio_input = False, + has_audio_input = getattr(llama_backend, "_has_audio_input", False), loading = [], loaded = [_display_model_id] if _display_model_id else [], inference = _inference_cfg, diff --git a/studio/backend/tests/test_audio_token_detection.py b/studio/backend/tests/test_audio_token_detection.py new file mode 100644 index 0000000000..a3ea7c89a7 --- /dev/null +++ b/studio/backend/tests/test_audio_token_detection.py @@ -0,0 +1,46 @@ +# 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 tokenizer-based audio_type detection patterns, covering both +Gemma 3n () and Gemma 4 (<|audio|>) audio-input tokens.""" + +from __future__ import annotations + +from utils.models.model_config import _AUDIO_TOKEN_PATTERNS, is_audio_input_type + + +def _classify(tokens: list[str]) -> str | None: + """Mirror _detect_audio_from_tokenizer._check_token_patterns: first match + in dict order wins.""" + for audio_type, check in _AUDIO_TOKEN_PATTERNS.items(): + if check(tokens): + return audio_type + return None + + +def test_gemma3n_audio_soft_token_is_audio_vlm(): + assert ( + _classify(["", "", ""]) == "audio_vlm" + ) + + +def test_gemma4_pipe_audio_token_is_audio_vlm(): + # Gemma 4 uses <|audio|> (and <|image|>) instead of *_soft_token. + assert _classify(["", "<|image|>", "<|audio|>"]) == "audio_vlm" + + +def test_csm_uppercase_audio_not_classified_as_audio_vlm(): + # csm uses uppercase <|AUDIO|> + <|audio_eos|>; must stay csm, not audio_vlm. + tokens = ["<|AUDIO|>", "<|audio_eos|>"] + assert _classify(tokens) == "csm" + + +def test_audio_vlm_and_whisper_accept_audio_input(): + assert is_audio_input_type("audio_vlm") is True + assert is_audio_input_type("whisper") is True + assert is_audio_input_type("snac") is False + assert is_audio_input_type(None) is False + + +def test_non_audio_tokens_classify_none(): + assert _classify(["", "", ""]) is None diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index cf1a17347f..e5040e306c 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -14,6 +14,7 @@ from utils.models.gguf_metadata import ( is_mmproj_by_metadata, pairing_score, read_gguf_general_metadata, + read_mmproj_audio_capability, ) @@ -21,6 +22,7 @@ _GGUF_MAGIC = 0x46554747 _VTYPE_STRING = 8 _VTYPE_UINT32 = 4 _VTYPE_ARRAY = 9 +_VTYPE_BOOL = 7 def _enc_string(s: str) -> bytes: @@ -38,6 +40,14 @@ def _enc_kv_uint32(key: str, value: int) -> bytes: ) +def _enc_kv_bool(key: str, value: bool) -> 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) + extra_bools = extra_bools or {} + kv_count = ( + len(general_strings) + + len(extra_uint32) + + len(extra_string_arrays) + + len(extra_bools) + ) body = b"" for k, v in general_strings.items(): body += _enc_kv_string(k, v) @@ -65,6 +82,8 @@ def _write_synthetic_gguf( body += _enc_kv_uint32(k, v) for k, v in extra_string_arrays.items(): body += _enc_kv_string_array(k, v) + for k, v in extra_bools.items(): + body += _enc_kv_bool(k, v) header = struct.pack( " Optional[_CacheKey]: try: @@ -193,6 +197,80 @@ def _skip_gguf_value(f, vtype: int) -> bool: return True +def _parse_gguf_bool(path: str, wanted_key: str) -> Optional[bool]: + """Bool value of ``wanted_key`` (GGUF vtype 7), or ``None`` if absent / + unreadable. Mirrors ``_parse_gguf_header`` for a single bool key.""" + 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(" Optional[bool]: + """Cached single-bool-key read, keyed by (path, mtime, size, wanted_key).""" + fkey = _cache_key(path) + if fkey is None: + return None + ckey = (fkey, wanted_key) + with _CACHE_LOCK: + if ckey in _BOOL_CACHE: + return _BOOL_CACHE[ckey] + result = _parse_gguf_bool(path, wanted_key) + with _CACHE_LOCK: + while len(_BOOL_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _BOOL_CACHE.pop(next(iter(_BOOL_CACHE))) + except StopIteration: + break + _BOOL_CACHE[ckey] = result + return result + + +def read_mmproj_audio_capability(path: str) -> Optional[bool]: + """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's gemma4ua): + ``True``/``False`` if present, ``None`` if absent / unreadable. Flags + audio-input models independently of tokenizer token names.""" + return _read_gguf_bool(path, "clip.has_audio_encoder") + + 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: diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index dc34444ccb..b488a19953 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -814,7 +814,8 @@ _audio_detection_cache: Dict[str, Optional[str]] = {} _AUDIO_TOKEN_PATTERNS = { "csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens, "whisper": lambda tokens: "<|startoftranscript|>" in tokens, - "audio_vlm": lambda tokens: "" in tokens, + # Gemma 3n: ; Gemma 4: <|audio|> (not csm's <|AUDIO|>). + "audio_vlm": lambda tokens: "" in tokens or "<|audio|>" in tokens, "bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens), "dac": lambda tokens: ( "<|audio_start|>" in tokens diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 1fcbe7ff70..9a7f95df87 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1149,13 +1149,13 @@ async function autoLoadSmallestModel(): Promise<{ toast("Downloading a small model…", { id: toastId, description: - "No downloaded models found. Fetching Gemma-4-E2B-it (UD-Q4_K_XL).", + "No downloaded models found. Fetching Qwen3.5-4B-MTP (UD-Q4_K_XL).", duration: 30000, }); try { if ( !(await canAutoLoad({ - model_path: "unsloth/gemma-4-E2B-it-GGUF", + model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", max_seq_length: 0, is_lora: false, gguf_variant: "UD-Q4_K_XL", @@ -1166,7 +1166,7 @@ async function autoLoadSmallestModel(): Promise<{ } loadAttempts += 1; const loadResp = await loadModel({ - model_path: "unsloth/gemma-4-E2B-it-GGUF", + model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", hf_token: hfToken, max_seq_length: 0, load_in_4bit: true, @@ -1176,7 +1176,7 @@ async function autoLoadSmallestModel(): Promise<{ }); useChatRuntimeStore .getState() - .setCheckpoint("unsloth/gemma-4-E2B-it-GGUF", "UD-Q4_K_XL"); + .setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL"); const store = useChatRuntimeStore.getState(); store.setModelRequiresTrustRemoteCode( loadResp.requires_trust_remote_code ?? false, @@ -1186,13 +1186,13 @@ async function autoLoadSmallestModel(): Promise<{ maxTokens: loadResp.context_length ?? 131072, }); const defaultModel: ChatModelSummary = { - id: "unsloth/gemma-4-E2B-it-GGUF", - name: loadResp.display_name ?? "gemma-4-E2B-it-GGUF", + id: "unsloth/Qwen3.5-4B-MTP-GGUF", + name: loadResp.display_name ?? "Qwen3.5-4B-MTP-GGUF", isVision: loadResp.is_vision ?? false, isLora: false, isGguf: true, }; - if (!store.models.some((m) => m.id === "unsloth/gemma-4-E2B-it-GGUF")) { + if (!store.models.some((m) => m.id === "unsloth/Qwen3.5-4B-MTP-GGUF")) { store.setModels([...store.models, defaultModel]); } useChatRuntimeStore.setState({ @@ -1212,7 +1212,7 @@ async function autoLoadSmallestModel(): Promise<{ chatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), }); - toast.success("Loaded Gemma-4-E2B-it (UD-Q4_K_XL)", { id: toastId }); + toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId }); return { loaded: true, blockedByTrustRemoteCode: false }; } catch { toast.dismiss(toastId);