Merge remote-tracking branch 'origin/nightly' into fix/dataset-mapping-vlm-text-datasets

README.md

@@ -45,12 +45,14 @@ This script will:
 ### Launch the studio
 
 ```bash
-# After setup, open a new terminal (or source ~/.bashrc), then:
+# After setup, open a new terminal (or source ~/.bashrc), then inside your working directory:
 unsloth-ui -H 0.0.0.0 -p 8000
 ```
 
 On **first launch**, a one-time setup token is printed to the console. Use it in the browser to create your admin account.
 
+As this repo is in continuous development, please make sure to run the setup.sh file everytime you pull new changes from the repo.
+
 ## API Reference
 
 All endpoints require a valid JWT `Authorization: Bearer <token>` header (except `/api/auth/*` and `/api/health`).

studio/backend/core/training/trainer.py

@@ -115,6 +115,14 @@ class UnslothTrainer:
                    is_dataset_multimodal: bool = False) -> bool:
         """Load model for training (supports both text and vision models)"""
         try:
+            if self.model is not None:
+                del self.model
+            if self.tokenizer is not None:
+                del self.tokenizer
+
+            if self.trainer is not None:
+                del self.trainer
+
             print("\nClearing GPU memory before training...")
             clear_gpu_cache()
 
@@ -214,7 +222,13 @@ class UnslothTrainer:
                 return True
 
             # LoRA/QLoRA mode - apply PEFT
-            if target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
+            # "all-linear" is a PEFT keyword that targets every linear layer
+            if isinstance(target_modules, list) and "all-linear" in target_modules:
+                if len(target_modules) == 1:
+                    target_modules = "all-linear"
+                else:
+                    target_modules = [m for m in target_modules if m != "all-linear"]
+            elif target_modules is None or (isinstance(target_modules, list) and len(target_modules) == 0):
                 target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                                 "gate_proj", "up_proj", "down_proj"]
 
@@ -701,6 +715,7 @@ class UnslothTrainer:
                 "output_dir": output_dir,
                 "report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none",
                 "include_num_input_tokens_seen": True,  # Enable token counting
+                "dataset_num_proc": max(1, os.cpu_count() // 3),
             }
             
             # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps

studio/backend/main.py

@@ -1,6 +1,7 @@
 """
 Main FastAPI application for Unsloth UI Backend
 """
+import os
 import secrets
 import shutil
 from contextlib import asynccontextmanager
@@ -15,7 +16,7 @@ from datetime import datetime
 # Import routers
 from routes import training_router, models_router, inference_router, datasets_router, auth_router, export_router
 from auth import storage
-from utils.hardware import detect_hardware
+from utils.hardware import detect_hardware, get_device, DeviceType
 import utils.hardware.hardware as _hw_module
 
 UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache"
@@ -27,6 +28,18 @@ async def lifespan(app: FastAPI):
     # Detect hardware first — sets DEVICE global used everywhere
     detect_hardware()
 
+    # Disable flex attention on Blackwell+ GPUs (sm_120 and above)
+    if get_device() == DeviceType.CUDA:
+        import torch
+        props = torch.cuda.get_device_properties(0)
+        sm_version = props.major * 10 + props.minor
+        if sm_version >= 120:
+            os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0"
+            import logging
+            logging.getLogger(__name__).info(
+                f"GPU sm_{sm_version} detected — setting UNSLOTH_FLEX_ATTENTION=0"
+            )
+
     if not storage.is_initialized():
         setup_token = secrets.token_urlsafe(32)
         storage.save_setup_token(setup_token)

studio/backend/routes/datasets.py

@@ -75,12 +75,16 @@ async def check_format(request: CheckFormatRequest):
     """
     Check if a dataset requires manual column mapping.
     
-    This is a lightweight check that loads only the first 10 rows,
+    This is a lightweight check that streams only the first N rows,
     runs format detection, and (if processable) returns processed
     preview samples. The full dataset is re-processed at training time.
+    
+    For HuggingFace datasets we use streaming mode so we never download
+    the entire dataset — only the rows we actually need are fetched.
     """
     try:
-        from datasets import load_dataset
+        from itertools import islice
+        from datasets import Dataset, load_dataset
         from utils.datasets import format_dataset
         
         PREVIEW_SIZE = 10
@@ -89,9 +93,10 @@ async def check_format(request: CheckFormatRequest):
         
         # Load dataset
         dataset_path = Path(request.dataset_name)
+        total_rows = None
         
         if dataset_path.exists():
-            # Local dataset
+            # Local dataset — direct load is fine (files are local)
             if dataset_path.suffix in ['.json', '.jsonl']:
                 dataset = load_dataset('json', data_files=str(dataset_path), split=request.train_split)
             elif dataset_path.suffix == '.csv':
@@ -103,18 +108,30 @@ async def check_format(request: CheckFormatRequest):
                     status_code=400,
                     detail=f"Unsupported file format: {dataset_path.suffix}"
                 )
+            total_rows = len(dataset)
+            preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
         else:
-            # HuggingFace dataset
-            load_kwargs = {"path": request.dataset_name, "split": request.train_split}
+            # HuggingFace dataset — use STREAMING to avoid downloading everything
+            load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True}
             if request.subset:
                 load_kwargs["name"] = request.subset
             if request.hf_token:
                 load_kwargs["token"] = request.hf_token
-            dataset = load_dataset(**load_kwargs)
-        
-        # Slice to top N rows — all detection and preview runs on this subset
-        total_rows = len(dataset)
-        preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
+            
+            streamed_ds = load_dataset(**load_kwargs)
+            
+            # Take only the first PREVIEW_SIZE rows from the stream
+            rows = list(islice(streamed_ds, PREVIEW_SIZE))
+            if not rows:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Dataset appears to be empty or could not be streamed"
+                )
+            
+            # Convert list-of-dicts into a proper Dataset for downstream compat
+            preview_slice = Dataset.from_list(rows)
+            # total_rows unknown in streaming mode
+            total_rows = None
         
         # Run lightweight format check on the preview slice
         result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)

studio/backend/routes/models.py

@@ -248,6 +248,7 @@ async def get_model_config(
     This endpoint wraps the backend load_model_defaults function.
     """
     try:
+        logger.info(f"Getting model config for: {model_name}")
         # Load model defaults from backend
         config_dict = load_model_defaults(model_name)
         
@@ -267,6 +268,7 @@ async def get_model_config(
             # If ModelConfig creation fails, use defaults
             pass
         
+        logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_lora={is_lora}, base_model={base_model}")
         return ModelDetails(
             id=model_name,
             model_name=model_name,
@@ -369,8 +371,10 @@ async def check_vision_model(
     This endpoint wraps the backend is_vision_model function.
     """
     try:
+        logger.info(f"Checking if vision model: {model_name}")
         is_vision = is_vision_model(model_name)
         
+        logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
         return VisionCheckResponse(
             model_name=model_name,
             is_vision=is_vision,

studio/backend/utils/datasets/chat_templates.py

@@ -285,7 +285,7 @@ def apply_chat_template_to_dataset(
             if not isinstance(dataset, IterableDataset):
                 from multiprocessing import cpu_count
                 if num_proc is None or type(num_proc) is not int:
-                    num_proc = cpu_count()
+                    num_proc = max(1, cpu_count() // 3)
                 dataset_map_kwargs['num_proc'] = num_proc
                 dataset_map_kwargs['desc'] = "Applying template to Alpaca format"
 
@@ -349,7 +349,7 @@ def apply_chat_template_to_dataset(
             if not isinstance(dataset, IterableDataset):
                 from multiprocessing import cpu_count
                 if num_proc is None or type(num_proc) is not int:
-                    num_proc = cpu_count()
+                    num_proc = max(1, cpu_count() // 3)
                 dataset_map_kwargs['num_proc'] = num_proc
                 dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}"
 

studio/backend/utils/datasets/format_conversion.py

@@ -147,7 +147,7 @@ def standardize_chat_format(
         from multiprocessing import cpu_count
 
         if num_proc is None or type(num_proc) is not int:
-            num_proc = cpu_count()
+            num_proc = max(1, cpu_count() // 3)
 
         dataset_map_kwargs['num_proc'] = num_proc
         dataset_map_kwargs['desc'] = "Standardizing chat format"
@@ -213,7 +213,7 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
         from multiprocessing import cpu_count
 
         if num_proc is None or type(num_proc) is not int:
-            num_proc = cpu_count()
+            num_proc = max(1, cpu_count() // 3)
 
         dataset_map_kwargs['num_proc'] = num_proc
         dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format"
@@ -261,7 +261,7 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
         from multiprocessing import cpu_count
 
         if num_proc is None or type(num_proc) is not int:
-            num_proc = cpu_count()
+            num_proc = max(1, cpu_count() // 3)
 
         dataset_map_kwargs['num_proc'] = num_proc
         dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"

studio/backend/utils/hardware/hardware.py

@@ -120,6 +120,7 @@ def clear_gpu_cache():
 
     if device == DeviceType.CUDA:
         import torch
+        torch.cuda.synchronize()
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
     elif device == DeviceType.MLX:

studio/frontend/src/components/navbar.tsx

@@ -21,6 +21,7 @@ import {
   ZapIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
+import { useTrainingRuntimeStore } from "@/features/training";
 import { Link, useRouterState } from "@tanstack/react-router";
 import { AnimatePresence, motion } from "motion/react";
 import { useState } from "react";
@@ -35,6 +36,7 @@ const NAV_ITEMS = [
 
 export function Navbar() {
   const pathname = useRouterState({ select: (s) => s.location.pathname });
+  const isTrainingRunning = useTrainingRuntimeStore((s) => s.isTrainingRunning);
   const [logoHovered, setLogoHovered] = useState(false);
   const [mobileOpen, setMobileOpen] = useState(false);
 
@@ -95,7 +97,9 @@ export function Navbar() {
         >
           {NAV_ITEMS.map((item) => {
             const active = pathname === item.href;
-            if (!item.enabled) {
+            const disabledByTraining =
+              isTrainingRunning && item.href !== "/studio";
+            if (!item.enabled || disabledByTraining) {
               return (
                 <span
                   key={item.href}
@@ -230,6 +234,18 @@ export function Navbar() {
               <div className="mt-6 flex flex-col gap-2">
                 {NAV_ITEMS.filter((item) => item.enabled).map((item) => {
                   const active = pathname === item.href;
+                  const disabledByTraining =
+                    isTrainingRunning && item.href !== "/studio";
+                  if (disabledByTraining) {
+                    return (
+                      <span
+                        key={item.href}
+                        className="rounded-md border border-border px-3 py-2 text-sm font-medium text-muted-foreground/40 cursor-not-allowed"
+                      >
+                        {item.label}
+                      </span>
+                    );
+                  }
                   return (
                     <Link
                       key={item.href}