diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index a8fea07bf7..13ecb992bf 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -2,8 +2,8 @@ "_comment": "Per-model-family inference parameter defaults. Sources: (1) Ollama params blobs, (2) Existing Unsloth Studio YAML configs. Patterns ordered longest-match-first.", "families": { "qwen3.5": { - "temperature": 1.0, - "top_p": 0.95, + "temperature": 0.7, + "top_p": 0.8, "top_k": 20, "min_p": 0.0, "repetition_penalty": 1.0 diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index c8e23deb09..d5e9ca2e97 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -3,7 +3,7 @@ """Default model lists for inference, split by platform.""" -import sys +import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ "unsloth/Llama-3.2-1B-Instruct-GGUF", @@ -25,6 +25,7 @@ DEFAULT_MODELS_STANDARD = [ def get_default_models() -> list[str]: - if sys.platform == "darwin": + hw.get_device() # ensure detect_hardware() has run + if hw.CHAT_ONLY: return list(DEFAULT_MODELS_GGUF) return list(DEFAULT_MODELS_STANDARD) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3a137b8b63..08126dedf7 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -49,6 +49,9 @@ class LlamaCppBackend: self._context_length: Optional[int] = None self._chat_template: Optional[str] = None self._supports_reasoning: bool = False + self._supports_tools: bool = False + self._cache_type_kv: Optional[str] = None + self._reasoning_default: bool = True self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -96,6 +99,18 @@ class LlamaCppBackend: def supports_reasoning(self) -> bool: return self._supports_reasoning + @property + def reasoning_default(self) -> bool: + return self._reasoning_default + + @property + def supports_tools(self) -> bool: + return self._supports_tools + + @property + def cache_type_kv(self) -> Optional[str]: + return self._cache_type_kv + # ── Binary discovery ────────────────────────────────────────── @staticmethod @@ -492,6 +507,18 @@ class LlamaCppBackend: logger.info( "GGUF metadata: model supports reasoning (DeepSeek thinking)" ) + # Detect tool calling support from chat template + tool_markers = [ + "{%- if tools %}", + "{% if tools %}", + '"role" == "tool"', + "'role' == 'tool'", + 'message.role == "tool"', + "message.role == 'tool'", + ] + if any(marker in tpl for marker in tool_markers): + self._supports_tools = True + logger.info("GGUF metadata: model supports tool calling") except Exception as e: logger.warning(f"Failed to read GGUF metadata: {e}") @@ -723,6 +750,7 @@ class LlamaCppBackend: is_vision: bool = False, n_ctx: int = 4096, chat_template_override: Optional[str] = None, + cache_type_kv: Optional[str] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused ) -> bool: @@ -828,6 +856,27 @@ class LlamaCppBackend: # Always enable Jinja chat template rendering for proper template support cmd.extend(["--jinja"]) + # KV cache data type + _valid_cache_types = { + "f16", + "bf16", + "q8_0", + "q4_0", + "q4_1", + "q5_0", + "q5_1", + "iq4_nl", + "f32", + } + if cache_type_kv and cache_type_kv in _valid_cache_types: + cmd.extend( + ["--cache-type-k", cache_type_kv, "--cache-type-v", cache_type_kv] + ) + self._cache_type_kv = cache_type_kv + logger.info(f"KV cache type: {cache_type_kv}") + else: + self._cache_type_kv = None + # Apply custom chat template override if provided if chat_template_override: import tempfile @@ -845,15 +894,31 @@ class LlamaCppBackend: f"Using custom chat template file: {self._chat_template_file.name}" ) - # For reasoning models, default to thinking ON (user can toggle per-request) + # For reasoning models, set default thinking mode. + # Qwen3.5 small models (0.8B, 2B, 4B, 9B) disable thinking by default + # per Qwen's recommendation. Larger models default to thinking ON. if self._supports_reasoning: + import re + + thinking_default = True + mid = (model_identifier or "").lower() + if "qwen3.5" in mid: + # Extract size like "0.8b", "4b", "35b" etc. + size_match = re.search(r"(\d+\.?\d*)\s*b", mid) + if size_match: + size_val = float(size_match.group(1)) + if size_val <= 2: + thinking_default = False + self._reasoning_default = thinking_default cmd.extend( [ "--chat-template-kwargs", - json.dumps({"enable_thinking": True}), + json.dumps({"enable_thinking": thinking_default}), ] ) - logger.info("Reasoning model: enabled enable_thinking=true by default") + logger.info( + f"Reasoning model: enable_thinking={thinking_default} by default" + ) if mmproj_path: if not Path(mmproj_path).is_file(): @@ -888,9 +953,28 @@ class LlamaCppBackend: env["PATH"] = ";".join(path_dirs) + ";" + existing_path else: # Linux: set LD_LIBRARY_PATH for shared libs next to the binary + # and CUDA runtime libs (libcudart, libcublas, etc.) + import platform + + lib_dirs = [binary_dir] + _arch = platform.machine() # x86_64, aarch64, etc. + for cuda_lib in [ + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{_arch}-linux/lib", + # Fallback CUDA compat paths (e.g. binary built with + # CUDA 12 on a system where default /usr/local/cuda + # points to CUDA 13+). + "/usr/local/cuda-12/lib64", + "/usr/local/cuda-12.8/lib64", + f"/usr/local/cuda-12/targets/{_arch}-linux/lib", + f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", + ]: + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) existing_ld = env.get("LD_LIBRARY_PATH", "") + new_ld = ":".join(lib_dirs) env["LD_LIBRARY_PATH"] = ( - f"{binary_dir}:{existing_ld}" if existing_ld else binary_dir + f"{new_ld}:{existing_ld}" if existing_ld else new_ld ) # Pin to selected GPU(s) via CUDA_VISIBLE_DEVICES @@ -952,6 +1036,8 @@ class LlamaCppBackend: self._context_length = None self._chat_template = None self._supports_reasoning = False + self._supports_tools = False + self._cache_type_kv = None # Clean up temp chat template file if hasattr(self, "_chat_template_file") and self._chat_template_file: try: @@ -1071,6 +1157,70 @@ class LlamaCppBackend: # ── Message building (OpenAI format) ────────────────────────── + @staticmethod + def _parse_tool_calls_from_text(content: str) -> list[dict]: + """ + Parse tool calls from XML markup in content text. + + Handles formats like: + {"name":"web_search","arguments":{"query":"..."}} + ... + Closing tag is optional (models sometimes omit it). + """ + import re + + tool_calls = [] + # Pattern 1: JSON inside tags (closing tag optional) + for match in re.finditer( + r"\s*(\{.*?\})\s*(?:)?", content, re.DOTALL + ): + try: + obj = json.loads(match.group(1)) + tc = { + "id": f"call_{len(tool_calls)}", + "type": "function", + "function": { + "name": obj.get("name", ""), + "arguments": obj.get("arguments", {}), + }, + } + if isinstance(tc["function"]["arguments"], dict): + tc["function"]["arguments"] = json.dumps( + tc["function"]["arguments"] + ) + tool_calls.append(tc) + except (json.JSONDecodeError, ValueError): + pass + + # Pattern 2: XML-style value + # Closing optional + if not tool_calls: + for match in re.finditer( + r"\s*(.*?)\s*(?:)?", + content, + re.DOTALL, + ): + func_name = match.group(1) + params_text = match.group(2) + arguments = {} + for param_match in re.finditer( + r"\s*(.*?)\s*", + params_text, + re.DOTALL, + ): + arguments[param_match.group(1)] = param_match.group(2) + tc = { + "id": f"call_{len(tool_calls)}", + "type": "function", + "function": { + "name": func_name, + "arguments": json.dumps(arguments), + }, + } + tool_calls.append(tc) + + return tool_calls + @staticmethod def _build_openai_messages( messages: list[dict], @@ -1166,6 +1316,8 @@ class LlamaCppBackend: ) buffer = "" + has_content_tokens = False + reasoning_text = "" for raw_chunk in response.iter_text(): if cancel_event is not None and cancel_event.is_set(): break @@ -1179,8 +1331,17 @@ class LlamaCppBackend: continue if line == "data: [DONE]": if in_thinking: - cumulative += "" - yield cumulative + if has_content_tokens: + # Real thinking + content: close the tag + cumulative += "" + yield cumulative + else: + # Only reasoning_content, no content tokens: + # the model put its entire reply in reasoning + # (e.g. Qwen3 always-think mode). Show it + # as the main response, not as a thinking block. + cumulative = reasoning_text + yield cumulative return if not line.startswith("data: "): continue @@ -1196,6 +1357,7 @@ class LlamaCppBackend: # Wrap in tags for the frontend parser reasoning = delta.get("reasoning_content", "") if reasoning: + reasoning_text += reasoning if not in_thinking: cumulative += "" in_thinking = True @@ -1204,6 +1366,7 @@ class LlamaCppBackend: token = delta.get("content", "") if token: + has_content_tokens = True if in_thinking: cumulative += "" in_thinking = False @@ -1221,6 +1384,254 @@ class LlamaCppBackend: return raise + # ── Tool-calling agentic loop ────────────────────────────── + + def generate_chat_completion_with_tools( + self, + messages: list[dict], + tools: list[dict], + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_tokens: Optional[int] = None, + repetition_penalty: float = 1.0, + stop: Optional[list[str]] = None, + cancel_event: Optional[threading.Event] = None, + enable_thinking: Optional[bool] = None, + max_tool_iterations: int = 5, + ) -> Generator[dict, None, None]: + """ + Agentic loop: let the model call tools, execute them, and continue. + + Yields dicts with: + {"type": "status", "text": "Searching: ..."} -- tool status updates + {"type": "content", "text": "token"} -- streamed content tokens (cumulative) + {"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative) + """ + from core.inference.tools import execute_tool + + if not self.is_loaded: + raise RuntimeError("llama-server is not loaded") + + conversation = list(messages) + url = f"{self.base_url}/v1/chat/completions" + + for iteration in range(max_tool_iterations): + if cancel_event is not None and cancel_event.is_set(): + return + + # Build payload for non-streaming tool detection pass + payload = { + "messages": conversation, + "stream": False, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k if top_k >= 0 else 0, + "min_p": min_p, + "repeat_penalty": repetition_penalty, + "tools": tools, + "tool_choice": "auto", + } + if self._supports_reasoning and enable_thinking is not None: + payload["chat_template_kwargs"] = {"enable_thinking": enable_thinking} + if max_tokens is not None: + payload["max_tokens"] = max_tokens + if stop: + payload["stop"] = stop + + try: + with httpx.Client(timeout = None) as client: + resp = client.post(url, json = payload) + if resp.status_code != 200: + raise RuntimeError( + f"llama-server returned {resp.status_code}: {resp.text}" + ) + data = resp.json() + except httpx.ConnectError: + raise RuntimeError("Lost connection to llama-server") + + choices = data.get("choices", []) + if not choices: + return + + choice = choices[0] + finish_reason = choice.get("finish_reason", "") + message = choice.get("message", {}) + + # If model wants to call tools + tool_calls = message.get("tool_calls") + + # Fallback: detect tool calls embedded as XML/text in content + # Some models output XML instead of structured tool_calls + content_text = message.get("content", "") or "" + if not tool_calls and "" in content_text: + tool_calls = self._parse_tool_calls_from_text(content_text) + if tool_calls: + # Strip the tool call markup from content + import re + + content_text = re.sub( + r".*?(?:|$)", + "", + content_text, + flags = re.DOTALL, + ).strip() + logger.info( + f"Parsed {len(tool_calls)} tool call(s) from content text" + ) + + if finish_reason == "tool_calls" or (tool_calls and len(tool_calls) > 0): + # Append the assistant message with tool_calls to conversation + assistant_msg = {"role": "assistant", "content": content_text} + if tool_calls: + assistant_msg["tool_calls"] = tool_calls + conversation.append(assistant_msg) + + # Execute each tool call + for tc in tool_calls or []: + func = tc.get("function", {}) + tool_name = func.get("name", "") + raw_args = func.get("arguments", {}) + + # Handle arguments as either string or dict + if isinstance(raw_args, str): + try: + arguments = json.loads(raw_args) + except (json.JSONDecodeError, ValueError): + arguments = {"query": raw_args} + else: + arguments = raw_args + + # Yield status update + query_text = arguments.get("query", tool_name) + yield {"type": "status", "text": f"Searching: {query_text}"} + + # Execute the tool + result = execute_tool(tool_name, arguments) + + # Append tool result to conversation + tool_msg = { + "role": "tool", + "name": tool_name, + "content": result, + } + tool_call_id = tc.get("id") + if tool_call_id: + tool_msg["tool_call_id"] = tool_call_id + conversation.append(tool_msg) + + # Continue the loop to let model respond with context + continue + + # No tool calls -- model answered directly. + # If no tools were executed at all, just yield the content + # from this response instead of making a redundant second request. + if iteration == 0 and content_text: + yield {"type": "status", "text": ""} + yield {"type": "content", "text": content_text} + return + + # Tools were called in previous iterations; do a final + # streaming pass so the model can synthesize a response + # incorporating the tool results. + break + + # Clear status + yield {"type": "status", "text": ""} + + # Final streaming pass with the full conversation context + stream_payload = { + "messages": conversation, + "stream": True, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k if top_k >= 0 else 0, + "min_p": min_p, + "repeat_penalty": repetition_penalty, + } + if self._supports_reasoning and enable_thinking is not None: + stream_payload["chat_template_kwargs"] = { + "enable_thinking": enable_thinking + } + if max_tokens is not None: + stream_payload["max_tokens"] = max_tokens + if stop: + stream_payload["stop"] = stop + + cumulative = "" + in_thinking = False + has_content_tokens = False + reasoning_text = "" + + try: + with httpx.Client(timeout = None) as client: + with client.stream("POST", url, json = stream_payload) as response: + if response.status_code != 200: + error_body = response.read().decode() + raise RuntimeError( + f"llama-server returned {response.status_code}: {error_body}" + ) + + buffer = "" + for raw_chunk in response.iter_text(): + if cancel_event is not None and cancel_event.is_set(): + break + + buffer += raw_chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.strip() + + if not line: + continue + if line == "data: [DONE]": + if in_thinking: + if has_content_tokens: + cumulative += "" + yield {"type": "content", "text": cumulative} + else: + cumulative = reasoning_text + yield {"type": "content", "text": cumulative} + return + if not line.startswith("data: "): + continue + + try: + chunk_data = json.loads(line[6:]) + choices = chunk_data.get("choices", []) + if choices: + delta = choices[0].get("delta", {}) + + reasoning = delta.get("reasoning_content", "") + if reasoning: + reasoning_text += reasoning + if not in_thinking: + cumulative += "" + in_thinking = True + cumulative += reasoning + yield {"type": "content", "text": cumulative} + + token = delta.get("content", "") + if token: + has_content_tokens = True + if in_thinking: + cumulative += "" + in_thinking = False + cumulative += token + yield {"type": "content", "text": cumulative} + except json.JSONDecodeError: + logger.debug( + f"Skipping malformed SSE line: {line[:100]}" + ) + + except httpx.ConnectError: + raise RuntimeError("Lost connection to llama-server") + except Exception as e: + if cancel_event is not None and cancel_event.is_set(): + return + raise + # ── TTS support ──────────────────────────────────────────── def detect_audio_type(self) -> Optional[str]: diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py new file mode 100644 index 0000000000..b633d377f2 --- /dev/null +++ b/studio/backend/core/inference/tools.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Tool definitions and executors for LLM tool calling. + +Currently supports web search via DuckDuckGo (ddgs package, no API key needed). +""" + +WEB_SEARCH_TOOL = { + "type": "function", + "function": { + "name": "web_search", + "description": "Search the web for current information, recent events, or facts you are uncertain about.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query", + } + }, + "required": ["query"], + }, + }, +} + +ALL_TOOLS = [WEB_SEARCH_TOOL] + + +def execute_tool(name: str, arguments: dict) -> str: + """Execute a tool by name with the given arguments. Returns result as a string.""" + if name == "web_search": + return _web_search(arguments.get("query", "")) + return f"Unknown tool: {name}" + + +def _web_search(query: str, max_results: int = 5) -> str: + """Search the web using DuckDuckGo and return formatted results.""" + if not query.strip(): + return "No query provided." + try: + from ddgs import DDGS + + results = DDGS().text(query, max_results = max_results) + if not results: + return "No results found." + parts = [] + for r in results: + parts.append( + f"Title: {r.get('title', '')}\n" + f"URL: {r.get('href', '')}\n" + f"Snippet: {r.get('body', '')}" + ) + return "\n\n---\n\n".join(parts) + except Exception as e: + return f"Search failed: {e}" diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index ecc11c30d5..330fb74747 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -2345,6 +2345,9 @@ class UnslothTrainer: status_message = f"Streamed {len(dataset)} rows from HuggingFace" ) else: + self._update_progress( + status_message = f"Downloading dataset: {dataset_source}..." + ) dataset = load_dataset(**load_kwargs) # Check if stopped during dataset loading @@ -2352,11 +2355,12 @@ class UnslothTrainer: logger.info("Stopped during dataset loading\n") return None + n_rows = len(dataset) if hasattr(dataset, "__len__") else 0 self._update_progress( - status_message = f"Loaded dataset from HuggingFace: {dataset_source}" + status_message = f"Downloaded {dataset_source} ({n_rows:,} rows)" ) logger.info( - f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n" + f"Loaded dataset from Hugging Face: {dataset_source} ({n_rows:,} rows)\n" ) # Resolve eval split from a separate HF split (explicit or auto-detected) @@ -2481,10 +2485,15 @@ class UnslothTrainer: self._update_progress(error = error_msg) return None + detected = dataset_info.get("detected_format", "unknown") + final_ds = dataset_info.get("dataset") + final_n = len(final_ds) if hasattr(final_ds, "__len__") else "?" self._update_progress( - status_message = f"Dataset formatted and ready for training" + status_message = f"Dataset ready ({final_n:,} samples, {detected} format)" + ) + logger.info( + f"Dataset formatted successfully ({final_n} samples, {detected})\n" ) - logger.info(f"Dataset formatted successfully\n") # ========== THEN SPLIT ========== if has_separate_eval_source and eval_dataset is not None: diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f0e8087672..c2801227ff 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -138,7 +138,22 @@ def run_training_process( ) return - # ── 1b. On Windows, check Triton availability (must be before import torch) ── + # ── 1b. Set fork start method so dataset.map() can multiprocess ── + # The parent launched us via spawn (clean process), but the compiled + # SFTTrainer checks get_start_method() and disables num_proc if not "fork". + # Linux only: fork is the default start method and is safe here (no CUDA + # context exists yet). macOS defaults to spawn since Python 3.8 because + # fork is unsafe with macOS frameworks (Metal/MPS, CoreFoundation) -- + # do NOT override on macOS. Windows has no fork at all. + if sys.platform == "linux": + import multiprocessing as _mp + + try: + _mp.set_start_method("fork", force = True) + except RuntimeError: + pass # Already set + + # ── 1c. On Windows, check Triton availability (must be before import torch) ── if sys.platform == "win32": try: import triton # noqa: F401 @@ -347,6 +362,31 @@ def run_training_process( ) return + # ── Start tqdm monitor early so it captures download + tokenization bars ── + import threading as _th + + _tqdm_stop = _th.Event() + + def _monitor_tqdm(): + from tqdm.auto import tqdm as _tqdm_cls + + while not _tqdm_stop.is_set(): + for bar in list(getattr(_tqdm_cls, "_instances", set())): + try: + n, total = bar.n or 0, bar.total or 0 + desc = getattr(bar, "desc", "") or "" + if total > 0 and n > 0 and desc: + pct = min(int(n * 100 / total), 100) + _send_status( + event_queue, f"{desc.strip()} {pct}% ({n:,}/{total:,})" + ) + except (AttributeError, ReferenceError): + pass + _tqdm_stop.wait(3) + + _tqdm_thread = _th.Thread(target = _monitor_tqdm, daemon = True) + _tqdm_thread.start() + # ── 4c. Load training model (uses VRAM — dataset already formatted) ── _send_status(event_queue, "Loading model...") success = trainer.load_model( @@ -477,6 +517,8 @@ def run_training_process( lr_scheduler_type = config.get("lr_scheduler_type", "linear"), ) + _tqdm_stop.set() + # Check final state progress = trainer.get_training_progress() if progress.error: diff --git a/studio/backend/main.py b/studio/backend/main.py index 3ab846306e..23c868b535 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -153,6 +153,7 @@ async def health_check(): "timestamp": datetime.now().isoformat(), "service": "Unsloth UI Backend", "device_type": device_type, + "chat_only": _hw_module.CHAT_ONLY, } diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0014ebfdd0..e7ae63adf1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -37,6 +37,10 @@ class LoadRequest(BaseModel): None, description = "Custom Jinja2 chat template to use instead of the model's default", ) + cache_type_kv: Optional[str] = Field( + None, + description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", + ) class UnloadRequest(BaseModel): @@ -128,6 +132,14 @@ class LoadResponse(BaseModel): False, description = "Whether model supports thinking/reasoning mode (enable_thinking)", ) + supports_tools: bool = Field( + False, + description = "Whether model supports tool calling (web search, etc.)", + ) + cache_type_kv: Optional[str] = Field( + None, + description = "KV cache data type for K and V (e.g. 'f16', 'bf16', 'q8_0')", + ) chat_template: Optional[str] = Field( None, description = "Jinja2 chat template string (from GGUF metadata or tokenizer)", @@ -284,6 +296,10 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Enable/disable thinking/reasoning mode for supported models", ) + enable_tools: Optional[bool] = Field( + None, + description = "[x-unsloth] Enable tool calling (web search) for supported models", + ) # ── Streaming response chunks ──────────────────────────────────── diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 7eccaa0faf..da0b4f382b 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -14,3 +14,4 @@ gradio>=4.0.0 huggingface-hub==0.36.2 structlog>=24.1.0 diceware +ddgs diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e93ba1d6d6..5a49a730ca 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -136,6 +136,7 @@ async def load_model( is_vision = config.is_vision, n_ctx = request.max_seq_length, chat_template_override = request.chat_template_override, + cache_type_kv = request.cache_type_kv, ) else: # Local mode: llama-server loads via -m @@ -147,6 +148,7 @@ async def load_model( is_vision = config.is_vision, n_ctx = request.max_seq_length, chat_template_override = request.chat_template_override, + cache_type_kv = request.cache_type_kv, ) if not success: @@ -183,6 +185,8 @@ async def load_model( inference = inference_config, context_length = llama_backend.context_length, supports_reasoning = llama_backend.supports_reasoning, + supports_tools = llama_backend.supports_tools, + cache_type_kv = llama_backend.cache_type_kv, chat_template = llama_backend.chat_template, ) @@ -931,6 +935,129 @@ async def openai_chat_completions( completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) + # ── Tool-calling path (agentic loop) ────────────────── + use_tools = ( + payload.enable_tools and llama_backend.supports_tools and not image_b64 + ) + + if use_tools: + from core.inference.tools import ALL_TOOLS + + def gguf_generate_with_tools(): + return llama_backend.generate_chat_completion_with_tools( + messages = gguf_messages, + tools = ALL_TOOLS, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_tokens = payload.max_tokens, + repetition_penalty = payload.repetition_penalty, + cancel_event = cancel_event, + enable_thinking = payload.enable_thinking, + ) + + _tool_sentinel = object() + + async def gguf_tool_stream(): + try: + first_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(role = "assistant"), + finish_reason = None, + ) + ], + ) + yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" + + # Iterate the synchronous generator in a thread so + # the event loop stays free for disconnect detection. + gen = gguf_generate_with_tools() + prev_text = "" + while True: + if await request.is_disconnected(): + cancel_event.set() + return + + event = await asyncio.to_thread(next, gen, _tool_sentinel) + if event is _tool_sentinel: + break + + if event["type"] == "status": + # Emit tool status as a custom SSE event + status_data = json.dumps( + { + "type": "tool_status", + "content": event["text"], + } + ) + yield f"data: {status_data}\n\n" + continue + + # "content" type -- cumulative text + cumulative = event.get("text", "") + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(content = new_text), + finish_reason = None, + ) + ], + ) + yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + + final_chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = "stop", + ) + ], + ) + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" + yield "data: [DONE]\n\n" + + except asyncio.CancelledError: + cancel_event.set() + raise + except Exception as e: + logger.error( + f"Error during GGUF tool streaming: {e}", exc_info = True + ) + error_chunk = { + "error": { + "message": "An internal error occurred", + "type": "server_error", + }, + } + yield f"data: {json.dumps(error_chunk)}\n\n" + + return StreamingResponse( + gguf_tool_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + # ── Standard GGUF path (no tools) ───────────────────── + def gguf_generate(): return llama_backend.generate_chat_completion( messages = gguf_messages, @@ -945,6 +1072,8 @@ async def openai_chat_completions( enable_thinking = payload.enable_thinking, ) + _gguf_sentinel = object() + if payload.stream: async def gguf_stream_chunks(): @@ -963,12 +1092,17 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - # Content chunks — llama backend yields cumulative text + # Iterate the synchronous generator in a thread so + # the event loop stays free for disconnect detection. + gen = gguf_generate() prev_text = "" - for cumulative in gguf_generate(): + while True: if await request.is_disconnected(): cancel_event.set() return + cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) + if cumulative is _gguf_sentinel: + break new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 23fc856772..ee68c6bc0c 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -127,6 +127,7 @@ def apply_chat_template_to_dataset( auto_detect_mapping = True, batch_size = 1000, num_proc = None, + progress_callback = None, ): """ Applies chat template to dataset based on its format. @@ -364,8 +365,38 @@ def apply_chat_template_to_dataset( dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}" + # Monitor tqdm progress from dataset.map() and relay to callback + _tqdm_monitor_stop = None + if progress_callback and not isinstance(dataset, IterableDataset): + import threading + from tqdm.auto import tqdm as _tqdm_cls + + _tqdm_monitor_stop = threading.Event() + _total = len(dataset) if hasattr(dataset, "__len__") else 0 + _desc = f"Applying chat template to {final_format}" + + def _poll_tqdm(): + while not _tqdm_monitor_stop.is_set(): + for bar in list(getattr(_tqdm_cls, "_instances", set())): + try: + n = bar.n or 0 + total = bar.total or _total + if total > 0 and n > 0: + pct = min(int(n * 100 / total), 100) + progress_callback( + status_message = f"{_desc}... {pct}% ({n:,}/{total:,})" + ) + except (AttributeError, ReferenceError): + pass + _tqdm_monitor_stop.wait(3) + + threading.Thread(target = _poll_tqdm, daemon = True).start() + formatted_dataset = dataset.map(_format_chatml, **dataset_map_kwargs) + if _tqdm_monitor_stop is not None: + _tqdm_monitor_stop.set() + return { "dataset": formatted_dataset, "success": True, diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 9d15b86ca1..fac8c3d295 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -80,9 +80,6 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: multimodal_info = detect_multimodal_dataset(dataset) is_audio = multimodal_info.get("is_audio", False) - if multimodal_info["is_image"]: - is_vlm = True # Route to VLM detection for image datasets - # Common audio fields for all return paths audio_fields = { "is_audio": is_audio, @@ -153,8 +150,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "suggested_mapping": heuristic_mapping, "detected_image_column": None, "detected_text_column": None, - "is_image": False, - "multimodal_columns": None, + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), **audio_fields, } else: @@ -166,8 +163,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "suggested_mapping": None, "detected_image_column": None, "detected_text_column": None, - "is_image": False, - "multimodal_columns": None, + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), "warning": ( f"Could not auto-detect column roles for columns: {columns}. " "Please assign roles manually, or use AI Assist." @@ -183,8 +180,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "suggested_mapping": None, "detected_image_column": None, "detected_text_column": None, - "is_image": False, - "multimodal_columns": None, + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), **audio_fields, } @@ -1092,6 +1089,9 @@ def format_and_template_dataset( # LLM FLOW (Existing code) else: # Step 1: Format the dataset + n_rows = len(dataset) if hasattr(dataset, "__len__") else None + if progress_callback and n_rows: + progress_callback(status_message = f"Formatting dataset ({n_rows:,} rows)...") dataset_info = format_dataset( dataset, format_type = format_type, @@ -1106,6 +1106,11 @@ def format_and_template_dataset( ) # Step 2: Apply chat template + detected = dataset_info.get("detected_format", "unknown") + if progress_callback and n_rows: + progress_callback( + status_message = f"Applying chat template to {detected} ({n_rows:,} rows)..." + ) # Gemma emits a leading that must be stripped for text-only chatml/sharegpt. is_alpaca = format_type == "alpaca" or ( format_type == "auto" and dataset_info["detected_format"] == "alpaca" @@ -1124,6 +1129,7 @@ def format_and_template_dataset( auto_detect_mapping = auto_detect_mapping, batch_size = batch_size, num_proc = num_proc, + progress_callback = progress_callback, ) # Step 3: Generate summary diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index c1f8b62010..f2a34ecde9 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -8,6 +8,7 @@ Hardware detection and GPU utilities from .hardware import ( DeviceType, DEVICE, + CHAT_ONLY, detect_hardware, get_device, is_apple_silicon, @@ -18,12 +19,14 @@ from .hardware import ( get_package_versions, get_gpu_utilization, get_physical_gpu_count, + get_visible_gpu_count, safe_num_proc, ) __all__ = [ "DeviceType", "DEVICE", + "CHAT_ONLY", "detect_hardware", "get_device", "is_apple_silicon", @@ -34,5 +37,6 @@ __all__ = [ "get_package_versions", "get_gpu_utilization", "get_physical_gpu_count", + "get_visible_gpu_count", "safe_num_proc", ] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index c743dcd897..7234dbd155 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -39,6 +39,7 @@ class DeviceType(str, Enum): # ========== Global State (set once by detect_hardware) ========== DEVICE: Optional[DeviceType] = None +CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) # ========== Detection ========== @@ -81,7 +82,8 @@ def detect_hardware() -> DeviceType: 2. MLX (Apple Silicon via MLX framework) 3. CPU (fallback) """ - global DEVICE + global DEVICE, CHAT_ONLY + CHAT_ONLY = True # reset -- only CUDA sets it to False # --- CUDA: try PyTorch --- if _has_torch(): @@ -89,6 +91,7 @@ def detect_hardware() -> DeviceType: if torch.cuda.is_available(): DEVICE = DeviceType.CUDA + CHAT_ONLY = False device_name = torch.cuda.get_device_properties(0).name print(f"Hardware detected: CUDA — {device_name}") return DEVICE @@ -410,6 +413,7 @@ def get_gpu_utilization() -> Dict[str, Any]: # ========== Multi-GPU Detection & Safe num_proc ========== _physical_gpu_count: Optional[int] = None +_visible_gpu_count: Optional[int] = None def get_physical_gpu_count() -> int: @@ -443,21 +447,57 @@ def get_physical_gpu_count() -> int: return _physical_gpu_count +def get_visible_gpu_count() -> int: + """ + Return the number of GPUs visible to this process. + + Respects ``CUDA_VISIBLE_DEVICES`` -- if set, only those GPUs count. + Falls back to physical count if the env var is unset or torch is + unavailable. Result is cached after the first call. + """ + global _visible_gpu_count + if _visible_gpu_count is not None: + return _visible_gpu_count + + import os + + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if cuda_visible is not None: + # "" means zero GPUs, "0" means 1, "0,1,2" means 3 + cuda_visible = cuda_visible.strip() + if cuda_visible == "" or cuda_visible == "-1": + _visible_gpu_count = 0 + else: + _visible_gpu_count = len([x for x in cuda_visible.split(",") if x.strip()]) + return _visible_gpu_count + + # CUDA_VISIBLE_DEVICES not set -- try torch, fall back to physical count + try: + import torch + + _visible_gpu_count = torch.cuda.device_count() + except Exception: + _visible_gpu_count = get_physical_gpu_count() + + return _visible_gpu_count + + def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. On Windows, always returns 1 because Python uses ``spawn`` instead of - ``fork`` for multiprocessing — the overhead of re-importing torch, + ``fork`` for multiprocessing -- the overhead of re-importing torch, transformers, unsloth etc. per worker is typically slower than single-process for normal dataset sizes. - On multi-GPU machines the NVIDIA driver spawns extra background threads, - making ``os.fork()`` prone to deadlocks when many workers are created. + On multi-GPU machines (where multiple GPUs are *visible* to this + process) the NVIDIA driver spawns extra background threads, making + ``os.fork()`` prone to deadlocks when many workers are created. This helper caps ``num_proc`` to 4 on such machines. - On single-GPU (or CPU-only) machines the original value is returned - unchanged. + When ``CUDA_VISIBLE_DEVICES`` restricts to a single GPU, the cap + does not apply. Args: desired: The num_proc you *want*. If None, auto-computes from @@ -469,7 +509,7 @@ def safe_num_proc(desired: Optional[int] = None) -> int: import os import sys - # Windows uses 'spawn' for multiprocessing — the overhead of re-importing + # Windows uses 'spawn' for multiprocessing -- the overhead of re-importing # torch/transformers/unsloth per worker is typically slower than single-process. if sys.platform == "win32": return 1 @@ -477,11 +517,12 @@ def safe_num_proc(desired: Optional[int] = None) -> int: if desired is None or not isinstance(desired, int): desired = max(1, os.cpu_count() // 3) - if get_physical_gpu_count() > 1: + visible = get_visible_gpu_count() + if visible > 1: capped = min(4, desired) logger.info( - f"⚙️ Multi-GPU detected ({get_physical_gpu_count()} GPUs) — " - f"capping num_proc {desired} → {capped} to avoid fork deadlocks" + f"Multi-GPU detected ({visible} visible GPUs) -- " + f"capping num_proc {desired} -> {capped} to avoid fork deadlocks" ) return capped diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 2284cd194f..34f0f9b533 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -396,9 +396,12 @@ export function HubModelPicker({ return s; }, [cachedGguf, cachedModels]); + const chatOnly = usePlatformStore((s) => s.isChatOnly()); + const recommendedIds = useMemo(() => { const all = dedupe([...models.map((model) => model.id), value ?? ""]) - .filter((id) => !downloadedSet.has(id.toLowerCase())); + .filter((id) => !downloadedSet.has(id.toLowerCase())) + .filter((id) => !chatOnly || isGgufRepo(id)); // Cap at 4 GGUFs + 4 non-GGUFs so the list stays manageable const gguf: string[] = []; const hub: string[] = []; @@ -407,7 +410,7 @@ export function HubModelPicker({ else if (!isGgufRepo(id) && hub.length < 4) hub.push(id); } return [...gguf, ...hub]; - }, [models, value, downloadedSet]); + }, [models, value, downloadedSet, chatOnly]); const { paramCountById: recommendedParamCountById } = useRecommendedModelVram(recommendedIds); @@ -415,8 +418,6 @@ export function HubModelPicker({ const showHfSection = debouncedQuery.trim().length > 0; const recommendedSet = useMemo(() => new Set(recommendedIds), [recommendedIds]); - const chatOnly = usePlatformStore((s) => s.isChatOnly()); - const hfIds = useMemo(() => { if (!showHfSection) return []; return results @@ -519,7 +520,7 @@ export function HubModelPicker({ Loading models… - ) : !showHfSection && (cachedGguf.length > 0 || cachedModels.length > 0) ? ( + ) : !showHfSection && (cachedGguf.length > 0 || (!chatOnly && cachedModels.length > 0)) ? ( <> {"\uD83E\uDDA5"} Downloaded {cachedGguf.map((c) => ( @@ -536,7 +537,7 @@ export function HubModelPicker({ )} ))} - {cachedModels.map((c) => ( + {!chatOnly && cachedModels.map((c) => ( = ({ !thread.isEmpty}> {!hideComposer && } + @@ -153,12 +156,26 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { /> {!hideComposer && } + ); }; +const GeneratingSpinner: FC = () => { + const status = useChatRuntimeStore((s) => s.generatingStatus); + if (!status) return null; + return ( +
+
+ + Generating +
+
+ ); +}; + const ComposerAnimated: FC = () => { return ( { + { ); }; +/** Qwen3/3.5 recommended params differ between thinking on/off. */ +function applyQwenThinkingParams(thinkingOn: boolean): void { + const store = useChatRuntimeStore.getState(); + const checkpoint = store.params.checkpoint?.toLowerCase() ?? ""; + if (!checkpoint.includes("qwen3")) return; + // Qwen3 & Qwen3.5 share the same recommended settings: + // Thinking ON (general): temp=1.0, top_p=0.95, top_k=20 + // Thinking OFF (general): temp=0.7, top_p=0.8, top_k=20 + const params = thinkingOn + ? { temperature: 0.6, topP: 0.95, topK: 20, minP: 0.0 } + : { temperature: 0.7, topP: 0.8, topK: 20, minP: 0.0 }; + store.setParams({ ...store.params, ...params }); +} + const ReasoningToggle: FC = () => { const supportsReasoning = useChatRuntimeStore((s) => s.supportsReasoning); const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled); @@ -274,7 +306,11 @@ const ReasoningToggle: FC = () => { return ( + ); +}; + +const ToolStatusDisplay: FC = () => { + const toolStatus = useChatRuntimeStore((s) => s.toolStatus); + if (!toolStatus) return null; + return ( +
+
+ + {toolStatus} +
+
+ ); +}; + const ComposerAction: FC = () => { return (
@@ -300,6 +374,7 @@ const ComposerAction: FC = () => { +
diff --git a/studio/frontend/src/components/navbar.tsx b/studio/frontend/src/components/navbar.tsx index 00693fa064..0f159add77 100644 --- a/studio/frontend/src/components/navbar.tsx +++ b/studio/frontend/src/components/navbar.tsx @@ -77,7 +77,7 @@ export function Navbar() { alt="Unsloth" className="hidden h-9 w-auto dark:block" /> - + BETA diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index e5b42012a4..91e17f6bb9 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -16,14 +16,16 @@ export type DeviceType = "mac" | "windows" | "linux" | string; interface PlatformState { deviceType: DeviceType; + chatOnly: boolean; fetched: boolean; isChatOnly: () => boolean; } export const usePlatformStore = create()((_, get) => ({ deviceType: "linux", + chatOnly: false, fetched: false, - isChatOnly: () => get().deviceType === "mac", + isChatOnly: () => get().chatOnly, })); export async function fetchDeviceType(): Promise { @@ -33,9 +35,10 @@ export async function fetchDeviceType(): Promise { try { const res = await fetch("/api/health"); if (res.ok) { - const data = (await res.json()) as { device_type?: string }; + const data = (await res.json()) as { device_type?: string; chat_only?: boolean }; const deviceType = data.device_type ?? "linux"; - usePlatformStore.setState({ deviceType, fetched: true }); + const chatOnly = data.chat_only ?? deviceType === "mac"; + usePlatformStore.setState({ deviceType, chatOnly, fetched: true }); return deviceType; } } catch (err) { diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 3d986f1e89..ef40cbf2a4 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -131,3 +131,26 @@ export const MODEL_TYPE_TO_HF_TASK: Record = { audio: "text-to-speech", embeddings: "feature-extraction", }; + + +export const PRIORITY_TRAINING_MODELS: readonly string[] = [ + "unsloth/Qwen3.5-2B", + "unsloth/Qwen3.5-9B", + "unsloth/gpt-oss-20b", + "unsloth/NVIDIA-Nemotron-3-Nano-4B", + "unsloth/Qwen3-0.6B", + "unsloth/gemma-3-4b-it", + "unsloth/embeddinggemma-300m", + "unsloth/orpheus-3b-0.1-ft", + "unsloth/Llama-3.1-8B-Instruct", + "unsloth/Llama-3.2-3B-Instruct", +]; + +/** Pin priority models to the top of a list of model IDs, preserving their defined order. */ +export function applyPriorityOrdering(ids: string[]): string[] { + const idSet = new Set(ids); + const pinned = PRIORITY_TRAINING_MODELS.filter((id) => idSet.has(id)); + const pinnedSet = new Set(pinned); + const rest = ids.filter((id) => !pinnedSet.has(id)); + return [...pinned, ...rest]; +} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 3710cd795c..5686a61d12 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -14,6 +14,7 @@ import { } from "./chat-api"; import { db } from "../db"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import type { ChatModelSummary } from "../types/runtime"; import { hasClosedThinkTag, parseAssistantContent, @@ -234,10 +235,27 @@ async function autoLoadSmallestModel(): Promise { const store = useChatRuntimeStore.getState(); store.setCheckpoint(repo.repo_id, variant.quant); store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); + // Add model to store so the selector shows the name + const autoModel: ChatModelSummary = { + id: repo.repo_id, + name: loadResp.display_name ?? repo.repo_id, + isVision: loadResp.is_vision ?? false, + isLora: loadResp.is_lora ?? false, + isGguf: loadResp.is_gguf ?? false, + isAudio: loadResp.is_audio ?? false, + audioType: loadResp.audio_type ?? null, + hasAudioInput: loadResp.has_audio_input ?? false, + }; + const existingModels = store.models; + if (!existingModels.some((m) => m.id === repo.repo_id)) { + store.setModels([...existingModels, autoModel]); + } useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, supportsReasoning: loadResp.supports_reasoning ?? false, reasoningEnabled: loadResp.supports_reasoning ?? false, + supportsTools: loadResp.supports_tools ?? false, + toolsEnabled: false, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, }); @@ -255,7 +273,7 @@ async function autoLoadSmallestModel(): Promise { const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes); for (const repo of sorted) { try { - await loadModel({ + const sfLoadResp = await loadModel({ model_path: repo.repo_id, hf_token: null, max_seq_length: 4096, @@ -267,6 +285,16 @@ async function autoLoadSmallestModel(): Promise { const store = useChatRuntimeStore.getState(); store.setCheckpoint(repo.repo_id); store.setParams({ ...store.params, maxTokens: 4096 }); + const sfModel: ChatModelSummary = { + id: repo.repo_id, + name: sfLoadResp.display_name ?? repo.repo_id, + isVision: sfLoadResp.is_vision ?? false, + isLora: sfLoadResp.is_lora ?? false, + isGguf: sfLoadResp.is_gguf ?? false, + }; + if (!store.models.some((m) => m.id === repo.repo_id)) { + store.setModels([...store.models, sfModel]); + } toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); return true; } catch { @@ -275,8 +303,50 @@ async function autoLoadSmallestModel(): Promise { } } - toast.dismiss(toastId); - return false; + // No cached models found — try downloading a small default GGUF + toast("Downloading a small model…", { + id: toastId, + description: "No downloaded models found. Fetching Qwen3.5-4B (UD-Q4_K_XL).", + duration: 30000, + }); + try { + const loadResp = await loadModel({ + model_path: "unsloth/Qwen3.5-4B-GGUF", + hf_token: null, + max_seq_length: 4096, + load_in_4bit: true, + is_lora: false, + gguf_variant: "UD-Q4_K_XL", + trust_remote_code: false, + }); + const store = useChatRuntimeStore.getState(); + store.setCheckpoint("unsloth/Qwen3.5-4B-GGUF", "UD-Q4_K_XL"); + store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); + const defaultModel: ChatModelSummary = { + id: "unsloth/Qwen3.5-4B-GGUF", + name: loadResp.display_name ?? "Qwen3.5-4B-GGUF", + isVision: loadResp.is_vision ?? false, + isLora: false, + isGguf: true, + }; + if (!store.models.some((m) => m.id === "unsloth/Qwen3.5-4B-GGUF")) { + store.setModels([...store.models, defaultModel]); + } + useChatRuntimeStore.setState({ + ggufContextLength: loadResp.context_length ?? 131072, + supportsReasoning: loadResp.supports_reasoning ?? false, + reasoningEnabled: loadResp.supports_reasoning ?? false, + supportsTools: loadResp.supports_tools ?? false, + toolsEnabled: false, + defaultChatTemplate: loadResp.chat_template ?? null, + chatTemplateOverride: null, + }); + toast.success("Loaded Qwen3.5-4B (UD-Q4_K_XL)", { id: toastId }); + return true; + } catch { + toast.dismiss(toastId); + return false; + } } catch { toast.dismiss(toastId); return false; @@ -306,6 +376,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } } + const { + supportsTools, + toolsEnabled, + } = runtime; + const outboundMessages = messages .map(toOpenAIMessage) .filter((message): message is NonNullable => @@ -415,14 +490,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (!waitingFirstChunk) return; if (abortSignal.aborted) return; warmupToastShown = true; - toast.promise(firstTokenPromise, { - loading: "Generating", - success: "Generating", - error: (err) => - err instanceof Error && err.message ? err.message : "Generation failed", - description: "Waiting for first token.", - duration: 900, - }); + runtime.setGeneratingStatus("waiting"); }, warmupDelayMs); runtime.setThreadRunning(threadKey, true); let cumulativeText = ""; @@ -446,11 +514,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { audio_base64: audioBase64, ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }), ...(supportsReasoning ? { enable_thinking: reasoningEnabled } : {}), + ...(supportsTools && toolsEnabled ? { enable_tools: true } : {}), }, abortSignal, ); for await (const chunk of stream) { + // Handle tool status events + const toolStatusText = (chunk as unknown as { _toolStatus?: string })._toolStatus; + if (toolStatusText !== undefined) { + runtime.setToolStatus(toolStatusText || null); + continue; + } + totalChunks += 1; const delta = chunk.choices?.[0]?.delta?.content; if (!delta) { @@ -460,6 +536,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { waitingFirstChunk = false; firstTokenTime = Date.now() - streamStartTime; settleFirstTokenOk(); + runtime.setGeneratingStatus(null); } cumulativeText += delta; @@ -501,17 +578,18 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }; } catch (err) { settleFirstTokenErr(err instanceof Error ? err : new Error("Generation failed")); - const isEarly = waitingFirstChunk; - if (!abortSignal.aborted && !(warmupToastShown && isEarly)) { + if (!abortSignal.aborted) { toast.error("Generation failed", { description: err instanceof Error ? err.message : "Unknown error", }); } throw err; } finally { + runtime.setGeneratingStatus(null); + runtime.setToolStatus(null); clearTimeout(warmupTimer); if (waitingFirstChunk) { - if (warmupToastShown && !firstTokenSettled) { + if (!firstTokenSettled) { if (abortSignal.aborted) { settleFirstTokenErr(new Error("Cancelled")); } else { diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 01bc762d86..effba93ef8 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -213,10 +213,16 @@ export async function* streamChatCompletions( const parsed = JSON.parse(dataText) as | OpenAIChatChunk - | { error?: { message?: string } }; + | { type?: string; content?: string; error?: { message?: string } }; if ("error" in parsed && parsed.error) { throw new Error(parsed.error.message || "Stream error"); } + // Tool status events are custom SSE payloads, not OpenAI chunks + if ("type" in parsed && parsed.type === "tool_status") { + yield { _toolStatus: parsed.content ?? "" } as unknown as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } yield parsed as OpenAIChatChunk; separatorIndex = buffer.search(/\r?\n\r?\n/); } diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 170429a44b..b14482308a 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -166,6 +166,8 @@ export function ChatSettingsPanel({ }: ChatSettingsPanelProps) { const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); + const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype); const [presets, setPresets] = useState(BUILTIN_PRESETS); const [activePreset, setActivePreset] = useState("Default"); const isBuiltinPreset = BUILTIN_PRESETS.some((p) => p.name === activePreset); @@ -382,6 +384,34 @@ export function ChatSettingsPanel({ onCheckedChange={set("trustRemoteCode")} />
+ {isGguf && ( +
+
+
KV Cache Dtype
+
+ Quantize KV cache to reduce VRAM. Reload to apply. +
+
+ +
+ )} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index b79c3073b2..545b7711ed 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -340,7 +340,7 @@ export function useChatModelRuntime() { previousWasUnloaded = true; } - const chatTemplateOverride = useChatRuntimeStore.getState().chatTemplateOverride; + const { chatTemplateOverride, kvCacheDtype } = useChatRuntimeStore.getState(); const loadResponse = await loadModel({ model_path: modelId, hf_token: null, @@ -350,6 +350,7 @@ export function useChatModelRuntime() { gguf_variant: ggufVariant ?? null, trust_remote_code: paramsBeforeLoad.trustRemoteCode ?? false, chat_template_override: chatTemplateOverride, + cache_type_kv: kvCacheDtype, }); // If cancelled while loading, don't update UI to show @@ -360,15 +361,37 @@ export function useChatModelRuntime() { setParams( mergeRecommendedInference(currentParams, loadResponse, modelId), ); + // Qwen3.5 small models (0.8B, 2B, 4B, 9B) disable thinking by default + let reasoningDefault = loadResponse.supports_reasoning ?? false; + if (reasoningDefault) { + const mid = modelId.toLowerCase(); + if (mid.includes("qwen3.5")) { + const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); + if (sizeMatch && parseFloat(sizeMatch[1]) <= 2) { + reasoningDefault = false; + } + } + } useChatRuntimeStore.setState({ ggufContextLength: loadResponse.is_gguf ? (loadResponse.context_length ?? 131072) : null, supportsReasoning: loadResponse.supports_reasoning ?? false, - reasoningEnabled: loadResponse.supports_reasoning ?? false, + reasoningEnabled: reasoningDefault, + supportsTools: loadResponse.supports_tools ?? false, + toolsEnabled: false, + kvCacheDtype: loadResponse.cache_type_kv ?? null, defaultChatTemplate: loadResponse.chat_template ?? null, chatTemplateOverride: null, }); + // Qwen3/3.5: apply thinking-mode-specific params after load + if (modelId.toLowerCase().includes("qwen3") && (loadResponse.supports_reasoning ?? false)) { + const store = useChatRuntimeStore.getState(); + const p = reasoningDefault + ? { temperature: 0.6, topP: 0.95, topK: 20, minP: 0.0 } + : { temperature: 0.7, topP: 0.8, topK: 20, minP: 0.0 }; + store.setParams({ ...store.params, ...p }); + } await refresh(); } catch (error) { // Skip rollback if user cancelled -- model is already being unloaded. diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 360f92a685..3eb1f699d2 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -6,7 +6,7 @@ import { Button } from "@/components/ui/button"; import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils"; import { useAui } from "@assistant-ui/react"; import { cn } from "@/lib/utils"; -import { ArrowUpIcon, HeadphonesIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; +import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { type KeyboardEvent, @@ -202,6 +202,9 @@ export function SharedComposer({ const supportsReasoning = useChatRuntimeStore((s) => s.supportsReasoning); const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled); const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled); + const supportsTools = useChatRuntimeStore((s) => s.supportsTools); + const toolsEnabled = useChatRuntimeStore((s) => s.toolsEnabled); + const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled); const setPendingAudioStore = useChatRuntimeStore((s) => s.setPendingAudio); const clearPendingAudioStore = useChatRuntimeStore((s) => s.clearPendingAudio); @@ -393,7 +396,19 @@ export function SharedComposer({ {supportsReasoning && ( )} + {supportsTools && ( + + )}
{dictationSupported && ( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 7fc1998b35..5715705725 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -46,6 +46,11 @@ type ChatRuntimeStore = { ggufContextLength: number | null; supportsReasoning: boolean; reasoningEnabled: boolean; + supportsTools: boolean; + toolsEnabled: boolean; + toolStatus: string | null; + generatingStatus: string | null; + kvCacheDtype: string | null; defaultChatTemplate: string | null; chatTemplateOverride: string | null; activeThreadId: string | null; @@ -63,6 +68,10 @@ type ChatRuntimeStore = { setActiveThreadId: (threadId: string | null) => void; clearCheckpoint: () => void; setReasoningEnabled: (enabled: boolean) => void; + setToolsEnabled: (enabled: boolean) => void; + setToolStatus: (status: string | null) => void; + setGeneratingStatus: (status: string | null) => void; + setKvCacheDtype: (dtype: string | null) => void; setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; @@ -79,6 +88,11 @@ export const useChatRuntimeStore = create((set) => ({ ggufContextLength: null, supportsReasoning: false, reasoningEnabled: true, + supportsTools: false, + toolsEnabled: false, + toolStatus: null, + generatingStatus: null, + kvCacheDtype: null, defaultChatTemplate: null, chatTemplateOverride: null, activeThreadId: null, @@ -124,10 +138,18 @@ export const useChatRuntimeStore = create((set) => ({ ggufContextLength: null, supportsReasoning: false, reasoningEnabled: true, + supportsTools: false, + toolsEnabled: false, + toolStatus: null, + kvCacheDtype: null, defaultChatTemplate: null, chatTemplateOverride: null, })), setReasoningEnabled: (reasoningEnabled) => set({ reasoningEnabled }), + setToolsEnabled: (toolsEnabled) => set({ toolsEnabled }), + setToolStatus: (toolStatus) => set({ toolStatus }), + setGeneratingStatus: (generatingStatus) => set({ generatingStatus }), + setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }), setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }), setPendingAudio: (base64, name) => set({ pendingAudioBase64: base64, pendingAudioName: name }), diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index f9eca3455d..d40b5fa853 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -40,6 +40,7 @@ export interface LoadModelRequest { /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trust_remote_code?: boolean; chat_template_override?: string | null; + cache_type_kv?: string | null; } export interface ValidateModelResponse { @@ -85,6 +86,8 @@ export interface LoadModelResponse { }; context_length?: number | null; supports_reasoning?: boolean; + supports_tools?: boolean; + cache_type_kv?: string | null; chat_template?: string | null; } @@ -139,6 +142,7 @@ export interface OpenAIChatCompletionsRequest { audio_base64?: string; use_adapter?: boolean | string | null; enable_thinking?: boolean | null; + enable_tools?: boolean | null; } export interface OpenAIChatDelta { diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 5e981fef3b..f05643c092 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -33,7 +33,7 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { MODEL_TYPE_TO_HF_TASK } from "@/config/training"; +import { MODEL_TYPE_TO_HF_TASK, PRIORITY_TRAINING_MODELS, applyPriorityOrdering } from "@/config/training"; import { useDebouncedValue, useGpuInfo, @@ -96,12 +96,16 @@ export function ModelSelectionStep() { task, accessToken: hfToken || undefined, excludeGguf: true, + priorityIds: PRIORITY_TRAINING_MODELS, }); const { error: tokenValidationError, isChecking: isCheckingToken } = useHfTokenValidation(hfToken); - const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]); + const resultIds = useMemo(() => { + const ids = hfResults.map((r) => r.id); + return applyPriorityOrdering(ids); + }, [hfResults]); // Match Studio behavior: only show exception signals (OOM/TIGHT) in training flows. const vramMap = useMemo(() => { diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 72d6abcf59..77edf68d24 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -28,7 +28,7 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { MODEL_TYPE_TO_HF_TASK } from "@/config/training"; +import { MODEL_TYPE_TO_HF_TASK, PRIORITY_TRAINING_MODELS, applyPriorityOrdering } from "@/config/training"; import { useDebouncedValue, useGpuInfo, @@ -162,6 +162,7 @@ export function ModelSection() { task, accessToken: hfToken || undefined, excludeGguf: true, + priorityIds: PRIORITY_TRAINING_MODELS, }); const { error: tokenValidationError, isChecking: isCheckingToken } = @@ -172,7 +173,8 @@ export function ModelSection() { if (selectedModel && !ids.includes(selectedModel)) { ids.push(selectedModel); } - return ids; + + return applyPriorityOrdering(ids); }, [hfResults, selectedModel]); // Filter out GGUF models — they can't be used for training diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx index d1cec00e22..a47654d290 100644 --- a/studio/frontend/src/features/studio/sections/training-section.tsx +++ b/studio/frontend/src/features/studio/sections/training-section.tsx @@ -46,7 +46,8 @@ export function TrainingSection() { const store = useTrainingConfigStore(); const { isStarting, startError, startTrainingRun } = useTrainingActions(); const isIncompatible = - !store.isVisionModel && store.isDatasetImage === true; + (!store.isVisionModel && store.isDatasetImage === true) || + (!store.isAudioModel && store.isDatasetAudio === true); const configValidation = validateTrainingConfig(store); const fileInputRef = useRef(null); @@ -155,10 +156,10 @@ export function TrainingSection() { data-tour="studio-start" className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600" onClick={() => void startTrainingRun()} - disabled={isStarting || isIncompatible || !configValidation.ok} + disabled={isStarting || isIncompatible || store.isCheckingDataset || !configValidation.ok} > - {isStarting ? "Starting..." : "Start Training"} + {isStarting ? "Starting..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"} {startError && (

{startError}

diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 212dc78e6c..e1b8d4b1a2 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -210,6 +210,7 @@ export const useTrainingConfigStore = create()( hfToken: state.hfToken.trim() || null, subset: state.datasetSubset, split, + isVlm: state.isVisionModel, }) .then((res) => { if (controller.signal.aborted) return; diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index f7b06ab65a..69ea4d3b83 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { PipelineType } from "@huggingface/hub"; -import { listModels } from "@huggingface/hub"; +import { listModels, modelInfo } from "@huggingface/hub"; import { useCallback, useMemo } from "react"; import { useHfPaginatedSearch } from "./use-hf-paginated-search"; @@ -148,17 +148,73 @@ async function* mergedModelIterator( } } +/** + * Creates an async generator that yields priority models (fetched individually + * via modelInfo for full metadata), then the general unsloth listing. + */ +async function* priorityThenListingIterator( + priorityIds: readonly string[], + task?: PipelineType, + accessToken?: string, +): AsyncGenerator { + const common = { + additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[], + fetch: withPopularitySort, + ...(accessToken ? { credentials: { accessToken } } : {}), + }; + + // Phase 1: fetch priority models in parallel via modelInfo + const seen = new Set(); + const settled = await Promise.allSettled( + priorityIds.map((id) => + modelInfo({ + name: id, + additionalFields: ["safetensors", "tags"], + ...(accessToken ? { credentials: { accessToken } } : {}), + }), + ), + ); + for (const result of settled) { + if (result.status === "fulfilled") { + const m = result.value as { name?: string; pipeline_tag?: string }; + // Skip models that don't match the selected task filter + if (task && m.pipeline_tag && m.pipeline_tag !== task) continue; + if (m.name) seen.add(m.name); + yield result.value; + } + } + + // Phase 2: yield general unsloth listing, skipping already-seen + const generalIter = listModels({ + search: { owner: "unsloth", ...(task ? { task } : {}) }, + ...common, + }); + for await (const model of generalIter) { + const m = model as { name?: string }; + if (m.name && seen.has(m.name)) continue; + yield model; + } +} + export function useHfModelSearch( query: string, - options?: { task?: PipelineType; accessToken?: string; excludeGguf?: boolean }, + options?: { + task?: PipelineType; + accessToken?: string; + excludeGguf?: boolean; + priorityIds?: readonly string[]; + }, ) { - const { task, accessToken, excludeGguf = false } = options ?? {}; + const { task, accessToken, excludeGguf = false, priorityIds } = options ?? {}; const createIter = useCallback( () => { const trimmed = query.trim(); if (!trimmed) { - // No query → show default unsloth models + // No query → show priority models first (with full metadata), then general unsloth listing + if (priorityIds && priorityIds.length > 0) { + return priorityThenListingIterator(priorityIds, task, accessToken) as AsyncGenerator; + } return listModels({ search: { owner: "unsloth", ...(task ? { task } : {}) }, additionalFields: ["safetensors", "tags"], @@ -169,7 +225,7 @@ export function useHfModelSearch( // Typed query: disable task filter so explicitly searched models still appear even if HF task metadata is wrong/missing. return mergedModelIterator(trimmed, undefined, accessToken) as AsyncGenerator; }, - [query, task, accessToken], + [query, task, accessToken, priorityIds], ); const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]); diff --git a/studio/setup.ps1 b/studio/setup.ps1 index d9e7c5a43b..15b5d217b0 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1054,6 +1054,10 @@ if (Test-Path $LlamaServerBin) { } # Common flags $CmakeArgs += '-DBUILD_SHARED_LIBS=OFF' + $CmakeArgs += '-DLLAMA_BUILD_TESTS=OFF' + $CmakeArgs += '-DLLAMA_BUILD_EXAMPLES=OFF' + $CmakeArgs += '-DLLAMA_BUILD_SERVER=ON' + $CmakeArgs += '-DGGML_NATIVE=ON' # HTTPS support via OpenSSL if ($OpenSslAvailable -and $OpenSslRoot) { $CmakeArgs += "-DOPENSSL_ROOT_DIR=$OpenSslRoot" diff --git a/studio/setup.sh b/studio/setup.sh index 5fbdce918f..7fe8dd7e51 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -168,7 +168,8 @@ for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)? continue fi # Get version string, e.g. "Python 3.12.5" - ver_str=$("$candidate" --version 2>&1 | awk '{print $2}') + ver_str=$("$candidate" --version 2>&1) || continue + ver_str=$(echo "$ver_str" | awk '{print $2}') py_major=$(echo "$ver_str" | cut -d. -f1) py_minor=$(echo "$ver_str" | cut -d. -f2) @@ -194,7 +195,7 @@ for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)? BEST_MINOR="$py_minor" fi done - +echo "finished finding best python" if [ -z "$BEST_PY" ]; then echo "❌ ERROR: No Python version between 3.${MIN_PY_MINOR} and 3.${MAX_PY_MINOR} found on this system." echo " Detected Python 3 installations:" @@ -296,7 +297,15 @@ rm -rf "$LLAMA_CPP_DIR" run_quiet "clone llama.cpp" git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$LLAMA_CPP_DIR" || BUILD_OK=false if [ "$BUILD_OK" = true ]; then - CMAKE_ARGS="" + # Skip tests/examples we don't need (faster build) + CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON" + + # Use ccache if available (dramatically faster rebuilds) + if command -v ccache &>/dev/null; then + CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache" + echo " Using ccache for faster compilation" + fi + # Detect CUDA: check nvcc on PATH, then common install locations NVCC_PATH="" if command -v nvcc &>/dev/null; then @@ -312,7 +321,7 @@ rm -rf "$LLAMA_CPP_DIR" if [ -n "$NVCC_PATH" ]; then echo " Building with CUDA support (nvcc: $NVCC_PATH)..." - CMAKE_ARGS="-DGGML_CUDA=ON" + CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON" # Detect GPU compute capability and limit CUDA architectures # Without this, cmake builds for ALL default archs (very slow) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 5c61a6eae2..a05186eee5 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1119,14 +1119,15 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): if "dataset_num_proc" in call_args: num_proc_check = ( "import multiprocessing as _mp\n" - "if _mp.get_start_method() != 'fork':\n" - " dataset_num_proc = None\n" - "elif dataset_num_proc is None:\n" - " import psutil\n" - " dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)\n" - " memory_gb_left = psutil.virtual_memory().available / (1024**3)\n" - " if memory_gb_left <= 2: dataset_num_proc = 1\n" - " else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))\n" + "if dataset_num_proc is None:\n" + " if _mp.get_start_method() != 'fork':\n" + " dataset_num_proc = None\n" + " else:\n" + " import psutil\n" + " dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)\n" + " memory_gb_left = psutil.virtual_memory().available / (1024**3)\n" + " if memory_gb_left <= 2: dataset_num_proc = 1\n" + " else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))\n" ) extra_args += num_proc_check