diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml new file mode 100644 index 0000000000..1ee6489209 --- /dev/null +++ b/.github/workflows/studio-export-capability-ci.yml @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS. +# +# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per +# platform) and the export backend must import without PyTorch, so this confirms the gating and +# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator +# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block +# torch/unsloth, so the job installs only a CPU PyTorch plus import deps. + +name: Studio export capability + +on: + pull_request: + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + capability: + name: capability (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + env: + # No accelerator on hosted runners; keep detection on the CPU path. + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Upgrade pip + run: python -m pip install --upgrade pip + - name: Install CPU PyTorch + # CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's + # transitive deps still resolve (matching the other workflows in this repo). + run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13" + - name: Install backend import deps + # Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and + # the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds). + run: python -m pip install + transformers peft accelerate safetensors huggingface_hub datasets + sentencepiece protobuf fastapi starlette structlog psutil + python-multipart pydantic httpx "numpy<3" pytest + - name: Export capability + import-safety tests + working-directory: studio/backend + run: python -m pytest tests/test_export_capability.py -q diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d0461dae95..c8be50b08b 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -13,7 +13,18 @@ import shutil import contextlib from pathlib import Path from typing import Optional, Tuple, List -from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + +# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable +# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash. +try: + from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + _UNSLOTH_IMPORT_ERROR = None +except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load + FastLanguageModel = None + FastVisionModel = None + _IS_MLX = False + _UNSLOTH_IMPORT_ERROR = _unsloth_exc + from huggingface_hub import HfApi, ModelCard from utils.hardware import clear_gpu_cache @@ -27,14 +38,46 @@ from utils.paths import ( ) from core.inference import get_inference_backend -# GPU-only imports — guarded for Apple Silicon where these aren't needed +# GPU/PyTorch-only imports, skipped on MLX and on a --no-torch install so the module stays +# importable; export then degrades to a clear "PyTorch is not installed" error. +torch = None +_TORCH_IMPORT_ERROR: Optional[BaseException] = None if not _IS_MLX: - from peft import PeftModel, PeftModelForCausalLM - from transformers.modeling_utils import PushToHubMixin - import torch + try: + from peft import PeftModel, PeftModelForCausalLM + from transformers.modeling_utils import PushToHubMixin + import torch + except Exception as _torch_exc: # ImportError, or a broken native torch load + _TORCH_IMPORT_ERROR = _torch_exc logger = get_logger(__name__) + +def _export_runtime_available() -> bool: + """True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host).""" + return bool(_IS_MLX) or (FastLanguageModel is not None) + + +def _export_runtime_message() -> str: + """Precise reason the export runtime is unavailable, mirroring hardware.export_capability().""" + if torch is None: + return ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + return ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on " + "CPU only.)" + ) + + +# Kept for call sites / tests referencing the PyTorch-missing text. +_PYTORCH_MISSING_MESSAGE = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." +) + _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False @@ -58,6 +101,28 @@ def _compressed_export_supported(): return False +def _torchao_export_supported(): + """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_torchao_method") + except Exception: + return False + + +def _has_nvidia_gpu(): + """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" + try: + from utils.hardware import hardware as _hw + return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM + except Exception: + try: + import torch + return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + except Exception: + return False + + def _hf_offline(timeout = 3): """True if export should avoid the Hub: honors the HF offline env vars, else does one cheap TCP reachability probe so a network-down load uses local files / the HF cache @@ -394,13 +459,17 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). Args: save_directory: Local directory to save model - format_type: "16-bit (FP16)" or "4-bit (FP4)" + format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label + compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8", + "fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides + format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES. push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID (username/model-name) hf_token: Hugging Face token @@ -409,38 +478,108 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - if not self.is_peft: - return ( - False, - "This is not a PEFT model. Use 'Export Base Model' instead.", - None, - ) + # Merged export works for PEFT adapters and non-PEFT Local/HF base models alike + # (save_pretrained_merged is a no-op merge that just saves the base). output_path: Optional[str] = None - # compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and - # write to a sibling "-" directory (for vLLM). - _COMPRESSED = { - "FP8 (compressed-tensors)": ("fp8", "fp8"), - "NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"), + # Quantized formats save to a sibling "-". Two backends: compressed-tensors + # (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias + # comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label. + _LABEL_TO_ALIAS = { + "FP8 (compressed-tensors)": "fp8", + "NVFP4 (compressed-tensors)": "nvfp4", } - is_compressed = format_type in _COMPRESSED + compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + compressed_suffix: Optional[str] = None + # Classify the alias: torchao-portable vs compressed-tensors. + torchao_info = None + if compressed_alias and _torchao_export_supported(): + try: + import unsloth.save as _us_t + torchao_info = _us_t._normalize_torchao_method(compressed_alias) + except Exception: + torchao_info = None + is_torchao = torchao_info is not None + is_compressed = compressed_alias is not None and not is_torchao try: - if _IS_MLX: - if is_compressed: - return False, "Compressed-tensors export is not supported on macOS/MLX.", None - mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" - elif is_compressed: + if _IS_MLX and (is_compressed or is_torchao): + return ( + False, + "Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. " + "Use 16-bit or GGUF.", + None, + ) + + if is_torchao: + # Portable torchao: no NVIDIA GPU, no calibration. + compressed_suffix = torchao_info[1] + + if is_compressed: + # compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed. + if not _has_nvidia_gpu(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other " + "hardware use the portable FP8/INT8 (torchao) formats or 16-bit.", + None, + ) if not _compressed_export_supported(): return ( False, - "Compressed-tensors (FP8/NVFP4) export requires an Unsloth build with " + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", None, ) - save_method = _COMPRESSED[format_type][0] + import unsloth.save as _us + + # Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models + # (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports + # through it when available; else fall back to the workspace 0.10.x path below. + _shadow_pp = None + try: + from utils.transformers_version import llmcompressor_shadow_pythonpath + _shadow_pp = llmcompressor_shadow_pythonpath() + except Exception as e: + logger.warning(f"llm-compressor-main shadow unavailable: {e}") + if _shadow_pp: + os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp + else: + # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its + # transformers ceiling, so fail fast for sidecar models; default-tier still works. + os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) + _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + return ( + False, + "FP8/FP4 compressed-tensors export is not available for this model: it " + f"runs under transformers {_tf_ver}, but the installed llm-compressor " + f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the " + "llm-compressor-main runtime could not be provisioned (offline or " + "UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.", + None, + ) + + try: + info = _us._normalize_compressed_method(compressed_alias) + except Exception as e: + return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + if info is None: + return ( + False, + f"'{compressed_alias}' is not a recognized compressed-tensors export.", + None, + ) + compressed_suffix = info[2] + + if _IS_MLX: + mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed or is_torchao: + save_method = compressed_alias elif format_type == "4-bit (FP4)": save_method = "merged_4bit_forced" elif self._audio_type == "whisper": @@ -464,10 +603,10 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - # Compressed export writes to the "-" sibling; report that as output. + # Compressed / torchao writes to the "-" sibling; report that as output. final_dir = ( - f"{save_directory}-{_COMPRESSED[format_type][1]}" - if is_compressed + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) else save_directory ) self._write_export_metadata(final_dir) @@ -507,10 +646,9 @@ class ExportBackend: token = hf_token, private = private, ) - elif is_compressed and output_path and Path(output_path).is_dir(): - # The compressed model was already built locally in output_path; upload it - # directly so we do not re-run the (expensive, OOM-prone) compression that - # push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time. + elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + # Already built in output_path; upload it directly instead of re-running the + # expensive quantization that push_to_hub_merged(save_method=...) would redo. hf_api = HfApi(token = hf_token) repo_id = PushToHubMixin._create_repo( PushToHubMixin, @@ -522,7 +660,7 @@ class ExportBackend: username = repo_id.split("/")[0], base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), model_type = getattr(self.current_model.config, "model_type", "llm"), - method = format_type, + method = compressed_alias or format_type, extra = "unsloth", ) ModelCard(content).push_to_hub( @@ -568,6 +706,8 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -686,7 +826,7 @@ class ExportBackend: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, @@ -697,7 +837,9 @@ class ExportBackend: Args: save_directory: Local directory to save model - quantization_method: GGUF quantization method (e.g., "Q4_K_M") + quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them + (e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single + model load (unsloth save_to_gguf loops internally). push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID hf_token: Hugging Face token @@ -705,11 +847,13 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - # Only forward imatrix_file to an unsloth build that accepts it; otherwise even a plain - # no-imatrix export would fail with an unexpected-keyword error against an older unsloth. + # Only forward imatrix_file to an unsloth build that accepts it, else older builds raise + # an unexpected-keyword error even for a plain no-imatrix export. if imatrix_file is not None and not _supports_kwarg( self.current_model.save_pretrained_gguf, "imatrix_file" ): @@ -724,8 +868,14 @@ class ExportBackend: output_path: Optional[str] = None model_tmp_to_cleanup: Optional[str] = None try: - # unsloth expects lowercase quant method - quant_method = quantization_method.lower() + # Normalize to a lowercased list so multiple quants come from one model load. + if isinstance(quantization_method, (list, tuple)): + quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + else: + quant_methods = [str(quantization_method).lower()] + if not quant_methods: + quant_methods = ["q4_k_m"] + quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0] # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it # can't drift past the pinned llama-quantize binary's gguf API. @@ -847,7 +997,7 @@ class ExportBackend: return ( True, - f"GGUF model exported successfully ({quantization_method})", + f"GGUF model exported successfully ({', '.join(quant_methods)})", output_path, ) @@ -867,19 +1017,56 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). + Args: + gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp + convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`. + gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32. + Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: return False, "This is not a PEFT model. No adapter to export.", None + _GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32") + if gguf: + if _IS_MLX: + return ( + False, + "GGUF LoRA adapter export is not supported on macOS/MLX. " + "Use the safetensors adapter instead.", + None, + ) + outtype = str(gguf_outtype).lower() + if outtype not in _GGUF_LORA_OUTTYPES: + return ( + False, + f"Invalid GGUF LoRA outtype '{gguf_outtype}'. " + f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.", + None, + ) + # getattr so an older build without save_pretrained_gguf returns a clean message + # instead of an AttributeError (a generic 500). + _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) + if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + return ( + False, + "This Unsloth build does not support GGUF LoRA adapter export. " + "Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.", + None, + ) + output_path: Optional[str] = None try: if save_directory: @@ -887,7 +1074,24 @@ class ExportBackend: logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) - if _IS_MLX: + if gguf: + # Writes the adapter files plus "-lora-.gguf". + _apply_wsl_sudo_patch() + self.current_model.save_pretrained_gguf( + save_directory, + self.current_tokenizer, + save_method = "lora", + quantization_method = outtype, + # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. + token = hf_token or None, + ) + final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + logger.info( + "LoRA GGUF export complete. Files in %s:\n %s", + save_directory, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + elif _IS_MLX: # MLX: save adapters.safetensors + tokenizer files self.current_model.save_lora_adapters(save_directory) self.current_tokenizer.save_pretrained(save_directory) @@ -907,7 +1111,24 @@ class ExportBackend: logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") - if _IS_MLX: + if gguf: + # Upload the locally-built GGUF folder; needs a local save_directory so the + # conversion is not re-run. + if not (output_path and Path(output_path).is_dir()): + return ( + False, + "GGUF LoRA Hub upload requires a local save directory; set one and " + "retry.", + None, + ) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + elif _IS_MLX: with tempfile.TemporaryDirectory() as tmp_dir: self.current_model.save_lora_adapters(tmp_dir) self.current_tokenizer.save_pretrained(tmp_dir) diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 052a47dd80..671ef363f5 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -456,6 +456,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( @@ -467,6 +468,7 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "compressed_method": compressed_method, }, ) @@ -495,13 +497,13 @@ class ExportOrchestrator: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: - """Export model in GGUF format.""" + """Export model in GGUF format. `quantization_method` may be a single method or a list.""" return self._run_export( "gguf", { @@ -521,8 +523,10 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: - """Export LoRA adapter only.""" + """Export LoRA adapter only (optionally also as a GGUF LoRA file).""" return self._run_export( "lora", { @@ -531,6 +535,8 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, }, ) @@ -557,9 +563,13 @@ class ExportOrchestrator: cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) + # GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them + # all in one op off a single merge, so scale the timeout by the quant count. + _qm = params.get("quantization_method") + _n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1 resp = self._wait_response( f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min + timeout = 3600 * max(1, _n), ) op_success = resp.get("success", False) op_message = resp.get("message", "") diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index d473dcb54f..7828116236 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -397,6 +397,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + compressed_method = cmd.get("compressed_method"), ) elif export_type == "base": success, message, output_path = backend.export_base_model( @@ -423,6 +424,8 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + gguf = cmd.get("gguf", False), + gguf_outtype = cmd.get("gguf_outtype", "q8_0"), ) else: success, message = False, f"Unknown export type: {export_type}" diff --git a/studio/backend/main.py b/studio/backend/main.py index 0613a5ae53..8762c43195 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1145,7 +1145,7 @@ def get_system_info(current_subject: str = Depends(get_current_subject)): import os import time import logging - from utils.hardware import get_device + from utils.hardware import get_device, export_capability from utils.hardware.hardware import _backend_label logger = logging.getLogger(__name__) @@ -1218,6 +1218,8 @@ def get_system_info(current_subject: str = Depends(get_current_subject)): }, "gpu": gpu_info, "ml_packages": ml_packages, + # Export capability + torch-aware reason. See /api/system/hardware. + **export_capability(), } @@ -1240,11 +1242,13 @@ def get_hardware_info( method auto-selection. Sync def (not async): hardware/detail probes can shell out, and FastAPI runs sync endpoints in a threadpool. """ - from utils.hardware import get_gpu_summary, get_package_versions + from utils.hardware import get_gpu_summary, get_package_versions, export_capability body = { "gpu": get_gpu_summary(), "versions": get_package_versions(), + # Export capability + torch-aware reason; the Export UI grays out with the message. + **export_capability(), } if include_details: from utils.llama_cpp_update import get_installed_llama_version diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 7e05373f11..9dc4d9451a 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -6,7 +6,7 @@ from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator -from typing import List, Optional, Literal, Dict, Any +from typing import List, Optional, Literal, Dict, Any, Union def _validate_save_directory(value: str) -> str: @@ -168,6 +168,15 @@ class ExportMergedModelRequest(ExportCommonOptions): description = "Export precision / format for the merged model. The compressed-tensors " "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", ) + compressed_method: Optional[str] = Field( + None, + description = "Optional quantized-export alias. Either a compressed-tensors scheme " + "(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) " + "from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias " + "('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. " + "When set, it overrides format_type. Lets the export UI expose the full set of formats " + "beyond the quick buttons.", + ) class ExportBaseModelRequest(ExportCommonOptions): @@ -189,9 +198,10 @@ class ExportGGUFRequest(BaseModel): def _check_save_directory(cls, v): return _validate_save_directory(v) - quantization_method: str = Field( + quantization_method: Union[str, List[str]] = Field( "Q4_K_M", - description = 'GGUF quantization method (e.g. "Q4_K_M")', + description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list ' + '(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.', ) push_to_hub: bool = Field( False, @@ -219,4 +229,13 @@ class ExportGGUFRequest(BaseModel): class ExportLoRAAdapterRequest(ExportCommonOptions): """Request for exporting only the LoRA adapter (not merged).""" - # Uses fields from ExportCommonOptions only + gguf: bool = Field( + False, + description = "If True, also convert the adapter to a GGUF LoRA file " + "(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.", + ) + gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field( + "q8_0", + description = "GGUF LoRA output float type (only used when gguf=True). " + "Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).", + ) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index cf2cb2fa70..a7fd7cbec7 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -46,6 +46,23 @@ router = APIRouter() logger = get_logger(__name__) +def _ensure_export_supported() -> None: + """Reject a mutating export request up front (HTTP 400) when the host can't export. + + Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints + (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + """ + from utils.hardware import export_capability + + cap = export_capability() + if not cap.get("export_supported", True): + raise HTTPException( + status_code = 400, + detail = cap.get("export_unsupported_message") + or "Export is not supported on this platform.", + ) + + @router.post("/load-checkpoint", response_model = ExportOperationResponse) async def load_checkpoint( request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject) @@ -58,6 +75,7 @@ async def load_checkpoint( a clear error instead of tearing down the user's other running workloads. """ try: + _ensure_export_supported() backend = get_export_backend() # Run in a worker thread (spawns and waits on a subprocess, can take # minutes) so the event loop stays free to serve the live log SSE stream. @@ -266,6 +284,7 @@ async def export_merged_model( Wraps ExportBackend.export_merged_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_merged_model, @@ -275,6 +294,7 @@ async def export_merged_model( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + compressed_method = request.compressed_method, ) if not success: @@ -304,6 +324,7 @@ async def export_base_model( Wraps ExportBackend.export_base_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_base_model, @@ -342,6 +363,7 @@ async def export_gguf( Wraps ExportBackend.export_gguf. """ try: + _ensure_export_supported() backend = get_export_backend() # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. imatrix_file = request.imatrix_path or (True if request.imatrix else None) @@ -382,6 +404,7 @@ async def export_lora_adapter( Wraps ExportBackend.export_lora_adapter. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_lora_adapter, @@ -390,6 +413,8 @@ async def export_lora_adapter( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + gguf = request.gguf, + gguf_outtype = request.gguf_outtype, ) if not success: diff --git a/studio/backend/tests/test_export_capability.py b/studio/backend/tests/test_export_capability.py new file mode 100644 index 0000000000..e04417f933 --- /dev/null +++ b/studio/backend/tests/test_export_capability.py @@ -0,0 +1,156 @@ +# 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 export capability gating. + +Export is supported iff ``get_device() in {CUDA, XPU, MLX}``, with a torch-aware reason otherwise +(pytorch_not_installed / no_accelerator / mlx_unavailable), and the backend must import without +PyTorch. The matrix mocks the hardware probes; wiring is checked with ast so it runs on CPU. +""" + +import ast +import builtins +from pathlib import Path + +import pytest + +import utils.hardware.hardware as hw + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- capability matrix -------------------------------------------------------------------------- + + +def _patch(monkeypatch, *, torch: bool, device, apple: bool): + monkeypatch.setattr(hw, "_has_torch", lambda: torch) + monkeypatch.setattr(hw, "get_device", lambda: device) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: apple) + + +def test_cpu_with_torch_unsupported_no_accelerator(monkeypatch): + # PyTorch present but no accelerator: unsupported with no_accelerator, not "PyTorch missing". + _patch(monkeypatch, torch = True, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "no_accelerator" + assert "accelerator" in cap["export_unsupported_message"].lower() + # Must NOT tell a user with PyTorch installed to install PyTorch. + assert "PyTorch is not installed" not in cap["export_unsupported_message"] + + +def test_cuda_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.CUDA, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is True + assert cap["export_unsupported_reason"] is None + assert cap["export_unsupported_message"] is None + + +def test_xpu_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.XPU, apple = False) + assert hw.export_capability()["export_supported"] is True + + +def test_mlx_without_torch_supports_export(monkeypatch): + # Apple Silicon MLX exports without PyTorch. + _patch(monkeypatch, torch = False, device = hw.DeviceType.MLX, apple = True) + assert hw.export_capability()["export_supported"] is True + + +def test_no_torch_non_apple_reports_pytorch_missing(monkeypatch): + _patch(monkeypatch, torch = False, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "pytorch_not_installed" + assert "PyTorch is not installed" in cap["export_unsupported_message"] + + +def test_apple_without_mlx_reports_mlx_unavailable(monkeypatch): + # Apple + CPU means the MLX stack is missing; reason is mlx_unavailable regardless of torch. + for has_torch in (False, True): + _patch(monkeypatch, torch = has_torch, device = hw.DeviceType.CPU, apple = True) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "mlx_unavailable" + assert "MLX" in cap["export_unsupported_message"] + + +# -- import safety without PyTorch -------------------------------------------------------------- + + +def test_export_backend_imports_without_torch(monkeypatch): + """core/export/export.py must import on a --no-torch host (unsloth/torch blocked) and return a + clean 'PyTorch is not installed' message from an export attempt, not crash at import.""" + import importlib + import sys + + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + top = name.split(".")[0] + if top in {"torch", "unsloth"}: + raise ImportError(f"simulated: {top} not installed") + return real_import(name, *args, **kwargs) + + # Drop any preloaded copies so the guarded import paths re-run under the block. + for m in [k for k in sys.modules if k.split(".")[0] in {"torch", "unsloth"}]: + monkeypatch.delitem(sys.modules, m, raising = False) + monkeypatch.delitem(sys.modules, "core.export.export", raising = False) + monkeypatch.setattr(builtins, "__import__", blocking_import) + + mod = importlib.import_module("core.export.export") + assert mod._IS_MLX is False + assert mod.torch is None + assert mod._export_runtime_available() is False + + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = None + be.current_tokenizer = None + be.is_peft = False + be._audio_type = None + ok, message, out = be.export_merged_model("/tmp/does-not-matter") + assert ok is False + assert "PyTorch is not installed" in message + + +# -- endpoint / backend wiring (ast) ------------------------------------------------------------ + + +def test_main_endpoints_expose_export_capability(): + m = _src("main.py") + # Both system endpoints spread export_capability() into their response. + assert m.count("**export_capability()") >= 2 + assert '"/api/system/hardware"' in m and '"/api/system"' in m + + +def test_routes_guard_mutating_endpoints(): + r = _src("routes/export.py") + assert "def _ensure_export_supported()" in r + # load + all four export endpoints call the guard. + assert r.count("_ensure_export_supported()") >= 6 + + +def test_export_methods_check_runtime(): + e = _src("core/export/export.py") + assert "def _export_runtime_available()" in e + # Each export method returns the clear message when the runtime is missing. + assert e.count("_export_runtime_available()") >= 5 + assert "_PYTORCH_MISSING_MESSAGE" in e + + +def test_export_capability_reads_no_torch_helper(): + cap = _func_src("utils/hardware/hardware.py", "export_capability") + assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py index d914ff8651..f499390add 100644 --- a/studio/backend/tests/test_export_imatrix_compressed.py +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -54,8 +54,7 @@ def test_merged_request_rejects_unknown_format(): def test_export_gguf_threads_imatrix_to_save_and_push(): - # imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the - # conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword. + # imatrix_file must reach both save paths, but only via the conditional **imatrix_kw. g = _func_src("core/export/export.py", "export_gguf") assert g.count("**imatrix_kw") >= 2 assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g @@ -109,8 +108,139 @@ def test_export_merged_maps_compressed_to_save_method(): def test_compressed_hub_push_uploads_local_dir_without_recompressing(): - # A compressed Hub push must upload the already-built output_path, not re-run compression - # via push_to_hub_merged (which would compress a second time). + # A compressed / torchao Hub push must upload the built output_path, not re-quantize. m = _func_src("core/export/export.py", "export_merged_model") - assert "elif is_compressed and output_path and Path(output_path).is_dir():" in m + assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m + + +# -- torchao portable FP8/INT8 (device-agnostic, no NVIDIA GPU) --------------------------------- + + +def test_merged_request_accepts_torchao_aliases(): + # Portable torchao aliases pass through compressed_method (validated in the backend registry). + for alias in ("torchao_fp8", "torchao_int8"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_routes_torchao_and_skips_nvidia_guard(): + m = _func_src("core/export/export.py", "export_merged_model") + # torchao is classified separately and its suffix comes from the torchao normalizer. + assert "_normalize_torchao_method(compressed_alias)" in m + assert "is_torchao = torchao_info is not None" in m + assert "is_compressed = compressed_alias is not None and not is_torchao" in m + # The NVIDIA guard applies to compressed-tensors only, not torchao. + assert "_has_nvidia_gpu()" in m + # torchao routes through save_method just like compressed. + assert "elif is_compressed or is_torchao:" in m + + +def test_export_merged_nvidia_guard_present(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "requires an NVIDIA GPU" in m + + +def test_has_nvidia_gpu_helper_reads_hardware_module(): + h = _func_src("core/export/export.py", "_has_nvidia_gpu") + assert "DeviceType.CUDA" in h and "IS_ROCM" in h + + +def test_export_merged_relaxes_is_peft_guard(): + # Non-PEFT (Local/HF base) models can now export merged; the old hard block must be gone. + m = _func_src("core/export/export.py", "export_merged_model") + assert "Use 'Export Base Model' instead." not in m + + +def test_unsloth_save_has_torchao_registry_and_path(): + # Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth. + save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8") + assert "def _normalize_torchao_method" in save_py + assert "def _unsloth_save_torchao" in save_py + assert "TORCHAO_EXPORT_SCHEMES = {" in save_py + # torchao aliases must map to (scheme, suffix) so the backend routes to the torchao path. + assert '"torchao_fp8": ("fp8", "torchao-fp8")' in save_py + assert '"torchao_int8": ("int8", "torchao-int8")' in save_py + + +# -- GGUF multi-quant list ---------------------------------------------------------------------- + + +def test_gguf_request_accepts_list_of_quants(): + r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"]) + assert r.quantization_method == ["Q4_K_M", "Q8_0"] + r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M") + assert r2.quantization_method == "Q4_K_M" + + +def test_export_gguf_normalizes_quant_list(): + g = _func_src("core/export/export.py", "export_gguf") + assert "isinstance(quantization_method, (list, tuple))" in g + assert "quant_methods" in g + + +# -- GGUF LoRA adapter export ------------------------------------------------------------------- + + +def test_lora_request_has_gguf_fields(): + from models.export import ExportLoRAAdapterRequest + + r = ExportLoRAAdapterRequest(save_directory = "/tmp/x") + assert r.gguf is False and r.gguf_outtype == "q8_0" + r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0") + assert r2.gguf is True and r2.gguf_outtype == "q8_0" + + +def test_lora_request_rejects_bad_outtype(): + from models.export import ExportLoRAAdapterRequest + with pytest.raises(ValidationError): + ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf_outtype = "q3_k") + + +def test_export_lora_wires_gguf_save_method(): + la = _func_src("core/export/export.py", "export_lora_adapter") + assert 'save_method = "lora"' in la + assert "quantization_method = outtype" in la + + +def test_orchestrator_and_worker_pass_lora_gguf(): + o = _func_src("core/export/orchestrator.py", "export_lora_adapter") + assert '"gguf": gguf' in o and '"gguf_outtype": gguf_outtype' in o + w = _src("core/export/worker.py") + assert 'gguf = cmd.get("gguf", False)' in w + assert 'gguf_outtype = cmd.get("gguf_outtype", "q8_0")' in w + + +def test_route_passes_lora_gguf(): + r = _src("routes/export.py") + assert "gguf = request.gguf" in r and "gguf_outtype = request.gguf_outtype" in r + + +# -- compressed_method ("all formats" dropdown) ------------------------------------------------- + + +def test_merged_request_accepts_compressed_method(): + # Defaults to None; any scheme alias is accepted (validation happens in the backend registry). + assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None + for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_resolves_alias_via_registry(): + # The scheme + suffix must come from unsloth.save's registry normalizer, not a hardcoded dict. + m = _func_src("core/export/export.py", "export_merged_model") + assert "compressed_method" in m + assert "_normalize_compressed_method(compressed_alias)" in m + assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m + + +def test_orchestrator_and_worker_pass_compressed_method(): + o = _func_src("core/export/orchestrator.py", "export_merged_model") + assert "compressed_method" in o and '"compressed_method": compressed_method' in o + assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py") + + +def test_route_passes_compressed_method(): + assert "compressed_method = request.compressed_method" in _src("routes/export.py") diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 5f2b2abbcf..62b537fbac 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -44,6 +44,12 @@ from .vram_estimation import ( estimate_training_vram, ) + +def export_capability() -> dict: + """Return live export capability from the hardware module.""" + return _hardware.export_capability() + + __all__ = [ "DeviceType", "DEVICE", @@ -51,6 +57,7 @@ __all__ = [ "IS_ROCM", "detect_hardware", "get_device", + "export_capability", "is_apple_silicon", "clear_gpu_cache", "get_gpu_memory_info", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index cde7070075..8d6c919ebd 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -263,6 +263,49 @@ def get_device() -> DeviceType: return DEVICE +def export_capability() -> dict: + """Whether model export can run here, with a torch-aware reason when it cannot. + + Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at + import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The + reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch. + + Returns {export_supported, export_unsupported_reason, export_unsupported_message}. + """ + if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): + return { + "export_supported": True, + "export_unsupported_reason": None, + "export_unsupported_message": None, + } + # No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch" + # would be wrong advice on a Mac even when torch is also absent. + if is_apple_silicon(): + reason = "mlx_unavailable" + message = ( + "Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run " + "`unsloth studio update` to restore MLX and enable export." + ) + elif not _has_torch(): + reason = "pytorch_not_installed" + message = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + else: + reason = "no_accelerator" + message = ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export " + "on CPU only.)" + ) + return { + "export_supported": False, + "export_unsupported_reason": reason, + "export_unsupported_message": message, + } + + def clear_gpu_cache(): """ Clear GPU memory cache for the current device. diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index ac0ecfbfcd..2a63caa5b2 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -226,6 +226,11 @@ _VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR +# llm-compressor-main shadow for FP8/FP4 export of newer-transformers models. Like the .venv_t5_* +# sidecars but also shadows llm-compressor main + compressed-tensors; installed --no-deps so it +# reuses the workspace torch (torch-agnostic). +_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") + # Tier precedence: higher rank wins in _higher_tier. _TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} @@ -1518,6 +1523,152 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- +# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize +# Qwen3.5 / Gemma-4 / Llama. +_LLMC_MAIN_TRANSFORMERS = "5.10.2" +_LLMC_MAIN_SHA = "973c9c539a84dd9efaf74e115ede5ca419704c18" +_LLMC_MAIN_COMPRESSED_TENSORS = "0.17.2a20260702" +# Installed --no-deps (torch untouched); the full runtime set llm-compressor main needs, pinned. +_VENV_LLMCOMPRESSOR_SPECS = ( + f"transformers=={_LLMC_MAIN_TRANSFORMERS}", + f"llmcompressor @ git+https://github.com/vllm-project/llm-compressor@{_LLMC_MAIN_SHA}", + f"compressed-tensors=={_LLMC_MAIN_COMPRESSED_TENSORS}", + "huggingface-hub==1.21.0", + "hf-xet==1.5.1", + "tokenizers==0.22.2", + "safetensors==0.8.0", + "accelerate==1.14.0", + "datasets==5.0.0", + "pydantic==2.13.4", + "pydantic-core==2.46.4", + "typing-inspection==0.4.2", + "loguru==0.7.3", + "pyyaml==6.0.3", + "nvidia-ml-py==13.610.43", + "pillow==12.3.0", + "auto-round==0.13.1", + "regex==2026.6.28", +) +# Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes. +_LLMC_SHADOW_FINGERPRINT = ( + f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1" +) +_LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint" + + +def _llmcompressor_main_disabled() -> bool: + """True if the operator forbids the llm-compressor-main shadow (air-gapped / locked-down).""" + return os.environ.get("UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def _llmcompressor_shadow_is_valid() -> bool: + """True if the shadow dir exists with a marker matching the current pin fingerprint.""" + marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER + try: + return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT + except Exception: + return False + + +def _ensure_venv_llmcompressor_exists() -> bool: + """Ensure .venv_llmcompressor/ has the pinned llm-compressor-main stack. Install if missing. + + All specs are installed with --no-deps into a --target dir (mirrors the transformers sidecars), + so the workspace torch is never touched. Returns True on success. + """ + if _llmcompressor_shadow_is_valid(): + return True + if _llmcompressor_main_disabled(): + logger.warning( + "llm-compressor-main shadow needed but UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN is set; " + "compressed export of newer-transformers models will fail fast." + ) + return False + if _env_offline(): + logger.warning( + "llm-compressor-main shadow missing and HF/offline mode is set; cannot provision it." + ) + return False + + logger.warning( + "Provisioning llm-compressor-main shadow at %s (one-time, ~a few hundred MB, no torch) ...", + _VENV_LLMCOMPRESSOR_DIR, + ) + shutil.rmtree(_VENV_LLMCOMPRESSOR_DIR, ignore_errors = True) + os.makedirs(_VENV_LLMCOMPRESSOR_DIR, exist_ok = True) + + # Prefer uv (faster) then pip; install every spec at once, --no-deps, prereleases allowed + # (compressed-tensors ships as a pre-release). + base = [ + "--target", + _VENV_LLMCOMPRESSOR_DIR, + "--no-deps", + "--prerelease=allow", + *_VENV_LLMCOMPRESSOR_SPECS, + ] + cmds = [] + if shutil.which("uv"): + cmds.append(["uv", "pip", "install", "--python", sys.executable, *base]) + cmds.append( + [ + sys.executable, + "-m", + "pip", + "install", + *[a for a in base if a != "--prerelease=allow"], + "--pre", + ] + ) + + last_out = "" + for cmd in cmds: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = child_env_without_native_path_secret(), + **_windows_hidden_subprocess_kwargs(), + ) + last_out = result.stdout or "" + if result.returncode == 0: + try: + (Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text( + _LLMC_SHADOW_FINGERPRINT + ) + except Exception: + pass + logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR) + return True + logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0]) + + logger.error( + "Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s", + _LLMC_MAIN_SHA, + last_out[-4000:], + ) + return False + + +def llmcompressor_shadow_pythonpath() -> str | None: + """Provision (lazily) the llm-compressor-main shadow and return its sys.path entry, or None. + + Returns None when the shadow is disabled (UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN), offline, or + provisioning failed - callers then fall back to the fail-fast path. + """ + if _llmcompressor_main_disabled(): + return None + if _ensure_venv_llmcompressor_exists(): + return _VENV_LLMCOMPRESSOR_DIR + return None + + def _activate_venv(venv_dir: str, label: str) -> None: """Prepend *venv_dir* to sys.path, purge stale modules, reimport.""" if venv_dir not in sys.path: diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index e5fa6f0191..8c6ddd197a 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -70,6 +70,9 @@ const CHAT_ONLY_ALLOWED = new Set([ "/login", "/signup", "/change-password", + // Export stays reachable on chat-only hosts so the page can show its own grayed-out reason + // instead of a silent redirect; it self-gates via export capability, so nothing runs. + "/export", ]); function isChatOnlyAllowed(pathname: string): boolean { diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index ef4178ea42..fb1a1fc9c7 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -290,13 +290,12 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason); - // When Train/Export are greyed out (chat-only host), explain why on hover - // instead of disabling them silently. mlx_unavailable is the common macOS case - // after a reinstall/update dropped MLX and is recoverable via `unsloth studio update`. - const trainExportDisabledHint: string | undefined = !chatOnly + // Explain a greyed-out Train (chat-only host) on hover instead of disabling silently. Export is + // no longer disabled here: it stays navigable so its page can show a precise grayed-out reason. + const trainDisabledHint: string | undefined = !chatOnly ? undefined : chatOnlyReason === "mlx_unavailable" - ? "Training needs MLX. Run `unsloth studio update` to enable Train and Export." + ? "Training needs MLX. Run `unsloth studio update` to enable Train." : chatOnlyReason === "intel_mac" ? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only." : chatOnlyReason === "no_gpu" @@ -1206,7 +1205,7 @@ export function AppSidebar() { pathname === "/studio" || pathname.startsWith("/studio/") } disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1235,7 +1234,7 @@ export function AppSidebar() { label={t("shell.navigation.train")} active={pathname === "/studio" || pathname.startsWith("/studio/")} disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1256,11 +1255,8 @@ export function AppSidebar() { icon={DownloadSquare01Icon} label={t("shell.navigation.export")} active={pathname === "/export" || pathname.startsWith("/export/")} - disabled={chatOnly} - tooltip={trainExportDisabledHint} spinner={exportInProgress} onClick={() => { - if (chatOnly) return; navigate({ to: "/export" }); closeMobileIfOpen(); }} diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index d1b4e88a6d..be9767a24f 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -127,6 +127,8 @@ export async function loadCheckpoint(params: { export async function exportMerged(params: { save_directory: string; format_type?: string; + /** Compressed-tensors scheme alias (e.g. "fp8", "w4a16", "mxfp4"); overrides format_type. */ + compressed_method?: string | null; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -158,7 +160,8 @@ export async function exportBase(params: { export async function exportGGUF(params: { save_directory: string; - quantization_method: string; + /** A single GGUF quant method or a list (list produces multiple GGUFs from one model load). */ + quantization_method: string | string[]; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -179,6 +182,10 @@ export async function exportLoRA(params: { repo_id?: string | null; hf_token?: string | null; private?: boolean; + /** Also convert the adapter to a GGUF LoRA file (llama.cpp `--lora`). */ + gguf?: boolean; + /** GGUF LoRA output float type (f32/f16/bf16/q8_0/auto); only used when gguf=true. */ + gguf_outtype?: string; }): Promise { const response = await authFetch("/api/export/export/lora", { method: "POST", diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 9e0edf7303..29ba5a703b 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -28,7 +28,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { EXPORT_METHODS, type ExportMethod } from "../constants"; +import { + EXPORT_METHODS, + type ExportMethod, + findMergedFormat, +} from "../constants"; import type { ExportLogEntry } from "../api/export-api"; import { getExportLogLineClass } from "../lib/log-style"; import { @@ -200,6 +204,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) { const summaryMethodLabel = summary?.methodLabel ?? methodTitle; const summaryQuants = summary?.quantLevels ?? quantLevels; const summaryMethod = summary?.method ?? exportMethod; + const summaryFormats = (summary?.mergedFormats ?? []).map( + (v) => findMergedFormat(v)?.label ?? v, + ); const showProgress = isExporting || isTerminal; return ( @@ -392,14 +399,32 @@ export function ExportRunPanel(props: ExportRunPanelProps) { ? "Export finished and pushed to Hugging Face Hub." : "Export finished successfully."} - {run.result?.outputPath ? ( - - {run.result.outputPath} - - ) : null} + {(() => { + // List every folder written; a multi-format merged run created one per format. + const paths = run.result?.outputPaths ?? []; + const items = + paths.length > 0 + ? paths + : run.result?.outputPath + ? [{ label: "", path: run.result.outputPath }] + : []; + const showLabels = items.length > 1; + return items.map((o, i) => ( +
+ {showLabels && o.label ? ( + + {o.label} + + ) : null} + + {o.path} + +
+ )); + })()} )} @@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) { Export Method {summaryMethodLabel} + {summaryMethod === "merged" && summaryFormats.length > 0 && ( +
+ Formats + + {summaryFormats.join(", ")} + +
+ )} {summaryMethod === "gguf" && summaryQuants.length > 0 && (
Quantizations diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index 058de1edc4..be9080ba14 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -55,34 +55,172 @@ export const QUANT_OPTIONS: { { value: "f16", label: "F16" }, ]; -/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */ -export type MergedFormat = - | "16-bit (FP16)" - | "FP8 (compressed-tensors)" - | "NVFP4 (compressed-tensors)"; +/** + * Merged-export precision formats, sorted by bit width. Three backends: + * - "plain": standard save (16-bit); `formatType` is the backend `format_type`. + * - "compressed": llm-compressor compressed-tensors (vLLM), NVIDIA-only; `value` is the alias. + * - "torchao": portable FP8/INT8, no NVIDIA GPU needed; `value` is the alias. + * `common` entries are quick pills, the rest the "More formats" dropdown; `needsNvidia` entries + * are hidden on non-NVIDIA hardware. + */ +export type MergedBackend = "plain" | "compressed" | "torchao"; -export const MERGED_FORMATS: { - value: MergedFormat; +export type MergedFormatOption = { + value: string; label: string; + bits: number; + backend: MergedBackend; + group: string; + common: boolean; + needsNvidia: boolean; + needsCalibration?: boolean; hint: string; -}[] = [ + /** Backend `format_type` for a "plain" save (unused for compressed/torchao). */ + formatType?: string; +}; + +/** Kept as a string alias for back-compat with callers that typed the old union. */ +export type MergedFormat = string; + +export const MERGED_FORMATS: MergedFormatOption[] = [ + // 16-bit { - value: "16-bit (FP16)", + value: "16-bit", label: "16-bit", + bits: 16, + backend: "plain", + group: "16-bit", + common: true, + needsNvidia: false, hint: "Full precision, runs anywhere.", + formatType: "16-bit (FP16)", + }, + // 8-bit + { + value: "fp8", + label: "FP8", + bits: 8, + backend: "compressed", + group: "FP8", + common: true, + needsNvidia: true, + hint: "Dynamic per-token FP8 (W8A8) for vLLM. Data-free.", }, { - value: "FP8 (compressed-tensors)", - label: "FP8 (vLLM)", - hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.", + value: "torchao_fp8", + label: "FP8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic FP8 (torchao). Produces on any hardware; loads in vLLM.", }, { - value: "NVFP4 (compressed-tensors)", - label: "NVFP4 (vLLM)", - hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.", + value: "w8a8", + label: "INT8 (W8A8)", + bits: 8, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "8-bit weights and 8-bit activations for vLLM. Data-free.", + }, + { + value: "torchao_int8", + label: "INT8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic INT8 (torchao). Produces on any hardware; loads in vLLM.", + }, + { + value: "fp8_static", + label: "FP8 Static", + bits: 8, + backend: "compressed", + group: "FP8", + common: false, + needsNvidia: true, + needsCalibration: true, + hint: "Static per-tensor FP8. Calibrates on data.", + }, + { + value: "w8a16", + label: "INT8 (W8A16)", + bits: 8, + backend: "compressed", + group: "INT", + common: false, + needsNvidia: true, + hint: "8-bit weight-only. Data-free.", + }, + { + value: "mxfp8", + label: "MXFP8", + bits: 8, + backend: "compressed", + group: "MXFP", + common: false, + needsNvidia: true, + hint: "Microscaling FP8. Needs a newer compressed-tensors stack.", + }, + // 4-bit + { + value: "w4a16", + label: "INT4 (W4A16)", + bits: 4, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "4-bit weight-only (GPTQ-style) for vLLM. Data-free.", + }, + { + value: "mxfp4", + label: "MXFP4", + bits: 4, + backend: "compressed", + group: "MXFP", + common: true, + needsNvidia: true, + hint: "Microscaling FP4 (W4A4) for vLLM. Data-free.", + }, + { + value: "nvfp4", + label: "NVFP4", + bits: 4, + backend: "compressed", + group: "FP4", + common: true, + needsNvidia: true, + needsCalibration: true, + hint: "NVIDIA FP4 (W4A4) for vLLM. Calibrates on data.", }, ]; +/** Look up a merged format option by its stable value. */ +export function findMergedFormat(value: string): MergedFormatOption | undefined { + return MERGED_FORMATS.find((f) => f.value === value); +} + +/** Backend payload for one merged format: plain -> formatType, compressed/torchao -> the alias. */ +export function mergedFormatPayload(value: string): { + formatType: string; + compressedMethod: string | null; +} { + const opt = findMergedFormat(value); + if (!opt || opt.backend === "plain") { + return { + formatType: opt?.formatType ?? "16-bit (FP16)", + compressedMethod: null, + }; + } + return { formatType: "16-bit (FP16)", compressedMethod: opt.value }; +} + /** * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16. * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0 diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 7b5c6ec6d4..07606a26ed 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -24,6 +24,19 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { + DropdownMenu, + DropdownMenuCheckboxItem, + DropdownMenuContent, + DropdownMenuLabel, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { + Alert, + AlertDescription, + AlertTitle, +} from "@/components/ui/alert"; import { Separator } from "@/components/ui/separator"; import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -63,11 +76,14 @@ import { type ExportMethod, GUIDE_STEPS, MERGED_FORMATS, - type MergedFormat, + type MergedFormatOption, + mergedFormatPayload, QUANT_OPTIONS, buildQuantSizeLabels, getEstimatedSize, } from "./constants"; +import { useHardwareInfo } from "@/hooks/use-hardware-info"; +import { usePlatformStore } from "@/config/env"; import { isExportPanelActive, useExportRuntimeStore, @@ -78,6 +94,10 @@ import { exportTourSteps } from "./tour"; const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]); +// GGUF LoRA output float types (Q8_0 first / default). Q8_0 falls back to F16 per tensor for dims +// not divisible by the block size (32); no "auto" - the choice is explicit. +const LORA_GGUF_OUTTYPES = ["q8_0", "f16", "bf16", "f32"] as const; + type SourceTab = "local" | "checkpoint" | "hf"; type SourceMode = "checkpoint" | "model"; @@ -109,7 +129,16 @@ function buildRelativeSaveDirectory( : sourceBaseModelName; return `${safePathSegment(rawName)}-GGUF`; } - return `${selectedModelIdx ?? "model"}/${checkpoint}`; + // Merged / LoRA: a checkpoint keeps the "/" layout under outputs. + if (sourceMode === "checkpoint" && selectedModelIdx && checkpoint) { + return `${selectedModelIdx}/${checkpoint}`; + } + // Local / HF source (no checkpoint): name from the model id to avoid "model/null". + const rawName = + sourceMode === "checkpoint" + ? checkpoint ?? selectedModelIdx ?? sourceBaseModelName + : sourceBaseModelName; + return `${safePathSegment(rawName)}-${exportMethod === "lora" ? "adapter" : "merged"}`; } function siblingGgufDirectory(sourcePath: string): string | null { @@ -177,9 +206,57 @@ export function ExportPage() { }); // GGUF importance matrix (required for the IQ quants) and merged-export precision. const [useImatrix, setUseImatrix] = useState(false); - const [mergedFormat, setMergedFormat] = useState("16-bit (FP16)"); - // IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit - // an IQ quant with no imatrix and llama.cpp would reject it. + // Merged precision: one or more MERGED_FORMATS values, exported in one run. Seed from a live run + // so navigating away and back (which remounts this page) keeps the selection, like exportMethod. + const [selectedFormats, setSelectedFormats] = useState(() => { + const s = useExportRuntimeStore.getState(); + return isExportPanelActive(s) && + s.summary?.method === "merged" && + s.summary.mergedFormats.length > 0 + ? s.summary.mergedFormats + : ["16-bit"]; + }); + // LoRA-only export: optionally also emit a GGUF LoRA adapter, and its output float type. + const [loraAsGguf, setLoraAsGguf] = useState(false); + const [loraGgufOuttype, setLoraGgufOuttype] = useState("q8_0"); + // GGUF method: export the full model as GGUF quants, or (for an adapter checkpoint) a GGUF LoRA. + const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model"); + + const hardware = useHardwareInfo(); + // GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host. + const isMacHost = usePlatformStore((s) => s.deviceType) === "mac"; + // Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats. + const hasNvidia = hardware.cuda != null && hardware.rocm == null; + // Only gray out on an authoritative unsupported response; while unloaded the backend route guard + // stays authoritative. The backend supplies the precise reason; the fallback below is a backstop. + const exportUnsupported = + hardware.loaded && hardware.exportSupported === false; + const exportUnsupportedMessage = + hardware.exportUnsupportedMessage ?? + "Export requires a supported accelerator (NVIDIA, AMD, or Intel GPU, or Apple Silicon) with PyTorch or MLX installed."; + const availableFormats = useMemo( + () => + MERGED_FORMATS.filter((f) => { + // compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU. + if (f.backend === "compressed") return hasNvidia; + // Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a + // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the + // backend rejects quantized export there). + if (f.backend === "torchao") return !hasNvidia && !isMacHost; + // Plain 16-bit is available everywhere. + return true; + }), + [hasNvidia, isMacHost], + ); + const toggleFormat = useCallback((value: string) => { + setSelectedFormats((prev) => + prev.includes(value) + ? prev.filter((v) => v !== value) + : [...prev, value], + ); + }, []); + // availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed. + // IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it. const requiresImatrix = quantLevels.some( (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, ); @@ -304,6 +381,11 @@ export function ExportPage() { const baseModelName = selectedModelData?.base_model ?? "—"; const isAdapter = !!selectedModelData?.peft_type; const isQuantized = !!selectedModelData?.is_quantized; + // isAdapter / isQuantized come from the checkpoint's metadata and are stale in "model" source + // mode (a direct base export), so treat both as false outside checkpoint mode to avoid wrongly + // gating the methods. + const effectiveIsAdapter = sourceMode === "checkpoint" && isAdapter; + const effectiveIsQuantized = sourceMode === "checkpoint" && isQuantized; const loraRank = selectedModelData?.lora_rank ?? null; const trainingMethodLabel = selectedModelData?.peft_type ? "LoRA / QLoRA" @@ -416,25 +498,30 @@ export function ExportPage() { setCheckpoint(null); }, [selectedModelIdx]); - // For a ?run= deep link, default to the run's main checkpoint. Declared after - // the reset effect above so it runs last and isn't clobbered back to null. + // Default to the newest checkpoint when none is chosen (checkpoints are sorted newest-first). + // Declared after the reset effect above so it runs last and isn't clobbered back to null. Covers + // both a ?run= deep link and a plain finetune opened without an explicit checkpoint pick. useEffect(() => { - if (appliedRunRef.current == null) return; - if (appliedRunRef.current !== selectedModelIdx) return; + if (sourceMode !== "checkpoint") return; if (checkpoint != null || checkpointsForModel.length === 0) return; setCheckpoint(checkpointsForModel[0].display_name); - }, [selectedModelIdx, checkpoint, checkpointsForModel]); + }, [sourceMode, selectedModelIdx, checkpoint, checkpointsForModel]); // Auto-reset export method if incompatible with the selected model type useEffect(() => { - if (!isAdapter && (exportMethod === "merged" || exportMethod === "lora")) { + // Only LoRA needs a real adapter; Merged and GGUF work for non-PEFT base models too. + if (!effectiveIsAdapter && exportMethod === "lora") { setExportMethod(null); } // Quantized non-PEFT models can't export to any format - if (!isAdapter && isQuantized && exportMethod !== null) { + if (!effectiveIsAdapter && effectiveIsQuantized && exportMethod !== null) { setExportMethod(null); } - }, [isAdapter, isQuantized, exportMethod]); + // The GGUF LoRA target only applies to an adapter checkpoint on a non-Mac host. + if ((!effectiveIsAdapter || isMacHost) && ggufTarget !== "model") { + setGgufTarget("model"); + } + }, [effectiveIsAdapter, effectiveIsQuantized, exportMethod, isMacHost, ggufTarget]); const handleSourceTabChange = useCallback((next: string) => { if (next === "checkpoint") { @@ -442,7 +529,7 @@ export function ExportPage() { } else if (next === "hf" || next === "local") { setSourceMode("model"); setModelSource(next); - setExportMethod("gguf"); + // Don't force GGUF: Local / HF sources can export Merged too; a stale LoRA pick auto-resets. } else { return; } @@ -508,10 +595,26 @@ export function ExportPage() { sourceMode, ]); const saveDirectory = customSaveDirectory?.trim() || defaultSaveDirectory; + // Each merged format uploads a full model to the repo root, so several to one repo would collide. + // GGUF method exporting an adapter checkpoint as a GGUF LoRA (vs full-model quants). Reuses the + // LoRA-adapter export path; no quant list needed. + const ggufAsLora = + exportMethod === "gguf" && + ggufTarget === "lora" && + effectiveIsAdapter && + !isMacHost; + + // Restrict a Hub merged export to a single format; multi-format stays available for local export. + const hubMultiFormat = + destination === "hub" && exportMethod === "merged" && selectedFormats.length > 1; + const canExport = !!( selectedExportSource && exportMethod && - (exportMethod !== "gguf" || quantLevels.length > 0) + !exportUnsupported && + !hubMultiFormat && + (exportMethod !== "gguf" || ggufAsLora || quantLevels.length > 0) && + (exportMethod !== "merged" || selectedFormats.length > 0) ); const applyHfSourceModel = useCallback((value: string) => { @@ -576,9 +679,14 @@ export function ExportPage() { const handleStart = useCallback(async () => { const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; if (!source || !exportMethod) return; - // A GGUF export with no quant selected runs zero exports yet would still - // settle as success with no file; require at least one (mirrors canExport). - if (exportMethod === "gguf" && quantLevels.length === 0) return; + // No supported accelerator (or PyTorch/MLX missing): the backend would reject anyway; don't submit. + if (exportUnsupported) return; + // GGUF with no quant, or merged with no format, would run an unintended/empty export; require + // at least one (mirrors canExport, in case the panel's Start button bypasses the outer one). + if (exportMethod === "gguf" && !ggufAsLora && quantLevels.length === 0) return; + if (exportMethod === "merged" && selectedFormats.length === 0) return; + // A Hub merged push writes each format to the repo root; several would collide (mirrors canExport). + if (hubMultiFormat) return; const selectedCp = sourceMode === "checkpoint" ? checkpointsForModel.find((cp) => cp.display_name === checkpoint) @@ -591,8 +699,13 @@ export function ExportPage() { ? `${hfUsername}/${modelName}` : undefined; const token = pushToHub && hfToken ? hfToken : undefined; - const methodLabel = - EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod; + // The GGUF method with the LoRA target reuses the LoRA-adapter export path. + const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod; + const emitLoraGguf = + ggufAsLora || (effectiveMethod === "lora" && loraAsGguf && !isMacHost); + const methodLabel = ggufAsLora + ? "GGUF LoRA adapter" + : (EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod); const adapterExport = sourceMode === "checkpoint" && isAdapter; // Consent gate for an HF source's custom (auto_map) code, run before we hand @@ -624,11 +737,16 @@ export function ExportPage() { trustRemoteCode, approvedRemoteCodeFingerprint, loadToken: hfToken || null, - exportMethod, + exportMethod: effectiveMethod, isAdapter: adapterExport, quantLevels, useImatrix: effectiveImatrix, - mergedFormat, + mergedSelections: selectedFormats.map((v) => ({ + ...mergedFormatPayload(v), + label: MERGED_FORMATS.find((f) => f.value === v)?.label ?? v, + })), + loraGguf: emitLoraGguf, + loraGgufOuttype, saveDirectory, destination, repoId, @@ -639,8 +757,9 @@ export function ExportPage() { baseModelName: sourceBaseModelName, checkpointLabel: selectedExportSource, methodLabel, - method: exportMethod, + method: effectiveMethod, quantLevels, + mergedFormats: exportMethod === "merged" ? selectedFormats : [], destination, }, }); @@ -656,7 +775,13 @@ export function ExportPage() { isAdapter, quantLevels, effectiveImatrix, - mergedFormat, + selectedFormats, + hubMultiFormat, + ggufAsLora, + loraAsGguf, + isMacHost, + loraGgufOuttype, + exportUnsupported, destination, saveDirectory, hfUsername, @@ -1163,75 +1288,303 @@ export function ExportPage() {
+ {exportUnsupported && ( + + + Export unavailable + {exportUnsupportedMessage} + + )} + - {exportMethod === "merged" && isAdapter && ( -
-
Precision
-
- {MERGED_FORMATS.map((f) => ( - - ))} -
-
- {MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint} + {exportMethod === "merged" && !exportUnsupported && ( +
+
+
+
Precision
+ + — select one or more + +
+
+ {availableFormats + .filter((f) => f.common) + .map((f) => { + const active = selectedFormats.includes(f.value); + return ( + + ); + })} + + {availableFormats.some((f) => !f.common) && ( + + + + + + + Additional formats + + + {availableFormats + .filter((f) => !f.common) + .map((f) => ( + toggleFormat(f.value)} + onSelect={(e) => e.preventDefault()} + > + + + {f.label} + {f.needsCalibration ? " *" : ""} + + + {f.hint} + + + + ))} + + + )} +
+ + {selectedFormats.length > 0 && ( +
+ + {selectedFormats.length} selected:{" "} + {selectedFormats + .map( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.label ?? v, + ) + .join(", ")} + + {selectedFormats.length > 1 && ( + + )} +
+ )} + + {hubMultiFormat && ( +
+ Hub export supports one format at a time (each writes to + the repository root). Select a single format, or export + locally to produce several at once. +
+ )} + + {selectedFormats.some( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.needsCalibration, + ) && ( +
+ * calibrates on data (uses a small calibration set). +
+ )} + + {!hasNvidia && ( +
+ No NVIDIA GPU detected: compressed-tensors formats are + hidden. 16-bit and portable FP8/INT8 (torchao) still + work here and load in vLLM. +
+ )}
)} - {exportMethod === "gguf" && ( - <> - -
-
-
- Importance matrix (imatrix) + {exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && ( +
+
+
Adapter format
+
+ + +
+
+ {isMacHost + ? "GGUF LoRA is not available on macOS/MLX; exporting the safetensors adapter." + : loraAsGguf + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Standard PEFT adapter files. Pair with the base model at inference."} +
+
+ + {loraAsGguf && ( +
+
Output type
+ +
+ )} +
+ )} + + {exportMethod === "gguf" && !exportUnsupported && ( +
+ {effectiveIsAdapter && !isMacHost && ( +
+
Export target
+
+ +
- {requiresImatrix - ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." - : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} + {ggufTarget === "lora" + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Merges the adapter into the base model, then quantizes the full model to GGUF."}
- -
- + )} + + {ggufAsLora ? ( +
+
Output type
+ +
+ ) : ( + <> + +
+
+
+ Importance matrix (imatrix) +
+
+ {requiresImatrix + ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." + : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} +
+
+ +
+ + )} +
)} {estimatedSize && (
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts index a87e7d93e9..6ee5c78994 100644 --- a/studio/frontend/src/features/export/stores/export-runtime-store.ts +++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts @@ -5,7 +5,6 @@ import { create } from "zustand"; import { cancelExport, cleanupExport, - exportBase, exportGGUF, exportLoRA, exportMerged, @@ -119,6 +118,8 @@ export interface ExportRunSummary { methodLabel: string; method: ExportMethod; quantLevels: string[]; + /** Merged: the selected format values (for the summary "Formats" row and to reseed the picker). */ + mergedFormats: string[]; destination: ExportDestination; } @@ -140,8 +141,16 @@ export interface RunExportParams { quantLevels: string[]; /** GGUF: use an importance matrix (auto-download); required for the IQ quants. */ useImatrix?: boolean; - /** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */ - mergedFormat?: string; + /** Merged: precision formats, each exported to its own sibling directory. Defaults to 16-bit. + * `label` is the display name for the success banner's per-format output line. */ + mergedSelections?: { + formatType: string; + compressedMethod: string | null; + label: string; + }[]; + /** LoRA: also emit a GGUF LoRA adapter (llama.cpp `--lora`), and its output float type. */ + loraGguf?: boolean; + loraGgufOuttype?: string; saveDirectory: string; destination: ExportDestination; repoId?: string; @@ -172,7 +181,13 @@ interface ExportRuntimeState { * settling the run by polling /api/export/status instead. Logs keep streaming. */ reconnecting: boolean; startedAt: number | null; - result: { outputPath: string | null; destination: ExportDestination } | null; + /** `outputPath` is the first path (back-compat); `outputPaths` is one entry per written folder + * so a multi-format merged run can list every sibling directory it created. */ + result: { + outputPath: string | null; + outputPaths: { label: string; path: string }[]; + destination: ExportDestination; + } | null; error: string | null; cancelRequested: boolean; hasHydrated: boolean; @@ -317,6 +332,10 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success" as const, result: { outputPath: status.last_op_output_path ?? null, + // A run recovered from the backend only knows the last output path. + outputPaths: status.last_op_output_path + ? [{ label: "", path: status.last_op_output_path }] + : [], destination: state.result?.destination ?? "local", }, }; @@ -351,7 +370,9 @@ export const useExportRuntimeStore = create()((set, get) => const quantTotal = params.exportMethod === "gguf" ? Math.max(1, params.quantLevels.length) - : 1; + : params.exportMethod === "merged" + ? Math.max(1, params.mergedSelections?.length ?? 1) + : 1; set({ runId, @@ -431,67 +452,72 @@ export const useExportRuntimeStore = create()((set, get) => } if (!isCurrent()) return; - // 2. Run the export. Capture the resolved output_path for the success - // banner; multi-quant GGUF shares one directory, so keep the last. + // 2. Run the export. Collect every resolved output_path so the success + // banner can list each sibling directory a multi-format run created. set({ phase: "exporting" }); - let lastOutputPath: string | null = null; + const outputs: { label: string; path: string }[] = []; if (params.exportMethod === "merged") { - if (params.isAdapter) { + // Each selected format writes its own sibling directory (PEFT or non-PEFT base alike). + const selections = + params.mergedSelections && params.mergedSelections.length > 0 + ? params.mergedSelections + : [{ formatType: "16-bit (FP16)", compressedMethod: null, label: "16-bit" }]; + for (let i = 0; i < selections.length; i += 1) { + if (!isCurrent()) return; + set({ quantIndex: i }); + const sel = selections[i]; const { outputPath } = await runRecoverableOp(() => exportMerged({ save_directory: params.saveDirectory, - format_type: params.mergedFormat, + format_type: sel.formatType, + compressed_method: sel.compressedMethod, push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, private: params.privateRepo, }), ); - lastOutputPath = outputPath; - } else { - const { outputPath } = await runRecoverableOp(() => - exportBase({ - save_directory: params.saveDirectory, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - private: params.privateRepo, - base_model_id: params.baseModelId, - }), - ); - lastOutputPath = outputPath; - } - } else if (params.exportMethod === "gguf") { - for (let i = 0; i < params.quantLevels.length; i += 1) { - if (!isCurrent()) return; - set({ quantIndex: i }); - const quant = params.quantLevels[i]; - const { outputPath } = await runRecoverableOp(() => - exportGGUF({ - save_directory: params.saveDirectory, - quantization_method: quant, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - imatrix: params.useImatrix, - }), - ); - lastOutputPath = outputPath ?? lastOutputPath; + if (outputPath) outputs.push({ label: sel.label, path: outputPath }); if (!isCurrent()) return; set({ quantIndex: i + 1 }); } + } else if (params.exportMethod === "gguf") { + // Send the whole quant list in ONE call: the model is merged once and every GGUF comes + // from that single merge (unsloth save_to_gguf loops internally). + const { outputPath } = await runRecoverableOp(() => + exportGGUF({ + save_directory: params.saveDirectory, + quantization_method: params.quantLevels, + push_to_hub: pushToHub, + repo_id: params.repoId, + hf_token: params.token, + imatrix: params.useImatrix, + }), + ); + if (outputPath) outputs.push({ label: "GGUF", path: outputPath }); + if (!isCurrent()) return; + set({ quantIndex: get().quantTotal }); } else if (params.exportMethod === "lora") { const { outputPath } = await runRecoverableOp(() => exportLoRA({ save_directory: params.saveDirectory, push_to_hub: pushToHub, repo_id: params.repoId, - hf_token: params.token, + // A local GGUF LoRA export still reloads a possibly-gated base config, so fall back to + // the load token when there is no hub-upload token (both are the same HF token). + hf_token: params.token ?? params.loadToken ?? null, private: params.privateRepo, + gguf: params.loraGguf ?? false, + gguf_outtype: params.loraGgufOuttype ?? "q8_0", }), ); - lastOutputPath = outputPath; + if (outputPath) { + outputs.push({ + label: params.loraGguf ? "GGUF LoRA adapter" : "LoRA adapter", + path: outputPath, + }); + } } if (!isCurrent()) return; @@ -499,7 +525,11 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success", isExporting: false, reconnecting: false, - result: { outputPath: lastOutputPath, destination: params.destination }, + result: { + outputPath: outputs[0]?.path ?? null, + outputPaths: outputs, + destination: params.destination, + }, }); } catch (err) { if (!isCurrent()) return; diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 36e3c532f3..4d63d4d6af 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -25,6 +25,13 @@ export interface HardwareInfo { transformers: string | null; unsloth: string | null; llamaCpp: string | null; + // Whether export can run here (true only on a supported accelerator), with a torch-aware + // reason. `null` until the authoritative response lands, so callers don't briefly enable + // export; `loaded` flips true once a real (non-error) response arrives. + exportSupported: boolean | null; + exportUnsupportedReason: string | null; + exportUnsupportedMessage: string | null; + loaded: boolean; } const DEFAULT: HardwareInfo = { @@ -38,6 +45,10 @@ const DEFAULT: HardwareInfo = { transformers: null, unsloth: null, llamaCpp: null, + exportSupported: null, + exportUnsupportedReason: null, + exportUnsupportedMessage: null, + loaded: false, }; // Module-level cache so multiple components share one fetch. @@ -87,6 +98,10 @@ async function fetchOnce(): Promise { transformers: data?.versions?.transformers ?? null, unsloth: data?.versions?.unsloth ?? null, llamaCpp: data?.llama_cpp ?? null, + exportSupported: data?.export_supported ?? null, + exportUnsupportedReason: data?.export_unsupported_reason ?? null, + exportUnsupportedMessage: data?.export_unsupported_message ?? null, + loaded: true, }; if (generation === cacheGeneration) { cached = info; diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 304de7a9af..209a8a06f1 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -256,7 +256,7 @@ with sync_playwright() as p: composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) - # Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat. + # Detect chat-only mode (/api/health.chat_only): /studio redirects to /chat while /export stays reachable and self-gated. health_resp = evaluate_fetch( page, f"{BASE}/api/health", @@ -404,15 +404,19 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── # 3. Export route. # ───────────────────────────────────────────────────── - step(f"Export route ({'chat-only redirect' if chat_only else 'form fields'})") + step(f"Export route ({'chat-only self-gated' if chat_only else 'form fields'})") page.goto(f"{BASE}/export") page.wait_for_timeout(1500) shoot("07-export") if chat_only: - if "/export" in page.url: - soft_fail(f"chat-only mode should redirect /export -> /chat; url={page.url}") + if "/export" not in page.url: + soft_fail(f"chat-only mode should keep /export reachable; url={page.url}") else: - info(f"OK chat-only redirected /export -> {page.url}") + unavailable = page.get_by_text(re.compile(r"Export unavailable", re.I)).first + if unavailable.count() == 0: + soft_fail("chat-only /export did not show the export unavailable gate") + else: + info("OK chat-only /export rendered the unavailable gate") else: # Non-chat-only: verify the export-cta button + HF token field. cta = page.locator('[data-tour="export-cta"]').first diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index 66ebdd75da..f0a843c380 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -245,6 +245,10 @@ def main(): # expert even if the sample set does not route tokens to all of them. is_moe = _is_moe(getattr(model, "config", None)) ignore = ["lm_head"] + # Skip the same modules RedHatAI/NVIDIA skip for the Qwen3.5 / Qwen3-Next family (these also have + # shapes not divisible by the grouped-scheme group_size, which would otherwise error). No-ops + # elsewhere. Hybrid linear attention, VLM vision tower, and the MTP/speculative head. + ignore += ["re:.*\\.linear_attn\\..*", "re:.*\\.visual\\..*", "re:.*mtp.*"] if is_moe: # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate. ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"] diff --git a/unsloth/save.py b/unsloth/save.py index 226ca5fed8..50ae4119bd 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -205,6 +205,24 @@ COMPRESSED_EXPORT_SCHEMES = { } +# torchao "portable" quant export: device-agnostic FP8 / INT8, no NVIDIA GPU needed. +# alias -> (kind, sibling suffix). FP8 saves to safetensors, INT8 to .bin; both load in vLLM. +TORCHAO_EXPORT_SCHEMES = { + "torchao_fp8": ("fp8", "torchao-fp8"), + "torchao_int8": ("int8", "torchao-int8"), + "portable_fp8": ("fp8", "torchao-fp8"), + "portable_int8": ("int8", "torchao-int8"), +} + + +def _normalize_torchao_method(save_method): + """Return (kind, suffix) if `save_method` is a torchao portable FP8/INT8 export, else None.""" + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + return TORCHAO_EXPORT_SCHEMES.get(key) + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -215,6 +233,9 @@ def _normalize_compressed_method(save_method): if not isinstance(save_method, str): return None key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + # torchao aliases route to the torchao path, so skip them before the "fp8" near-miss check. + if key in TORCHAO_EXPORT_SCHEMES: + return None if key in COMPRESSED_EXPORT_SCHEMES: return COMPRESSED_EXPORT_SCHEMES[key] if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): @@ -1368,6 +1389,53 @@ def install_python_non_blocking(packages = []): # bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). _LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" +# Highest transformers release llm-compressor 0.10.x/0.12.x can run against (its metadata pins +# transformers<=4.57.6). Models that require a newer-transformers sidecar (e.g. Qwen3.5 needs +# transformers 5.3.0) cannot be quantized by llm-compressor at all: it imports +# transformers.modeling_utils.TORCH_INIT_FUNCTIONS, which was removed in transformers 5.x, so the +# compressed-export subprocess dies with a cryptic ImportError AFTER the expensive 16bit merge. +# Detect that up front and fail fast with an actionable message. Bump this in lockstep with a +# llm-compressor release that supports newer transformers. +_LLM_COMPRESSOR_MAX_TRANSFORMERS = "4.57.6" + + +def _transformers_exceeds_llm_compressor_ceiling(transformers_version = None): + """Return (exceeds, active_version) comparing the active transformers to the llm-compressor ceiling. + + `exceeds` is True only when we can parse both versions and the active transformers is strictly + newer than `_LLM_COMPRESSOR_MAX_TRANSFORMERS`. Any parse failure returns False (fail open) so a + real quantization attempt still surfaces the underlying error rather than a false positive. + """ + if transformers_version is None: + try: + import transformers as _tf + transformers_version = _tf.__version__ + except Exception: + return False, "unknown" + try: + from packaging.version import parse as _parse + + # Drop any local build suffix ("4.57.6+abc") so it does not skew the comparison. + active = _parse(str(transformers_version).split("+", 1)[0]) + ceiling = _parse(_LLM_COMPRESSOR_MAX_TRANSFORMERS) + return active > ceiling, str(transformers_version) + except Exception: + return False, str(transformers_version) + + +# A caller (e.g. Unsloth Studio) can enable FP8/FP4 export of newer-transformers models (Qwen3.5, +# Gemma-4, ...) by provisioning a dedicated llm-compressor-main "shadow" (transformers>=5.9 layered +# over the existing torch) and pointing us at its sys.path entry via this env var. When set, the +# quantization subprocess uses it instead of the workspace llm-compressor and the ceiling fail-fast +# is bypassed. +_COMPRESSED_QUANTIZE_PYTHONPATH_ENV = "UNSLOTH_COMPRESSED_QUANTIZE_PYTHONPATH" + + +def _compressed_quantize_pythonpath(): + """Return the llm-compressor-main shadow PYTHONPATH, or None if not set.""" + pp = os.environ.get(_COMPRESSED_QUANTIZE_PYTHONPATH_ENV, "").strip() + return pp or None + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. @@ -2023,10 +2091,40 @@ def unsloth_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -2107,6 +2205,37 @@ def unsloth_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -2114,6 +2243,7 @@ def unsloth_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3812,10 +3942,40 @@ def unsloth_generic_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3896,6 +4056,37 @@ def unsloth_generic_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -3903,6 +4094,7 @@ def unsloth_generic_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -4114,22 +4306,36 @@ def _unsloth_save_compressed_tensors( if not is_main_process: return None - # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported - # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. - install_llm_compressor() - if not _scheme_is_available(scheme): - try: - import transformers as _tf - tf_ver = _tf.__version__ - except Exception: - tf_ver = "unknown" - raise RuntimeError( - f"Unsloth: scheme '{scheme}' is not available in your installed " - f"compressed-tensors / llm-compressor.\n" - f"It requires a newer llm-compressor that needs transformers>=5.9 " - f"(you have transformers {tf_ver}).\n" - "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." - ) + # 1) Prepare the quantization runtime BEFORE merging, so an unusable config fails fast instead of + # writing a full 16bit checkpoint first. With the llm-compressor-main shadow the subprocess + # validates everything itself, so skip the workspace install / ceiling / scheme checks; without + # it, install the workspace llm-compressor and fail fast past its transformers ceiling. + _shadow_pythonpath = _compressed_quantize_pythonpath() + if _shadow_pythonpath is None: + install_llm_compressor() + # llm-compressor cannot run under a newer transformers than its ceiling: the quantization + # subprocess would die with a cryptic ImportError (TORCH_INIT_FUNCTIONS) only AFTER the costly + # 16bit merge. Detect and fail fast with an actionable message instead. + _exceeds, _tf_ver = _transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + raise RuntimeError( + f"Unsloth: FP8/FP4 compressed-tensors export is not available for this model. It runs " + f"under transformers {_tf_ver}, but llm-compressor supports transformers " + f"<= {_LLM_COMPRESSOR_MAX_TRANSFORMERS}. Export to GGUF or 16-bit instead." + ) + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and # quantize inside an isolated temp dir instead of writing ./ into the cwd. @@ -4307,8 +4513,15 @@ def _unsloth_save_compressed_tensors( env["HF_TOKEN"] = token env["HUGGING_FACE_HUB_TOKEN"] = token + # Clean PYTHONPATH = shadow only. torch still comes from the interpreter's site-packages; + # transformers 5.x + llm-compressor main come from the shadow. Dropping the inherited + # PYTHONPATH removes any parent transformers sidecar so the shadow's is authoritative. + if _shadow_pythonpath is not None: + env["PYTHONPATH"] = _shadow_pythonpath + print( f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + f"{'(llm-compressor-main shadow) ' if _shadow_pythonpath is not None else ''}" "(in a separate process)..." ) try: @@ -4377,6 +4590,255 @@ def _unsloth_save_compressed_tensors( torch.cuda.empty_cache() +def _unsloth_save_torchao( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + kind: str, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + **merge_kwargs, +): + """Export a device-agnostic torchao FP8 / INT8 "portable" checkpoint (no NVIDIA GPU needed). + + Merges LoRA to 16bit in a staging dir, then applies torchao weight-only quantization via + `TorchAoConfig` into `save_directory + "-" + suffix`. No calibration, subprocess, or CUDA. + `kind` is "fp8" (safetensors) or "int8" (.bin; torchao only whitelists float8 for safetensors). + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + if token is None: + token = get_token() + + # Only the main process merges, quantizes, and uploads; other ranks return at once. + if not is_main_process: + return None + + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoProcessor, + TorchAoConfig, + ) + from torchao.quantization import Float8WeightOnlyConfig, Int8WeightOnlyConfig + + if kind == "fp8": + quant_type = Float8WeightOnlyConfig() + safe_serialization = True + elif kind == "int8": + quant_type = Int8WeightOnlyConfig() + safe_serialization = False # torchao only supports safetensors for float8 configs + else: + raise RuntimeError(f"Unsloth: unknown torchao export kind '{kind}' (expected fp8/int8).") + + # Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected + # 16-bit export written to save_directory is not overwritten or deleted; the torchao output is + # the sibling "-" (or the repo id on a hub push). + repo_id, work_tmp, model_dev = None, None, None + work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-") + if push_to_hub: + repo_id = os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + out_dir = staging + "-" + suffix + else: + base = os.fspath(save_directory).rstrip("/\\") or os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(base) or "model") + out_dir = base + "-" + suffix + + api = None + try: + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 1) Merge to 16bit at a staging dir (LoRA and base alike). The reload reads default + # weight filenames, so never write variant-named shards here. + merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before torchao {kind} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = staging, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint. + # A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off + # vision_config / a vision-named architecture only, like the compressed path. + is_vlm = False + trust_remote_code = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + trust_remote_code = bool(getattr(model.config, "auto_map", None)) + # Custom code can be declared only in the tokenizer/processor config, so also honor an + # auto_map in any staged config (the original load already had the user's consent). + if not trust_remote_code: + for _cfg in ( + "config.json", + "tokenizer_config.json", + "processor_config.json", + "preprocessor_config.json", + ): + try: + _p = os.path.join(staging, _cfg) + if os.path.exists(_p): + with open(_p, "r", encoding = "utf-8") as _f: + if "auto_map" in json.load(_f): + trust_remote_code = True + break + except Exception: + pass + # Reload with the class that matches the checkpoint: an image-text VLM class (with a + # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own + # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and + # AutoModelForCausalLM would fail to load them); otherwise causal-LM. + if is_vlm: + try: + from transformers import AutoModelForImageTextToText as _reload_model + except ImportError: + from transformers import AutoModelForVision2Seq as _reload_model + auto_model = _reload_model + elif getattr(getattr(model, "config", None), "is_encoder_decoder", False): + import transformers as _tf + auto_model = next( + ( + getattr(_tf, _arch) + for _arch in (getattr(model.config, "architectures", None) or []) + if getattr(_tf, _arch, None) is not None + ), + AutoModelForCausalLM, + ) + else: + auto_model = AutoModelForCausalLM + auto_processor = AutoProcessor if is_vlm else AutoTokenizer + + # 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk. + # Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit + # resident alongside the reloaded copy and OOM a device that fit the model once. + _has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() + try: + if ( + (torch.cuda.is_available() or _has_xpu) + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if _has_xpu: + torch.xpu.empty_cache() + + # 4) Reload the staged 16bit checkpoint with torchao applied. bfloat16 is required; + # device_map="auto" falls back to CPU, so this works on any hardware. + print(f"Unsloth: Quantizing the merged model to torchao {kind}...") + dtype_kw = {"torch_dtype": torch.bfloat16} if HAS_TORCH_DTYPE else {"dtype": torch.bfloat16} + quantized_model = auto_model.from_pretrained( + staging, + device_map = "auto", + quantization_config = TorchAoConfig(quant_type = quant_type), + trust_remote_code = trust_remote_code, + **dtype_kw, + ) + staged_tokenizer = auto_processor.from_pretrained( + staging, trust_remote_code = trust_remote_code + ) + + quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) + staged_tokenizer.save_pretrained(out_dir) + del quantized_model + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # 5) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: torchao {kind} export failed - no quantization_config written to " + f"{cfg_path}" + ) + + # 6) Optional hub upload of the quantized artifact (the temp staging is cleaned in finally). + if push_to_hub: + print(f"Unsloth: Uploading torchao {kind} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + result = repo_id if push_to_hub else out_dir + print( + f"Unsloth: Saved torchao {kind} checkpoint to '{result}'.\n" + f"Unsloth: This is portable (produced on any device, no NVIDIA GPU required). Load it " + f"with vLLM or transformers; FP8/INT8 acceleration is available on supported GPUs." + ) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after torchao " + "export; it may remain on CPU." + ) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike],