Studio: scale export GGUF size estimates from the real model size (#6418)
* Studio: scale export GGUF size estimates from the real model size The Export page showed hardcoded, model-independent GGUF quant size labels (Q8_0 ~8.2 GB, BF16 ~14.2 GB, ...) calibrated for an ~8B model. For a 35B MoE model like Qwen3.6-35B-A3B (67 GiB bf16, Q8 ~34 GiB) the picker wrongly reported Q8 ~8.2 GB. Only the displayed estimate was wrong; the actual export via save_pretrained_gguf was always correct. Add GET /api/models/export-size, which returns a model's MoE-aware fp16/bf16-equivalent size and total params using the existing estimate_fp16_model_size_bytes (safetensors -> config -> local -> vllm). The result is memoized and degrades to nulls so a size hint can never break the Export page. The Export picker now scales each quant from that size (bytes ~= fp16_bytes * bits_per_weight / 16, GiB units to match the model selector), and renders no size when it is unknown rather than a misleading fixed number. The Est. size summary in the page and dialog is restored now that the value comes from the backend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio export-size: address review feedback - Run the size estimate off the event loop with asyncio.to_thread so a slow Hugging Face request cannot stall other API or SSE endpoints. - Cache only successful estimates; a transient failure (offline, gated before credentials) is no longer pinned as unavailable until restart. - Forward the HF token so private and gated models can be sized, and refetch when the token changes. - Clamp the size formatter index so sub-1-byte values cannot pick an out-of-range unit. * Studio export-size: address second review pass - Send the HF token in an X-HF-Token header instead of the query string, so it never lands in URLs, logs, or browser history. - Key the estimate cache by model id only (the fp16 size is token independent), so HF tokens are never retained in the cache. - Restrict local-path sizing to known Studio roots (outputs/exports/cache/home) so an authenticated caller cannot trigger a scan of an arbitrary directory. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio export-size: fix CI (import-hoist + isolated-load test stubs) - Import ExportSizeResponse from models.models in routes/models.py instead of re-exporting it through models/__init__.py, so the import-hoist lint does not flag a newly added but un-loaded re-export (models/__init__.py is unchanged). - Add Header and ExportSizeResponse to the stubbed fastapi / models.models in test_export_absolute_paths.py, which loads routes/models.py in isolation. * Studio: validate export-size local path before filesystem access CodeQL flagged the export-size local-path guard as path injection: the user-provided model path was resolved and stat-ed before it was checked for containment under a Studio data root. Decide containment by lexical normalization (normpath/abspath/expanduser, no filesystem access) and only touch the filesystem once the path is proven to sit under a trusted root, so an unvalidated value never reaches a filesystem call. Add a direct containment unit test (under-root, root itself, missing, /etc, and '..' traversal). * Studio: trim export-size comments to be more concise Shorten docstrings and comments on the export-size endpoint, helpers, tests, and frontend size utilities; drop comments that just restate the code. Verified code-identical (comments only) via AST/TS-compiler check. No behavior change. * Studio: harden export-size local-path handling Address review feedback on the export-size endpoint's local sizing: - Resolve symlinks and re-verify containment in _is_sizable_local_path so a symlink inside a Studio root can't point the sizer outside it. - Re-validate the resolved LoRA base before sizing, so a crafted adapter whose base_model points outside the roots can't redirect the scan. - Skip nested checkpoint-*/global_step* snapshots when summing local weight sizes so a run dir's intermediate checkpoints don't inflate the estimate. - Size the checkpoint directory for full fine-tune checkpoint exports (whose base may be a local/custom path), keeping base-model sizing for adapters. Adds tests for the adapter-base escape, symlink escape, and nested-checkpoint exclusion. * [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: danielhanchen <michaelhan2050@gmail.com>
This commit is contained in:
Daniel Han
GitHub
parent
3bfc83781d
commit
22e6d64493
11 changed files with 648 additions and 42 deletions
|
|
@ -53,6 +53,24 @@ class CheckpointListResponse(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
class ExportSizeResponse(BaseModel):
|
||||
"""Model fp16/bf16-equivalent size; size fields are null when unknown."""
|
||||
|
||||
model: str = Field(..., description = "Model id or path the estimate was computed for")
|
||||
fp16_bytes: Optional[int] = Field(
|
||||
None,
|
||||
description = "Estimated FP16/BF16-equivalent on-disk size in bytes, or null if unknown",
|
||||
)
|
||||
total_params: Optional[int] = Field(
|
||||
None,
|
||||
description = "Estimated total parameter count (fp16_bytes // 2), or null if unknown",
|
||||
)
|
||||
source: str = Field(
|
||||
"unavailable",
|
||||
description = "How the estimate was derived (e.g. safetensors, config, local, vllm, unavailable)",
|
||||
)
|
||||
|
||||
|
||||
class ModelDetails(BaseModel):
|
||||
"""Model configuration and metadata; used for both list and detail views"""
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
"""Model management API routes."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -10,7 +11,7 @@ import shutil
|
|||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from fastapi import APIRouter, Body, Depends, HTTPException, Query
|
||||
from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
|
||||
from typing import List, Optional
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
|
|
@ -149,6 +150,7 @@ from models import (
|
|||
from models.models import (
|
||||
BrowseEntry,
|
||||
BrowseFoldersResponse,
|
||||
ExportSizeResponse,
|
||||
GgufVariantDetail,
|
||||
GgufVariantsResponse,
|
||||
ModelType,
|
||||
|
|
@ -3008,3 +3010,121 @@ async def list_checkpoints(
|
|||
event = "models.list_checkpoints_failed",
|
||||
log = logger,
|
||||
)
|
||||
|
||||
|
||||
# Successful estimates only, keyed by model id (token-independent, never stored).
|
||||
# Failures are not cached so a transient offline/gated error can recover later.
|
||||
_EXPORT_SIZE_CACHE: dict[str, tuple[int, int, str]] = {}
|
||||
|
||||
|
||||
def _is_sizable_local_path(model: str) -> bool:
|
||||
"""True only for local paths under a Studio data root.
|
||||
|
||||
Containment is decided lexically (no filesystem access) before the path is
|
||||
touched, then the path is symlink-resolved and re-checked so a symlink
|
||||
inside a root can't point the sizer outside it. A user-controlled path thus
|
||||
can't trigger a scan of an arbitrary dir.
|
||||
"""
|
||||
from utils.paths import outputs_root, exports_root, studio_root
|
||||
from utils.paths.storage_roots import cache_root
|
||||
|
||||
def _lexical(p: str) -> str:
|
||||
# Lexical only (no filesystem read); normpath collapses '..'.
|
||||
return os.path.normpath(os.path.abspath(os.path.expanduser(p)))
|
||||
|
||||
raw_roots = [studio_root(), outputs_root(), exports_root(), cache_root()]
|
||||
roots = []
|
||||
for root in raw_roots:
|
||||
try:
|
||||
roots.append(_lexical(str(root)))
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
continue
|
||||
|
||||
try:
|
||||
candidate = _lexical(model)
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
return False
|
||||
for root in roots:
|
||||
if candidate == root or candidate.startswith(root + os.sep):
|
||||
# Contained lexically; resolve symlinks and re-verify the real path
|
||||
# is still under a root before touching the filesystem.
|
||||
try:
|
||||
real = os.path.realpath(candidate)
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
return False
|
||||
for raw in raw_roots:
|
||||
try:
|
||||
real_root = os.path.realpath(str(raw))
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
continue
|
||||
if real == real_root or real.startswith(real_root + os.sep):
|
||||
return os.path.exists(real)
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def _export_size_cached(
|
||||
model: str, hf_token: Optional[str]
|
||||
) -> tuple[Optional[int], Optional[int], str]:
|
||||
"""Estimate a model's fp16/bf16-equivalent size in bytes (+ total params).
|
||||
|
||||
Memoizes successful results by model id; never raises (failures return
|
||||
(None, None, "unavailable") and are not cached). Blocking I/O; call off-thread.
|
||||
"""
|
||||
cached = _EXPORT_SIZE_CACHE.get(model)
|
||||
if cached is not None:
|
||||
return cached
|
||||
try:
|
||||
from utils.hardware.hardware import (
|
||||
_resolve_model_identifier_for_gpu_estimate,
|
||||
estimate_fp16_model_size_bytes,
|
||||
)
|
||||
|
||||
# A local LoRA adapter is sized via its base model, which the sizer
|
||||
# reads from the adapter config; re-validate that resolved base so a
|
||||
# crafted adapter can't redirect the local scan outside the roots.
|
||||
if is_local_path(model):
|
||||
base = _resolve_model_identifier_for_gpu_estimate(model, hf_token = hf_token)
|
||||
if is_local_path(base) and not _is_sizable_local_path(base):
|
||||
return None, None, "unavailable"
|
||||
|
||||
fp16_bytes, source = estimate_fp16_model_size_bytes(model, hf_token = hf_token)
|
||||
if not fp16_bytes or fp16_bytes <= 0:
|
||||
return None, None, source or "unavailable"
|
||||
result = (int(fp16_bytes), int(fp16_bytes) // 2, source)
|
||||
_EXPORT_SIZE_CACHE[model] = result
|
||||
return result
|
||||
except Exception as e: # a size hint must never break export
|
||||
logger.warning("Could not estimate export size for '%s': %s", model, e)
|
||||
return None, None, "unavailable"
|
||||
|
||||
|
||||
@router.get("/export-size", response_model = ExportSizeResponse)
|
||||
async def get_export_size(
|
||||
model: str = Query(..., description = "Base model id or local model path to size"),
|
||||
hf_token: Optional[str] = Header(None, alias = "X-HF-Token"),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Estimate a model's fp16/bf16-equivalent size for the Export page.
|
||||
|
||||
Returns nulls with HTTP 200 when the size can't be determined. The HF token
|
||||
(for gated repos) comes from the X-HF-Token header so it never hits URLs/logs.
|
||||
"""
|
||||
if is_local_path(model):
|
||||
if not _is_sizable_local_path(model):
|
||||
return ExportSizeResponse(
|
||||
model = model, fp16_bytes = None, total_params = None, source = "unavailable"
|
||||
)
|
||||
resolved = model
|
||||
else:
|
||||
resolved = resolve_cached_repo_id_case(model)
|
||||
# Blocking network/disk I/O: run off the event loop.
|
||||
fp16_bytes, total_params, source = await asyncio.to_thread(
|
||||
_export_size_cached, resolved, hf_token
|
||||
)
|
||||
return ExportSizeResponse(
|
||||
model = resolved,
|
||||
fp16_bytes = fp16_bytes,
|
||||
total_params = total_params,
|
||||
source = source,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
|
|||
fastapi.APIRouter = lambda: _Router()
|
||||
fastapi.Body = lambda default = None, **_kwargs: default
|
||||
fastapi.Depends = lambda dependency = None, **_kwargs: dependency
|
||||
fastapi.Header = lambda default = None, **_kwargs: default
|
||||
fastapi.HTTPException = _HTTPException
|
||||
fastapi.Query = lambda default = None, **_kwargs: default
|
||||
fastapi.Request = object
|
||||
|
|
@ -191,6 +192,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
|
|||
for name in (
|
||||
"BrowseEntry",
|
||||
"BrowseFoldersResponse",
|
||||
"ExportSizeResponse",
|
||||
"GgufVariantDetail",
|
||||
"GgufVariantsResponse",
|
||||
"ScanFolderInfo",
|
||||
|
|
|
|||
244
studio/backend/tests/test_export_size_estimate.py
Normal file
244
studio/backend/tests/test_export_size_estimate.py
Normal file
|
|
@ -0,0 +1,244 @@
|
|||
# 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 GET /api/models/export-size (the Export page size estimate).
|
||||
|
||||
The endpoint must never raise and must degrade to nulls when size is unknown.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
# Real Qwen3.6-35B-A3B: 35.95B params -> ~67 GiB bf16 (UI wrongly showed Q8 ~8.2 GB).
|
||||
_QWEN35_PARAMS = 35_951_822_704
|
||||
_QWEN35_FP16_BYTES = _QWEN35_PARAMS * 2
|
||||
|
||||
|
||||
def _load_route_module(name: str, relative_path: str):
|
||||
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class TestExportSizeEndpoint(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models_route = _load_route_module(
|
||||
"models_route_module_for_export_size_test",
|
||||
"routes/models.py",
|
||||
)
|
||||
|
||||
def setUp(self):
|
||||
self.models_route._EXPORT_SIZE_CACHE.clear()
|
||||
|
||||
def _call(self, model: str = "unsloth/Qwen3.6-35B-A3B"):
|
||||
with (
|
||||
patch.object(self.models_route, "is_local_path", return_value = False),
|
||||
patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m),
|
||||
):
|
||||
return asyncio.run(
|
||||
self.models_route.get_export_size(
|
||||
model = model, hf_token = None, current_subject = "test-user"
|
||||
)
|
||||
)
|
||||
|
||||
def test_known_model_returns_bytes_and_params(self):
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (_QWEN35_FP16_BYTES, "safetensors"),
|
||||
):
|
||||
resp = self._call()
|
||||
self.assertEqual(resp.fp16_bytes, _QWEN35_FP16_BYTES)
|
||||
self.assertEqual(resp.total_params, _QWEN35_PARAMS)
|
||||
self.assertEqual(resp.source, "safetensors")
|
||||
self.assertEqual(resp.model, "unsloth/Qwen3.6-35B-A3B")
|
||||
|
||||
def test_moe_via_config_fallback(self):
|
||||
# MoE sized via the sizer's config path -> source "config".
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (67 * (1024**3), "config"),
|
||||
):
|
||||
resp = self._call()
|
||||
self.assertEqual(resp.fp16_bytes, 67 * (1024**3))
|
||||
self.assertEqual(resp.total_params, (67 * (1024**3)) // 2)
|
||||
self.assertEqual(resp.source, "config")
|
||||
|
||||
def test_unknown_size_returns_nulls_not_error(self):
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (None, "unavailable"),
|
||||
):
|
||||
resp = self._call()
|
||||
self.assertIsNone(resp.fp16_bytes)
|
||||
self.assertIsNone(resp.total_params)
|
||||
self.assertEqual(resp.source, "unavailable")
|
||||
|
||||
def test_zero_size_treated_as_unknown(self):
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (0, "safetensors"),
|
||||
):
|
||||
resp = self._call()
|
||||
self.assertIsNone(resp.fp16_bytes)
|
||||
self.assertIsNone(resp.total_params)
|
||||
|
||||
def test_sizer_exception_is_swallowed(self):
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
side_effect = RuntimeError("boom"),
|
||||
):
|
||||
resp = self._call()
|
||||
self.assertIsNone(resp.fp16_bytes)
|
||||
self.assertEqual(resp.source, "unavailable")
|
||||
|
||||
def test_result_is_memoized_per_model(self):
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (_QWEN35_FP16_BYTES, "safetensors"),
|
||||
) as mock_sizer:
|
||||
first = self._call()
|
||||
second = self._call()
|
||||
self.assertEqual(first.fp16_bytes, second.fp16_bytes)
|
||||
self.assertEqual(mock_sizer.call_count, 1)
|
||||
|
||||
def test_failures_are_not_cached(self):
|
||||
# A transient failure must not poison the cache; a later call recovers.
|
||||
with patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
side_effect = [(None, "unavailable"), (_QWEN35_FP16_BYTES, "safetensors")],
|
||||
) as mock_sizer:
|
||||
first = self._call()
|
||||
second = self._call()
|
||||
self.assertIsNone(first.fp16_bytes)
|
||||
self.assertEqual(second.fp16_bytes, _QWEN35_FP16_BYTES)
|
||||
self.assertEqual(mock_sizer.call_count, 2)
|
||||
|
||||
def test_token_is_forwarded_to_sizer(self):
|
||||
with (
|
||||
patch.object(self.models_route, "is_local_path", return_value = False),
|
||||
patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m),
|
||||
patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (_QWEN35_FP16_BYTES, "safetensors"),
|
||||
) as mock_sizer,
|
||||
):
|
||||
asyncio.run(
|
||||
self.models_route.get_export_size(
|
||||
model = "unsloth/Private",
|
||||
hf_token = "secret-token",
|
||||
current_subject = "test-user",
|
||||
)
|
||||
)
|
||||
self.assertEqual(mock_sizer.call_args.kwargs.get("hf_token"), "secret-token")
|
||||
|
||||
def test_arbitrary_local_path_is_not_scanned(self):
|
||||
# Unsafe local paths must not be scanned -> unavailable.
|
||||
with (
|
||||
patch.object(self.models_route, "is_local_path", return_value = True),
|
||||
patch.object(self.models_route, "_is_sizable_local_path", return_value = False),
|
||||
patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer,
|
||||
):
|
||||
resp = asyncio.run(
|
||||
self.models_route.get_export_size(
|
||||
model = "/etc", hf_token = None, current_subject = "test-user"
|
||||
)
|
||||
)
|
||||
self.assertIsNone(resp.fp16_bytes)
|
||||
self.assertEqual(resp.source, "unavailable")
|
||||
mock_sizer.assert_not_called()
|
||||
|
||||
def test_sizable_local_path_is_sized(self):
|
||||
with (
|
||||
patch.object(self.models_route, "is_local_path", return_value = True),
|
||||
patch.object(self.models_route, "_is_sizable_local_path", return_value = True),
|
||||
patch(
|
||||
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
||||
side_effect = lambda m, **_kw: m,
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (_QWEN35_FP16_BYTES, "local"),
|
||||
),
|
||||
):
|
||||
resp = asyncio.run(
|
||||
self.models_route.get_export_size(
|
||||
model = "/root/.unsloth/studio/outputs/run",
|
||||
hf_token = None,
|
||||
current_subject = "test-user",
|
||||
)
|
||||
)
|
||||
self.assertEqual(resp.fp16_bytes, _QWEN35_FP16_BYTES)
|
||||
self.assertEqual(resp.source, "local")
|
||||
|
||||
def test_local_adapter_base_escaping_roots_is_rejected(self):
|
||||
# A local adapter under a root whose resolved base points outside the
|
||||
# roots (e.g. "/") must not be sized: the resolved base is re-validated.
|
||||
adapter = "/root/.unsloth/studio/outputs/adapter"
|
||||
with (
|
||||
patch.object(self.models_route, "is_local_path", return_value = True),
|
||||
patch.object(
|
||||
self.models_route, "_is_sizable_local_path", side_effect = lambda p: p == adapter
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
||||
return_value = "/",
|
||||
),
|
||||
patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer,
|
||||
):
|
||||
resp = asyncio.run(
|
||||
self.models_route.get_export_size(
|
||||
model = adapter, hf_token = None, current_subject = "test-user"
|
||||
)
|
||||
)
|
||||
self.assertIsNone(resp.fp16_bytes)
|
||||
self.assertEqual(resp.source, "unavailable")
|
||||
mock_sizer.assert_not_called()
|
||||
|
||||
def test_is_sizable_local_path_containment(self):
|
||||
# Only paths under a trusted root are sizable; '..' can't escape.
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp) / "outputs"
|
||||
inside = root / "run-1"
|
||||
inside.mkdir(parents = True)
|
||||
with (
|
||||
patch("utils.paths.studio_root", return_value = root),
|
||||
patch("utils.paths.outputs_root", return_value = root),
|
||||
patch("utils.paths.exports_root", return_value = root),
|
||||
patch("utils.paths.storage_roots.cache_root", return_value = root),
|
||||
):
|
||||
is_sizable = self.models_route._is_sizable_local_path
|
||||
self.assertTrue(is_sizable(str(inside)))
|
||||
self.assertTrue(is_sizable(str(root)))
|
||||
self.assertFalse(is_sizable(str(root / "missing")))
|
||||
self.assertFalse(is_sizable("/etc"))
|
||||
self.assertFalse(is_sizable(str(root / ".." / "etc")))
|
||||
# A symlink inside a root pointing outside it cannot escape.
|
||||
escape = root / "escape"
|
||||
os.symlink(tmp, escape)
|
||||
self.assertFalse(is_sizable(str(escape)))
|
||||
|
||||
def test_local_weight_size_skips_nested_checkpoints(self):
|
||||
# A run dir's intermediate checkpoint-*/global_step* snapshots must not
|
||||
# be counted; only the model files at the root are summed.
|
||||
from utils.hardware.hardware import _get_local_weight_size_bytes
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
run = Path(tmp)
|
||||
(run / "model.safetensors").write_bytes(b"\0" * 1000)
|
||||
for sub, size in (("checkpoint-60", 5000), ("global_step10", 7000)):
|
||||
d = run / sub
|
||||
d.mkdir()
|
||||
(d / "model.safetensors").write_bytes(b"\0" * size)
|
||||
self.assertEqual(_get_local_weight_size_bytes(str(run)), 1000)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -1111,10 +1111,18 @@ def _get_local_weight_size_bytes(model_name: str) -> Optional[int]:
|
|||
return None
|
||||
|
||||
weight_exts = (".safetensors", ".bin", ".pt", ".pth")
|
||||
# Skip intermediate training checkpoints: a run dir can hold several
|
||||
# checkpoint-*/global_step* snapshots, but export loads only the model at
|
||||
# the root, so counting them would multiply the estimate.
|
||||
skip_prefixes = ("checkpoint-", "global_step")
|
||||
total = 0
|
||||
for file in model_path.rglob("*"):
|
||||
if file.is_file() and file.suffix in weight_exts:
|
||||
total += file.stat().st_size
|
||||
if not file.is_file() or file.suffix not in weight_exts:
|
||||
continue
|
||||
rel = file.relative_to(model_path)
|
||||
if any(part.startswith(skip_prefixes) for part in rel.parts):
|
||||
continue
|
||||
total += file.stat().st_size
|
||||
return total if total > 0 else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -33,6 +33,13 @@ export interface CheckpointListResponse {
|
|||
models: ModelCheckpoints[];
|
||||
}
|
||||
|
||||
export interface ExportSizeEstimate {
|
||||
/** Estimated FP16/BF16-equivalent on-disk size, or null when unknown. */
|
||||
fp16_bytes: number | null;
|
||||
total_params: number | null;
|
||||
source: string;
|
||||
}
|
||||
|
||||
export interface ExportOperationResponse {
|
||||
success: boolean;
|
||||
message: string;
|
||||
|
|
@ -48,6 +55,24 @@ export async function fetchCheckpoints(): Promise<CheckpointListResponse> {
|
|||
return parseJson<CheckpointListResponse>(response);
|
||||
}
|
||||
|
||||
/** Estimate a model's fp16-equivalent size to scale the GGUF quant labels; nulls (not error) when unknown. */
|
||||
export async function fetchExportSize(
|
||||
modelId: string,
|
||||
hfToken?: string | null,
|
||||
signal?: AbortSignal,
|
||||
): Promise<ExportSizeEstimate> {
|
||||
// Token in a header (not the query string) so it never lands in URLs/logs.
|
||||
const headers: Record<string, string> = {};
|
||||
if (hfToken) {
|
||||
headers["X-HF-Token"] = hfToken;
|
||||
}
|
||||
const response = await authFetch(
|
||||
`/api/models/export-size?model=${encodeURIComponent(modelId)}`,
|
||||
{ signal, headers },
|
||||
);
|
||||
return parseJson<ExportSizeEstimate>(response);
|
||||
}
|
||||
|
||||
export async function loadCheckpoint(params: {
|
||||
checkpoint_path: string;
|
||||
max_seq_length?: number;
|
||||
|
|
|
|||
|
|
@ -252,7 +252,7 @@ export function ExportDialog({
|
|||
checkpoint,
|
||||
exportMethod,
|
||||
quantLevels,
|
||||
estimatedSize: _estimatedSize,
|
||||
estimatedSize,
|
||||
baseModelName,
|
||||
isAdapter,
|
||||
destination,
|
||||
|
|
@ -576,11 +576,14 @@ export function ExportDialog({
|
|||
</span>
|
||||
</div>
|
||||
)}
|
||||
{/* TODO: unhide once estimated size comes from the backend API */}
|
||||
{/* <div className="flex justify-between">
|
||||
<span>Est. size</span>
|
||||
<span className="font-medium text-foreground">{estimatedSize}</span>
|
||||
</div> */}
|
||||
{estimatedSize && (
|
||||
<div className="flex justify-between">
|
||||
<span>Est. size</span>
|
||||
<span className="font-medium text-foreground">
|
||||
{estimatedSize}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Live export output panel */}
|
||||
|
|
|
|||
|
|
@ -18,9 +18,11 @@ import { QUANT_OPTIONS } from "../constants";
|
|||
interface QuantPickerProps {
|
||||
value: string[];
|
||||
onChange: (v: string[]) => void;
|
||||
/** quant value -> "~X GB"; blank/missing when the model size is unknown. */
|
||||
sizes?: Record<string, string>;
|
||||
}
|
||||
|
||||
export function QuantPicker({ value, onChange }: QuantPickerProps) {
|
||||
export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
|
||||
const toggle = (qv: string) => {
|
||||
onChange(
|
||||
value.includes(qv) ? value.filter((q) => q !== qv) : [...value, qv],
|
||||
|
|
@ -66,6 +68,7 @@ export function QuantPicker({ value, onChange }: QuantPickerProps) {
|
|||
<div className="flex flex-wrap gap-2 py-1 pl-1">
|
||||
{QUANT_OPTIONS.map((q) => {
|
||||
const active = value.includes(q.value);
|
||||
const sizeLabel = sizes?.[q.value] ?? "";
|
||||
return (
|
||||
<button
|
||||
key={q.value}
|
||||
|
|
@ -85,7 +88,9 @@ export function QuantPicker({ value, onChange }: QuantPickerProps) {
|
|||
/>
|
||||
)}
|
||||
{q.label}
|
||||
<span className="text-[10px] opacity-60">{q.size}</span>
|
||||
{sizeLabel && (
|
||||
<span className="text-[10px] opacity-60">{sizeLabel}</span>
|
||||
)}
|
||||
{q.recommended && !active && (
|
||||
<span className="rounded-full bg-emerald-100 px-1.5 py-0 text-[9px] font-semibold text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300">
|
||||
rec
|
||||
|
|
|
|||
|
|
@ -35,39 +35,123 @@ export const EXPORT_METHODS: {
|
|||
},
|
||||
];
|
||||
|
||||
export const QUANT_OPTIONS = [
|
||||
{ value: "q2_k_l", label: "Q2_K_L", size: "~2.9 GB" },
|
||||
{ value: "q3_k_m", label: "Q3_K_M", size: "~3.5 GB" },
|
||||
{ value: "q4_k_m", label: "Q4_K_M", size: "~4.8 GB", recommended: true },
|
||||
{ value: "q5_k_m", label: "Q5_K_M", size: "~5.6 GB" },
|
||||
{ value: "q6_k", label: "Q6_K", size: "~6.6 GB" },
|
||||
{ value: "q8_0", label: "Q8_0", size: "~8.2 GB" },
|
||||
{ value: "bf16", label: "BF16", size: "~14.2 GB" },
|
||||
{ value: "f16", label: "F16", size: "~14.2 GB" },
|
||||
export const QUANT_OPTIONS: {
|
||||
value: string;
|
||||
label: string;
|
||||
recommended?: boolean;
|
||||
}[] = [
|
||||
{ value: "q2_k_l", label: "Q2_K_L" },
|
||||
{ value: "q3_k_m", label: "Q3_K_M" },
|
||||
{ value: "q4_k_m", label: "Q4_K_M", recommended: true },
|
||||
{ value: "q5_k_m", label: "Q5_K_M" },
|
||||
{ value: "q6_k", label: "Q6_K" },
|
||||
{ value: "q8_0", label: "Q8_0" },
|
||||
{ value: "bf16", label: "BF16" },
|
||||
{ value: "f16", label: "F16" },
|
||||
];
|
||||
|
||||
/**
|
||||
* 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
|
||||
* embeddings). Approximate ("~"), not exact file sizes.
|
||||
*/
|
||||
export const GGUF_BPW: Record<string, number> = {
|
||||
q2_k_l: 3.35,
|
||||
q3_k_m: 3.91,
|
||||
q4_k_m: 4.83,
|
||||
q5_k_m: 5.67,
|
||||
q6_k: 6.56,
|
||||
q8_0: 8.5,
|
||||
bf16: 16,
|
||||
f16: 16,
|
||||
};
|
||||
|
||||
const FP16_BPW = 16;
|
||||
|
||||
/**
|
||||
* Human-readable base-1024 size ("67 GB"), matching the model-selector picker.
|
||||
* Do NOT use the base-1000 hub formatBytes here -- it would disagree ("72 GB").
|
||||
*/
|
||||
export function formatModelSize(bytes: number): string {
|
||||
if (!Number.isFinite(bytes) || bytes <= 0) {
|
||||
return "";
|
||||
}
|
||||
const units = ["B", "KB", "MB", "GB", "TB"];
|
||||
// clamp: bytes < 1 would give a negative index
|
||||
const i = Math.max(
|
||||
0,
|
||||
Math.min(Math.floor(Math.log(bytes) / Math.log(1024)), units.length - 1),
|
||||
);
|
||||
const value = bytes / 1024 ** i;
|
||||
return `${value.toFixed(value < 10 ? 1 : 0)} ${units[i]}`;
|
||||
}
|
||||
|
||||
/** Estimated on-disk bytes for one GGUF quant, scaled from the real fp16 size. */
|
||||
export function estimateQuantBytes(
|
||||
fp16Bytes: number | null | undefined,
|
||||
quant: string,
|
||||
): number | null {
|
||||
if (!fp16Bytes || fp16Bytes <= 0) {
|
||||
return null;
|
||||
}
|
||||
const bpw = GGUF_BPW[quant];
|
||||
if (bpw == null) {
|
||||
return null;
|
||||
}
|
||||
return fp16Bytes * (bpw / FP16_BPW);
|
||||
}
|
||||
|
||||
/** "~X GB" label for a quant, or "" when the real model size is unknown. */
|
||||
export function formatQuantSize(
|
||||
fp16Bytes: number | null | undefined,
|
||||
quant: string,
|
||||
): string {
|
||||
const bytes = estimateQuantBytes(fp16Bytes, quant);
|
||||
return bytes == null ? "" : `~${formatModelSize(bytes)}`;
|
||||
}
|
||||
|
||||
/** value -> "~X GB" for every quant option (blank when size unknown). */
|
||||
export function buildQuantSizeLabels(
|
||||
fp16Bytes: number | null | undefined,
|
||||
): Record<string, string> {
|
||||
const out: Record<string, string> = {};
|
||||
for (const q of QUANT_OPTIONS) {
|
||||
out[q.value] = formatQuantSize(fp16Bytes, q.value);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* Estimated total export size for the summary line; scales from the model's
|
||||
* real fp16 size, returns "" when unknown so the UI can hide a wrong number.
|
||||
*/
|
||||
export function getEstimatedSize(
|
||||
method: ExportMethod | null,
|
||||
quantLevels: string[],
|
||||
) {
|
||||
const sizeOf = (v: string) =>
|
||||
QUANT_OPTIONS.find((q) => q.value === v)?.size ?? "—";
|
||||
fp16Bytes: number | null | undefined,
|
||||
): string {
|
||||
if (method === "gguf" && quantLevels.length > 0) {
|
||||
if (quantLevels.length === 1) {
|
||||
return sizeOf(quantLevels[0]);
|
||||
const perQuant = quantLevels.map((q) => estimateQuantBytes(fp16Bytes, q));
|
||||
if (perQuant.some((b) => b == null)) {
|
||||
return ""; // unknown -> blank
|
||||
}
|
||||
const total = quantLevels
|
||||
.map((q) => Number.parseFloat(sizeOf(q).replace(/[^0-9.]/g, "")))
|
||||
.reduce((a, b) => a + b, 0);
|
||||
return `~${total.toFixed(1)} GB (${quantLevels.length} files)`;
|
||||
let total = 0;
|
||||
for (const b of perQuant) {
|
||||
total += b ?? 0;
|
||||
}
|
||||
const label = `~${formatModelSize(total)}`;
|
||||
return quantLevels.length === 1
|
||||
? label
|
||||
: `${label} (${quantLevels.length} files)`;
|
||||
}
|
||||
if (method === "merged") {
|
||||
return "~14.2 GB";
|
||||
return fp16Bytes && fp16Bytes > 0 ? `~${formatModelSize(fp16Bytes)}` : "";
|
||||
}
|
||||
if (method === "lora") {
|
||||
// Adapter size is bounded by LoRA rank, not the base model size.
|
||||
return "~100 MB";
|
||||
}
|
||||
return "—";
|
||||
return "";
|
||||
}
|
||||
|
||||
export const METHOD_LABELS: Record<TrainingMethod, string> = {
|
||||
|
|
|
|||
|
|
@ -71,8 +71,10 @@ import { QuantPicker } from "./components/quant-picker";
|
|||
import {
|
||||
type ExportMethod,
|
||||
GUIDE_STEPS,
|
||||
buildQuantSizeLabels,
|
||||
getEstimatedSize,
|
||||
} from "./constants";
|
||||
import { useExportSizeEstimate } from "./hooks/use-export-size-estimate";
|
||||
import { GuidedTour, useGuidedTourController } from "@/features/tour";
|
||||
import { exportTourSteps } from "./tour";
|
||||
|
||||
|
|
@ -243,6 +245,26 @@ export function ExportPage() {
|
|||
? selectedSourceModel ?? "—"
|
||||
: baseModelName;
|
||||
|
||||
// For a full fine-tune checkpoint the weights live in the checkpoint dir
|
||||
// itself (its base_model may be a local/custom path that can't be sized), so
|
||||
// size that dir; for LoRA adapters the export merges into the base model.
|
||||
const sizeTargetModel = useMemo(() => {
|
||||
if (sourceMode === "checkpoint" && !isAdapter) {
|
||||
const cp = checkpointsForModel.find((c) => c.display_name === checkpoint);
|
||||
if (cp?.path) {
|
||||
return cp.path;
|
||||
}
|
||||
}
|
||||
return sourceBaseModelName;
|
||||
}, [sourceMode, isAdapter, checkpointsForModel, checkpoint, sourceBaseModelName]);
|
||||
|
||||
// Real (MoE-aware) fp16 size, used to scale the GGUF quant estimates.
|
||||
const { fp16Bytes } = useExportSizeEstimate(sizeTargetModel, debouncedHfToken);
|
||||
const quantSizeLabels = useMemo(
|
||||
() => buildQuantSizeLabels(fp16Bytes),
|
||||
[fp16Bytes],
|
||||
);
|
||||
|
||||
const {
|
||||
results: hfResults,
|
||||
isLoading: isLoadingHfModels,
|
||||
|
|
@ -367,7 +389,7 @@ export function ExportPage() {
|
|||
}
|
||||
};
|
||||
|
||||
const estimatedSize = getEstimatedSize(exportMethod, quantLevels);
|
||||
const estimatedSize = getEstimatedSize(exportMethod, quantLevels, fp16Bytes);
|
||||
const selectedExportSource =
|
||||
sourceMode === "checkpoint" ? checkpoint : selectedSourceModel;
|
||||
const defaultSaveDirectory = useMemo(() => {
|
||||
|
|
@ -1087,21 +1109,28 @@ export function ExportPage() {
|
|||
<AnimatePresence>
|
||||
{exportMethod === "gguf" && (
|
||||
<motion.div {...collapseAnim} className="overflow-visible">
|
||||
<QuantPicker value={quantLevels} onChange={setQuantLevels} />
|
||||
<QuantPicker
|
||||
value={quantLevels}
|
||||
onChange={setQuantLevels}
|
||||
sizes={quantSizeLabels}
|
||||
/>
|
||||
</motion.div>
|
||||
)}
|
||||
</AnimatePresence>
|
||||
|
||||
<Separator />
|
||||
<div className="flex items-center justify-end">
|
||||
{/* TODO: unhide once estimated size comes from the backend API */}
|
||||
{/* <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3.5"
|
||||
/>
|
||||
<span>Est. size: {estimatedSize} · Free disk space: 120 GB</span>
|
||||
</div> */}
|
||||
<div className="flex items-center justify-between">
|
||||
{estimatedSize ? (
|
||||
<div className="flex items-center gap-1.5 text-xs text-muted-foreground">
|
||||
<HugeiconsIcon
|
||||
icon={InformationCircleIcon}
|
||||
className="size-3.5"
|
||||
/>
|
||||
<span>Est. size: {estimatedSize}</span>
|
||||
</div>
|
||||
) : (
|
||||
<span />
|
||||
)}
|
||||
<Button
|
||||
data-tour="export-cta"
|
||||
disabled={!canExport}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,68 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { useEffect, useState } from "react";
|
||||
import { type ExportSizeEstimate, fetchExportSize } from "../api/export-api";
|
||||
|
||||
export interface ExportSizeState {
|
||||
data: ExportSizeEstimate | null;
|
||||
loading: boolean;
|
||||
/** Real FP16/BF16-equivalent bytes, or null when unknown. */
|
||||
fp16Bytes: number | null;
|
||||
}
|
||||
|
||||
const EMPTY: ExportSizeState = { data: null, loading: false, fp16Bytes: null };
|
||||
|
||||
// Export page's "no model" sentinel (em dash U+2014); char code keeps this file ASCII.
|
||||
const EMPTY_MODEL_SENTINEL = String.fromCharCode(0x2014);
|
||||
|
||||
function normalizeModelId(modelId: string | null | undefined): string {
|
||||
const id = (modelId ?? "").trim();
|
||||
return id && id !== EMPTY_MODEL_SENTINEL ? id : "";
|
||||
}
|
||||
|
||||
/**
|
||||
* Fetch the selected model's fp16-equivalent size so the GGUF picker can scale
|
||||
* its per-quant estimates. Null size on any failure; refetches on token change.
|
||||
*/
|
||||
export function useExportSizeEstimate(
|
||||
modelId: string | null | undefined,
|
||||
hfToken?: string | null,
|
||||
): ExportSizeState {
|
||||
const modelKey = normalizeModelId(modelId);
|
||||
const token = hfToken?.trim() || "";
|
||||
// Refetch when model or token changes; "|" can't appear in an id/token.
|
||||
const key = modelKey ? `${modelKey}|${token}` : "";
|
||||
const [state, setState] = useState<{ key: string; value: ExportSizeState }>(
|
||||
() => ({ key: "", value: EMPTY }),
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (!modelKey) {
|
||||
setState({ key: "", value: EMPTY });
|
||||
return;
|
||||
}
|
||||
const controller = new AbortController();
|
||||
setState({ key, value: { ...EMPTY, loading: true } });
|
||||
void fetchExportSize(modelKey, token, controller.signal)
|
||||
.then((data) => {
|
||||
if (controller.signal.aborted) {
|
||||
return;
|
||||
}
|
||||
setState({
|
||||
key,
|
||||
value: { data, loading: false, fp16Bytes: data.fp16_bytes ?? null },
|
||||
});
|
||||
})
|
||||
.catch(() => {
|
||||
if (controller.signal.aborted) {
|
||||
return;
|
||||
}
|
||||
setState({ key, value: EMPTY });
|
||||
});
|
||||
return () => controller.abort();
|
||||
}, [key, modelKey, token]);
|
||||
|
||||
// Guard against a stale result rendering for a newer key.
|
||||
return state.key === key ? state.value : EMPTY;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue