vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: multi-select export formats, portable FP8/INT8, GGUF LoRA, and source parity (#6767)

Browse source 
* Studio: expose full compressed-tensors scheme set in an export formats dropdown

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: multi-select export formats, portable torchao FP8/INT8, GGUF LoRA, source parity

Export page overhaul on top of the formats dropdown:

- Unify merged precision into one sorted multi-select list (16-bit first, then
  8-bit, then 4-bit). Drop "vLLM" from labels, add INT8 (W8A8), INT8 (W8A16),
  INT4 (W4A16), MXFP4, MXFP8. Quick formats render as toggle pills; the rest live
  in a multi-select "More formats" dropdown, so several formats export in one run.
- Add a portable torchao FP8/INT8 save path (Float8WeightOnlyConfig /
  Int8WeightOnlyConfig) that needs no NVIDIA GPU to produce and loads in vLLM.
  FP8 serializes to safetensors, INT8 to .bin. Wired into save_pretrained_merged
  and push_to_hub_merged via a TORCHAO_EXPORT_SCHEMES registry and
  _unsloth_save_torchao, parallel to the compressed-tensors path.
- Hide NVIDIA-only compressed-tensors formats when no NVIDIA GPU is present; keep
  16-bit and portable FP8/INT8. The backend also rejects a compressed request on
  non-NVIDIA hardware so it stays authoritative.
- Relax merged export to non-PEFT models so Local Model and Hugging Face sources
  get the same 16-bit / compressed / portable options.
- GGUF: send the whole quant list in one call (merge once, quantize many).
- LoRA: add a GGUF adapter option (convert_lora_to_gguf.py) with an outtype
  select (f16/bf16/f32/q8_0/auto), alongside the safetensors adapter.
- Thread the new fields through models, routes, orchestrator, and worker; extend
  the export tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: gate export by accelerator with a torch-aware reason; fix export save dir naming

Export runs through Unsloth, which requires a compute accelerator (NVIDIA/AMD/Intel
GPU or Apple MLX) and has no CPU code path, so a bare-CPU host cannot export even
with PyTorch installed. Add export_capability() in utils/hardware that reports
export_supported plus a precise reason so the UI stops showing a generic "no GPU":

  - pytorch_not_installed: a --no-torch install (even a physical GPU is unusable)
  - no_accelerator: PyTorch present but no supported accelerator (bare CPU)
  - mlx_unavailable: Apple Silicon where the MLX stack is missing or too old

Expose the fields on /api/system/hardware and /api/system, and guard the mutating
export routes (load-checkpoint, export/merged|base|gguf|lora) with HTTP 400 and the
reason, leaving read-only endpoints usable so the Export page still renders.

Make core/export/export.py import without PyTorch and without a usable accelerator
(the Unsloth import is caught) so the export worker degrades to a clear message
instead of crashing at import.

Frontend: keep /export reachable on chat-only hosts and gray out the method and
format options with the backend reason (Alert plus disabled MethodPicker) instead
of silently redirecting to /chat, so users see why export is unavailable.

Also fix the export save directory producing "model/null" for Local Model and
Hugging Face sources that have no run/checkpoint, naming the folder from the model id.

* CI: validate Studio export capability gating on Linux, Windows and macOS

Add a small pytest matrix that runs studio/backend/tests/test_export_capability.py
on ubuntu-latest, windows-latest and macos-latest. It confirms, on each real OS,
that hardware.export_capability() reports the right decision and reason
(pytorch_not_installed, no_accelerator, or mlx_unavailable) and that the export
backend imports without PyTorch and degrades to a clear message instead of crashing.

Hosted runners have no GPU/MLX, so this covers the "export unavailable, here is why"
path a Mac/Windows user without an accelerator sees; a real accelerator export is
validated separately. The job installs only a CPU PyTorch plus the backend import
deps (no unsloth, triton, or llama.cpp), so it runs in seconds with no GPU.

* Studio export: address Codex review (source-aware gating, GGUF LoRA token/MLX/guard)

Frontend (export-page):
- Gate LoRA and quantized-model restrictions on the active source. isAdapter /
  isQuantized come from the selected checkpoint; in Local Model / Hugging Face
  ("model") source mode they were stale, so LoRA stayed wrongly enabled for a
  direct base model (backend then rejects "No adapter to export") and a stale
  "quantized" flag disabled every method for an unrelated, exportable model. Add
  effectiveIsAdapter / effectiveIsQuantized (false outside checkpoint mode) and use
  them in the method-reset effect and the MethodPicker disabled state.
- Hide the GGUF LoRA option on a macOS/MLX host (the backend rejects GGUF LoRA on
  MLX), so users no longer pick it, wait through the load, and always fail. Disable
  the "GGUF adapter" button on a Mac host and never send loraGguf there.

Backend (core/export/export.py):
- Pass the HF token into the GGUF LoRA conversion (save_pretrained_gguf), so a
  gated/private base model's config fetch in convert_lora_to_gguf.py is
  authenticated; without it the load can succeed but the conversion fails.
- Guard the save_pretrained_gguf capability check with getattr so an older Unsloth
  model that lacks the method returns the clean "not supported" message instead of
  an AttributeError that surfaces as a generic 500.

* Studio export: address 2nd Codex review (CI index, empty merged, test import)

- studio-export-capability-ci.yml: add --extra-index-url https://pypi.org/simple to
  the torch install so torch's transitive deps still resolve; --index-url alone
  replaces PyPI with only the CPU wheel index, which does not serve all of them.
- export-page handleStart: reject an empty merged selection (mirrors canExport), so
  clicking the panel's Start button with every precision pill deselected no longer
  submits mergedSelections: [] and launches an unintended default 16-bit export.
- test_export_imatrix_compressed: the torchao-registry test now reads unsloth/save.py
  as text (like the other ast/string checks) instead of `import unsloth.save`, which
  raised ModuleNotFoundError in the CPU studio-backend suite that has no unsloth
  installed.

* Studio export: make comments succinct across the export changes

* Studio export: use load token for local GGUF LoRA export of gated bases

* Studio export: harden portable torchao path and gate multi-format Hub push

torchao (_unsloth_save_torchao):
- merge to an isolated temp staging dir so a co-selected 16-bit output at save_directory is not deleted
- narrow VLM detection to vision_config / ForVisionText2Text so T5/BART/Whisper are not misrouted
- forward trust_remote_code (from auto_map) to the reload so custom-code models export

Export UI:
- hide portable torchao formats on macOS/MLX (backend rejects quantized export there)
- restrict a Hub merged export to a single format (each writes to the repo root)

* Studio export: torchao tokenizer remote-code + XPU offload, scale GGUF timeout

torchao (_unsloth_save_torchao):
- honor auto_map in the staged tokenizer/processor configs (not just model.config) when
  deriving trust_remote_code, so custom-code tokenizers reload after the merge
- offload single-device XPU models to CPU (and empty the XPU cache) before the reload, matching
  the CUDA path, so an Intel GPU that fits the model once does not OOM on the second copy

Export orchestrator:
- scale the GGUF wait timeout by the number of requested quants so a multi-quant list export of a
  large model does not time out at a flat 3600s

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio export: show portable torchao formats only on non-NVIDIA (CPU) hosts

Portable torchao FP8/INT8 is the fallback for hosts without the NVIDIA compressed-tensors path.
On an NVIDIA GPU the compressed-tensors FP8/FP4/INT formats are the intended path (llm-compressor
auto-installs), so hide the portable duplicates there; keep them on CPU / non-NVIDIA hosts and
continue hiding them on macOS/MLX.

* Studio export: report all output folders and the exported formats

- Multi-format merged export now collects every sibling output directory (one per selected
  precision) instead of only the last; the success banner lists them all.
- Show the selected precision formats in the run summary (a Formats row, like GGUF Quantizations),
  so the panel says what is being exported rather than just 'Merged Model'.
- Persist the selected formats in the run summary and seed them on mount, so navigating away and
  back (or toggling the export method) restores the selection instead of resetting to 16-bit.

* Studio export: list all output formats, add GGUF LoRA target, default Q8_0, auto-select newest checkpoint

- Progress/summary panel now shows a Formats row with the selected merged
  formats, and the success banner lists every output folder a multi-format
  merged run creates (one line per format) instead of only the last one.
- Merged format selection is seeded from the active run, so navigating away
  and back (or switching method cards) no longer resets it to 16-bit.
- GGUF / Llama.cpp now offers an Export target toggle (Full model or LoRA
  adapter) for adapter checkpoints, reusing the LoRA GGUF export path.
- Removed the Auto GGUF LoRA output type and defaulted to Q8_0 in the UI,
  the request model, and the backend defaults; the outtype list is now
  Q8_0/F16/BF16/F32. Core save.py still accepts auto for external callers.
- When a finetune has no checkpoint selected, auto-select the newest one.

* Studio torchao export: robust reload class + optional VLM import

Two fixes to the portable torchao FP8/INT8 export reload, from review of the
narrowed VLM detection:

- Encoder-decoder seq2seq checkpoints (T5/BART/Whisper) are not causal LMs.
  With the narrowed is_vlm test they now correctly skip the image-text class,
  but fell through to AutoModelForCausalLM and failed to reload after the merge.
  Reload them with their own architecture class from the config instead.
- AutoModelForImageTextToText was imported unconditionally at the top of the
  torchao path, so on Transformers builds without that class the import aborted
  every torchao export (even text-only). Import it lazily only for a VLM, with
  the AutoModelForVision2Seq fallback used elsewhere in Unsloth.

* Studio: enable FP8/FP4 compressed export for newer-transformers models

The shipped llm-compressor 0.10.x pins transformers<=4.57.6, so FP8/FP4 export failed
for models needing a transformers 5.x sidecar (Qwen3.5, Gemma-4, Qwen3-Next): the
quantization subprocess crashed importing the removed TORCH_INIT_FUNCTIONS.

Run the quantization against a dedicated llm-compressor-main "shadow": a --target
package dir (transformers 5.10.2 + llm-compressor main + compressed-tensors) layered
over the existing torch. It installs --no-deps so torch is never touched (works on any
Studio torch build), is provisioned lazily and fingerprint-cached, and can be turned
off with UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN.

- transformers_version.py: provision + validate .venv_llmcompressor.
- export.py: route all compressed exports through the shadow when available; else keep
  the workspace 0.10.x path and fail fast past its transformers ceiling.
- save.py: launch _compressed_quantize.py with a clean PYTHONPATH = shadow.
- _compressed_quantize.py: skip linear_attn / vision tower / MTP modules (matches the
  RedHatAI and NVIDIA reference quants, and is required by the grouped schemes).

Verified all four schemes (fp8, w8a8, w4a16, mxfp4) on Qwen3.5-9B and Llama-3.2-1B, and
fp8 on Gemma-4, end to end through Studio.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix GGUF LoRA export tests

* Fix export CI expectations

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
This commit is contained in:
Daniel Han 2026-07-03 08:25:10 -07:00 • committed by GitHub
parent 308ea5a93c
commit 026141a4a4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 2120 additions and 230 deletions
Download patch file Download diff file Expand all files Collapse all files

76
.github/workflows/studio-export-capability-ci.yml vendored Normal file
View file

@ -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

305
studio/backend/core/export/export.py
View file

@ -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 "<dir>-<suffix>" directory (for vLLM).
_COMPRESSED = {
"FP8 (compressed-tensors)": ("fp8", "fp8"),
"NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"),
# Quantized formats save to a sibling "<dir>-<suffix>". 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 "<dir>-<suffix>" sibling; report that as output.
# Compressed / torchao writes to the "<dir>-<suffix>" 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 "<base>-lora-<outtype>.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)

18
studio/backend/core/export/orchestrator.py
View file

@ -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", "")

3
studio/backend/core/export/worker.py
View file

@ -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}"

8
studio/backend/main.py
View file

@ -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

27
studio/backend/models/export.py
View file

@ -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).",
)

25
studio/backend/routes/export.py
View file

@ -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:

156
studio/backend/tests/test_export_capability.py Normal file
View file

@ -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

140
studio/backend/tests/test_export_imatrix_compressed.py
View file

@ -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")

7
studio/backend/utils/hardware/__init__.py
View file

@ -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",

43
studio/backend/utils/hardware/hardware.py
View file

@ -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.

151
studio/backend/utils/transformers_version.py
View file

@ -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:

3
studio/frontend/src/app/routes/__root.tsx
View file

@ -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 {

16
studio/frontend/src/components/app-sidebar.tsx
View file

@ -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();
}}

9
studio/frontend/src/features/export/api/export-api.ts
View file

@ -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<ExportOperationResponse> {
const response = await authFetch("/api/export/export/lora", {
method: "POST",

51
studio/frontend/src/features/export/components/export-run-panel.tsx
View file

@ -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."}
</span>
{run.result?.outputPath ? (
<code
className="select-all break-all font-mono text-[12px] text-foreground/90"
title={run.result.outputPath}
>
{run.result.outputPath}
</code>
) : 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) => (
<div key={`${o.path}-${i}`} className="flex min-w-0 flex-col gap-0.5">
{showLabels && o.label ? (
<span className="text-xs text-emerald-700/80 dark:text-emerald-300/80">
{o.label}
</span>
) : null}
<code
className="select-all break-all font-mono text-[12px] text-foreground/90"
title={o.path}
>
{o.path}
</code>
</div>
));
})()}
</div>
</div>
)}
@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
<span>Export Method</span>
<span className="font-medium text-foreground">{summaryMethodLabel}</span>
</div>
{summaryMethod === "merged" && summaryFormats.length > 0 && (
<div className="flex justify-between gap-3">
<span>Formats</span>
<span className="font-medium text-foreground text-right">
{summaryFormats.join(", ")}
</span>
</div>
)}
{summaryMethod === "gguf" && summaryQuants.length > 0 && (
<div className="flex justify-between">
<span>Quantizations</span>

168
studio/frontend/src/features/export/constants.ts
View file

@ -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

501
studio/frontend/src/features/export/export-page.tsx
View file

@ -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 "<run>/<checkpoint>" 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<MergedFormat>("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<string[]>(() => {
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<string>("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<MergedFormatOption[]>(
() =>
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() {
</div>
</div>
{exportUnsupported && (
<Alert variant="destructive">
<HugeiconsIcon icon={AlertCircleIcon} className="size-4" />
<AlertTitle>Export unavailable</AlertTitle>
<AlertDescription>{exportUnsupportedMessage}</AlertDescription>
</Alert>
)}
<MethodPicker
value={exportMethod}
onChange={handleMethodChange}
disabledMethods={
!isAdapter && isQuantized
exportUnsupported
? ["merged", "lora", "gguf"]
: !isAdapter || sourceMode === "model"
? ["merged", "lora"]
: []
: !effectiveIsAdapter && effectiveIsQuantized
? ["merged", "lora", "gguf"]
: !effectiveIsAdapter
? ["lora"]
: []
}
disabledReason={
!isAdapter && isQuantized
? "Pre-quantized (BNB 4-bit) models cannot be exported without LoRA adapters"
: sourceMode === "model"
? "Only GGUF export is available for direct model export"
: !isAdapter
? "Not available for full fine-tune checkpoints (no LoRA adapters)"
exportUnsupported
? exportUnsupportedMessage
: !effectiveIsAdapter && effectiveIsQuantized
? "Pre-quantized (BNB 4-bit) models cannot be exported without LoRA adapters"
: !effectiveIsAdapter
? "LoRA-only export needs a LoRA adapter checkpoint"
: undefined
}
/>
{exportMethod === "merged" && isAdapter && (
<div className="space-y-2">
<div className="text-sm font-medium">Precision</div>
<div className="flex flex-wrap gap-2">
{MERGED_FORMATS.map((f) => (
<Button
key={f.value}
type="button"
variant={mergedFormat === f.value ? "default" : "outline"}
size="sm"
onClick={() => setMergedFormat(f.value)}
title={f.hint}
>
{f.label}
</Button>
))}
</div>
<div className="text-xs text-muted-foreground">
{MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint}
{exportMethod === "merged" && !exportUnsupported && (
<div className="space-y-3">
<div className="space-y-2">
<div className="flex items-center justify-between">
<div className="text-sm font-medium">Precision</div>
<span className="text-[11px] text-muted-foreground/70">
— select one or more
</span>
</div>
<div className="flex flex-wrap gap-2">
{availableFormats
.filter((f) => f.common)
.map((f) => {
const active = selectedFormats.includes(f.value);
return (
<Button
key={f.value}
type="button"
variant={active ? "default" : "outline"}
size="sm"
onClick={() => toggleFormat(f.value)}
title={f.hint}
>
{f.label}
{f.needsCalibration ? " *" : ""}
</Button>
);
})}
{availableFormats.some((f) => !f.common) && (
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button type="button" variant="outline" size="sm">
More formats
{selectedFormats.some((v) =>
availableFormats.find(
(f) => f.value === v && !f.common,
),
)
? " ✓"
: "…"}
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="w-64">
<DropdownMenuLabel>
Additional formats
</DropdownMenuLabel>
<DropdownMenuSeparator />
{availableFormats
.filter((f) => !f.common)
.map((f) => (
<DropdownMenuCheckboxItem
key={f.value}
checked={selectedFormats.includes(f.value)}
onCheckedChange={() => toggleFormat(f.value)}
onSelect={(e) => e.preventDefault()}
>
<span className="flex flex-col">
<span>
{f.label}
{f.needsCalibration ? " *" : ""}
</span>
<span className="text-[10px] text-muted-foreground">
{f.hint}
</span>
</span>
</DropdownMenuCheckboxItem>
))}
</DropdownMenuContent>
</DropdownMenu>
)}
</div>
{selectedFormats.length > 0 && (
<div className="flex flex-wrap items-center gap-x-3 gap-y-1">
<span className="text-[11px] text-muted-foreground">
{selectedFormats.length} selected:{" "}
{selectedFormats
.map(
(v) =>
MERGED_FORMATS.find((f) => f.value === v)
?.label ?? v,
)
.join(", ")}
</span>
{selectedFormats.length > 1 && (
<button
type="button"
onClick={() => setSelectedFormats(["16-bit"])}
className="text-[11px] text-muted-foreground/70 hover:text-foreground transition-colors"
>
Reset to 16-bit
</button>
)}
</div>
)}
{hubMultiFormat && (
<div className="text-[11px] text-amber-600 dark:text-amber-500">
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.
</div>
)}
{selectedFormats.some(
(v) =>
MERGED_FORMATS.find((f) => f.value === v)
?.needsCalibration,
) && (
<div className="text-[11px] text-muted-foreground">
* calibrates on data (uses a small calibration set).
</div>
)}
{!hasNvidia && (
<div className="text-[11px] text-muted-foreground">
No NVIDIA GPU detected: compressed-tensors formats are
hidden. 16-bit and portable FP8/INT8 (torchao) still
work here and load in vLLM.
</div>
)}
</div>
</div>
)}
{exportMethod === "gguf" && (
<>
<QuantPicker
value={quantLevels}
onChange={setQuantLevels}
sizes={quantSizeLabels}
/>
<div className="flex items-center justify-between gap-3 rounded-lg border p-3">
<div className="space-y-0.5">
<div className="text-sm font-medium">
Importance matrix (imatrix)
{exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && (
<div className="space-y-3">
<div className="space-y-2">
<div className="text-sm font-medium">Adapter format</div>
<div className="flex flex-wrap gap-2">
<Button
type="button"
variant={!loraAsGguf ? "default" : "outline"}
size="sm"
onClick={() => setLoraAsGguf(false)}
title="Standard PEFT adapter (adapter_model.safetensors)."
>
Adapter (safetensors)
</Button>
<Button
type="button"
variant={loraAsGguf ? "default" : "outline"}
size="sm"
disabled={isMacHost}
onClick={() => setLoraAsGguf(true)}
title={
isMacHost
? "GGUF LoRA export is not available on macOS/MLX. Use the safetensors adapter."
: "llama.cpp GGUF LoRA, loadable with `llama-cli --lora`."
}
>
GGUF adapter
</Button>
</div>
<div className="text-xs text-muted-foreground">
{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."}
</div>
</div>
{loraAsGguf && (
<div className="space-y-1.5">
<div className="text-sm font-medium">Output type</div>
<Select
value={loraGgufOuttype}
onValueChange={(v) => setLoraGgufOuttype(v)}
>
<SelectTrigger className="w-full sm:w-56">
<SelectValue />
</SelectTrigger>
<SelectContent>
{LORA_GGUF_OUTTYPES.map((t) => (
<SelectItem key={t} value={t}>
{t.toUpperCase()}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
)}
</div>
)}
{exportMethod === "gguf" && !exportUnsupported && (
<div className="space-y-3">
{effectiveIsAdapter && !isMacHost && (
<div className="space-y-2">
<div className="text-sm font-medium">Export target</div>
<div className="flex flex-wrap gap-2">
<Button
type="button"
variant={ggufTarget === "model" ? "default" : "outline"}
size="sm"
onClick={() => setGgufTarget("model")}
title="Merge the adapter into the base model, then quantize the full model to GGUF."
>
Full model
</Button>
<Button
type="button"
variant={ggufTarget === "lora" ? "default" : "outline"}
size="sm"
onClick={() => setGgufTarget("lora")}
title="Export just the adapter as a GGUF LoRA (llama.cpp `--lora`); the base model stays separate."
>
LoRA adapter
</Button>
</div>
<div className="text-xs text-muted-foreground">
{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."}
</div>
</div>
<Switch
checked={effectiveImatrix}
onCheckedChange={setUseImatrix}
disabled={requiresImatrix}
/>
</div>
</>
)}
{ggufAsLora ? (
<div className="space-y-1.5">
<div className="text-sm font-medium">Output type</div>
<Select
value={loraGgufOuttype}
onValueChange={(v) => setLoraGgufOuttype(v)}
>
<SelectTrigger className="w-full sm:w-56">
<SelectValue />
</SelectTrigger>
<SelectContent>
{LORA_GGUF_OUTTYPES.map((t) => (
<SelectItem key={t} value={t}>
{t.toUpperCase()}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
) : (
<>
<QuantPicker
value={quantLevels}
onChange={setQuantLevels}
sizes={quantSizeLabels}
/>
<div className="flex items-center justify-between gap-3 rounded-lg border p-3">
<div className="space-y-0.5">
<div className="text-sm font-medium">
Importance matrix (imatrix)
</div>
<div className="text-xs text-muted-foreground">
{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."}
</div>
</div>
<Switch
checked={effectiveImatrix}
onCheckedChange={setUseImatrix}
disabled={requiresImatrix}
/>
</div>
</>
)}
</div>
)}
{estimatedSize && (
<div className="flex items-center gap-1.5 text-xs text-muted-foreground">

116
studio/frontend/src/features/export/stores/export-runtime-store.ts
View file

@ -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<ExportRuntimeStore>()((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<ExportRuntimeStore>()((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<ExportRuntimeStore>()((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<ExportRuntimeStore>()((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;

15
studio/frontend/src/hooks/use-hardware-info.ts
View file

@ -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<HardwareInfo> {
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;

14
tests/studio/playwright_extra_ui.py
View file

@ -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

4
unsloth/_compressed_quantize.py
View file

@ -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$"]

494
unsloth/save.py
View file

@ -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 ./<repo_id> 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 "<save_directory>-<suffix>" (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],