From d96e3a7096ccbc0096be2f70f35adae0a4fe6f46 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 1 Apr 2026 04:55:51 +0000 Subject: [PATCH] fix(vram): use training-aware estimates and backend arch-based VRAM for model fitness - Replace loading-only VRAM formula with full training estimate (weights + LoRA adapters + optimizer states + gradients + activations + overhead) for all three methods: QLoRA, LoRA, full fine-tuning - Expose architecture-based VRAM estimates from backend /api/models/config, reusing already-loaded AutoConfig to avoid extra HF round-trip - Store per-method estimates in training config state; selected model badge uses authoritative backend estimate (handles MoE like gpt-oss-20b correctly) - Replace file-size heuristic in autoSelectTrainingMethod with backend estimates - Use total VRAM (not free) since chat models are offloaded before training --- studio/backend/models/models.py | 9 ++ studio/backend/routes/models.py | 74 +++++++++++-- .../assistant-ui/model-selector/pickers.tsx | 6 +- .../studio/sections/model-section.tsx | 35 +++++- .../src/features/training/api/models-api.ts | 3 + .../training/stores/training-config-store.ts | 77 +++++++------ .../src/features/training/types/config.ts | 3 + studio/frontend/src/lib/vram.ts | 101 +++++++++--------- 8 files changed, 204 insertions(+), 104 deletions(-) diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index f67014a17b..3f2ef63fbe 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -95,6 +95,15 @@ class ModelDetails(BaseModel): model_size_bytes: Optional[int] = Field( None, description = "Total size of model weight files in bytes" ) + vram_estimate_qlora_gb: Optional[float] = Field( + None, description = "Estimated training VRAM (GB) for QLoRA (4-bit) with default params" + ) + vram_estimate_lora_gb: Optional[float] = Field( + None, description = "Estimated training VRAM (GB) for LoRA (fp16) with default params" + ) + vram_estimate_full_gb: Optional[float] = Field( + None, description = "Estimated training VRAM (GB) for full fine-tuning with default params" + ) class LoRAInfo(BaseModel): diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 445cf0e7f4..adbef40ba2 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -585,6 +585,52 @@ def _get_model_size_bytes( return None +def _compute_vram_estimates( + hf_config, +) -> tuple[Optional[float], Optional[float], Optional[float]]: + """Estimate training VRAM (GB) for qlora, lora, and full fine-tuning. + + Uses architecture-based estimation (weights + LoRA + optimizer + gradients + + activations + overhead) with default training params: + batch=4, seq=2048, rank=16, adamw_8bit, unsloth gradient checkpointing. + + Returns (qlora_gb, lora_gb, full_gb). Returns (None, None, None) on failure. + """ + try: + from utils.hardware.vram_estimation import ( + DEFAULT_TARGET_MODULES, + TrainingVramConfig, + estimate_training_vram, + extract_arch_config, + ) + + if hf_config is None: + return None, None, None + + arch = extract_arch_config(hf_config) + if arch is None: + return None, None, None + + estimates = [] + for method, load_in_4bit in (("qlora", True), ("lora", False), ("full", False)): + vram_config = TrainingVramConfig( + training_method = method, + batch_size = 4, + max_seq_length = 2048, + lora_rank = 16, + target_modules = list(DEFAULT_TARGET_MODULES), + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = load_in_4bit, + ) + breakdown = estimate_training_vram(arch, vram_config) + estimates.append(round(breakdown.total / (1024 ** 3), 3)) + return estimates[0], estimates[1], estimates[2] + except Exception as e: + logger.warning(f"Could not compute VRAM estimates: {e}") + return None, None, None + + @router.get("/config/{model_name:path}") async def get_model_config( model_name: str, @@ -625,21 +671,24 @@ async def get_model_config( except Exception: pass - # Fallback: try AutoConfig directly if not found yet - if max_position_embeddings is None: - try: - from transformers import AutoConfig as _AutoConfig + # Load AutoConfig — used for max_position_embeddings (fallback) and VRAM estimation. + _ac = None + try: + from transformers import AutoConfig as _AutoConfig - _trust = model_name.lower().startswith("unsloth/") - _ac = _AutoConfig.from_pretrained( - model_name, trust_remote_code = _trust, token = hf_token - ) + _trust = model_name.lower().startswith("unsloth/") + _ac = _AutoConfig.from_pretrained( + model_name, trust_remote_code = _trust, token = hf_token + ) + if max_position_embeddings is None: max_position_embeddings = _get_max_position_embeddings(_ac) - except Exception: - pass + except Exception: + pass + + vram_qlora, vram_lora, vram_full = _compute_vram_estimates(_ac) logger.info( - f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}, max_position_embeddings={max_position_embeddings}" + f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}, max_position_embeddings={max_position_embeddings}, vram_qlora={vram_qlora}, vram_lora={vram_lora}, vram_full={vram_full}" ) return ModelDetails( id = model_name, @@ -655,6 +704,9 @@ async def get_model_config( base_model = base_model, max_position_embeddings = max_position_embeddings, model_size_bytes = _get_model_size_bytes(model_name, hf_token), + vram_estimate_qlora_gb = vram_qlora, + vram_estimate_lora_gb = vram_lora, + vram_estimate_full_gb = vram_full, ) except Exception as e: diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index cf8b4cd54e..51f1a125eb 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -41,7 +41,7 @@ import { } from "@/hooks"; import { cn, formatCompact } from "@/lib/utils"; import type { VramFitStatus } from "@/lib/vram"; -import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; +import { checkVramFit, estimateTrainingVram } from "@/lib/vram"; import { Search01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Trash2Icon } from "lucide-react"; @@ -659,7 +659,7 @@ export function HubModelPicker({ for (const r of results) { const detail = r.totalParams ? formatCompact(r.totalParams) : null; if (r.totalParams) { - const est = estimateLoadingVram(r.totalParams, "qlora"); + const est = estimateTrainingVram(r.totalParams, "qlora"); const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null; @@ -680,7 +680,7 @@ export function HubModelPicker({ for (const id of ids) { const totalParams = recommendedParamCountById.get(id); if (totalParams) { - const est = estimateLoadingVram(totalParams, "qlora"); + const est = estimateTrainingVram(totalParams, "qlora"); const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null; diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 755c93c5f0..f93fdf4f13 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -50,6 +50,7 @@ import { type VramFitStatus, type TrainingMethod as VramTrainingMethod, buildModelVramMap, + checkVramFit, } from "@/lib/vram"; import type { TrainingMethod } from "@/types/training"; import { @@ -94,6 +95,9 @@ export function ModelSection() { setTrainingMethod, hfToken, setHfToken, + vramEstimateQloraGb, + vramEstimateLoraGb, + vramEstimateFullGb, } = useTrainingConfigStore( useShallow( ({ @@ -104,6 +108,9 @@ export function ModelSection() { setTrainingMethod, hfToken, setHfToken, + vramEstimateQloraGb, + vramEstimateLoraGb, + vramEstimateFullGb, }) => ({ modelType, selectedModel, @@ -112,6 +119,9 @@ export function ModelSection() { setTrainingMethod, hfToken, setHfToken, + vramEstimateQloraGb, + vramEstimateLoraGb, + vramEstimateFullGb, }), ), ); @@ -230,12 +240,12 @@ export function ModelSection() { }); }, [localMetaById, localModelInput, localResultIds]); - // Pre-compute VRAM fit status for every model in the current result set. - // Keyed by model id so the render callback is a simple O(1) lookup. - // // Pre-compute VRAM fit status for every model in the current result set. // Keyed by model id so the render callback is a simple O(1) lookup. // Re-computes when the training method changes (QLoRA=4-bit vs LoRA/Full=fp16). + // + // For the currently selected model, we override the frontend estimate with + // the authoritative backend estimate (architecture-aware, handles MoE). const vramMap = useMemo(() => { const fitMap = buildModelVramMap( hfResults, @@ -251,6 +261,23 @@ export function ModelSection() { ? formatCompact(r.totalParams) : extractParamLabel(r.id); const fit = fitMap.get(r.id); + + // Override with backend estimate for the selected model. + if (r.id === selectedModel) { + const backendEst = + trainingMethod === "qlora" ? vramEstimateQloraGb + : trainingMethod === "lora" ? vramEstimateLoraGb + : vramEstimateFullGb; + if (backendEst != null) { + map.set(r.id, { + est: backendEst, + status: gpu.available ? checkVramFit(backendEst, gpu.memoryTotalGb) : null, + detail, + }); + continue; + } + } + map.set(r.id, { est: fit?.est ?? 0, status: fit?.status ?? null, @@ -258,7 +285,7 @@ export function ModelSection() { }); } return map; - }, [hfResults, gpu, trainingMethod]); + }, [hfResults, gpu, trainingMethod, selectedModel, vramEstimateQloraGb, vramEstimateLoraGb, vramEstimateFullGb]); const comboboxAnchorRef = useRef(null); const localComboboxAnchorRef = useRef(null); diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index e512b9e28d..dad3485243 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -73,6 +73,9 @@ export interface ModelConfigResponse { model_type?: "text" | "vision" | "audio" | "embeddings" | null; max_position_embeddings?: number | null; model_size_bytes?: number | null; + vram_estimate_qlora_gb?: number | null; + vram_estimate_lora_gb?: number | null; + vram_estimate_full_gb?: number | null; } export interface LocalModelInfo { diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 8214b0eb2a..065a08ce89 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -17,33 +17,29 @@ const MIN_STEP: StepNumber = 1; const MAX_STEP: StepNumber = STEPS.length as StepNumber; /** - * Auto-select LoRA (16-bit) vs QLoRA (4-bit) based on model size and GPU memory. + * Auto-select LoRA (16-bit) vs QLoRA (4-bit) based on backend VRAM estimates and total GPU VRAM. * - * Rule: if model_size_gb * 1.5 * context_scale fits in free VRAM, use "lora" (16-bit). - * Otherwise use "qlora" (4-bit). + * Uses architecture-aware estimates from the backend (weights + LoRA adapters + + * optimizer states + gradients + activations + CUDA overhead) rather than a + * file-size heuristic. We use total VRAM (not free) because any loaded models + * (e.g. from the chat tab) are offloaded before training begins. * - * Context scale: <=8192 = 1.0, >8192 = 1.7, >=16384 = 2.0, >=32768 = 4.0 + * Full fine-tuning is never auto-selected — the user must opt in explicitly. */ async function autoSelectTrainingMethod( - modelSizeBytes: number, - contextLength: number, + vramEstimateLoraGb: number | null, + vramEstimateQloraGb: number | null, ): Promise { + if (vramEstimateLoraGb == null && vramEstimateQloraGb == null) return null; try { const res = await authFetch("/api/system/hardware"); if (!res.ok) return null; const data = await res.json(); - const freeGb: number | null = data?.gpu?.vram_free_gb ?? null; - if (freeGb == null) return null; + const totalGb: number | null = data?.gpu?.vram_total_gb ?? null; + if (totalGb == null) return null; - const modelSizeGb = modelSizeBytes / (1024 ** 3); - - let contextScale = 1.0; - if (contextLength >= 32768) contextScale = 4.0; - else if (contextLength >= 16384) contextScale = 2.0; - else if (contextLength > 8192) contextScale = 1.7; - - const estimatedUsage = modelSizeGb * 1.5 * contextScale; - return estimatedUsage <= freeGb ? "lora" : "qlora"; + if (vramEstimateLoraGb != null && vramEstimateLoraGb <= totalGb) return "lora"; + return "qlora"; } catch { return null; } @@ -87,6 +83,9 @@ const initialState: TrainingConfigState = { isDatasetAudio: false, maxPositionEmbeddings: null, ...DEFAULT_HYPERPARAMS, + vramEstimateQloraGb: null, + vramEstimateLoraGb: null, + vramEstimateFullGb: null, }; // AbortController for in-flight dataset multimodal checks. @@ -121,6 +120,9 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set "isDatasetAudio", "trainOnCompletions", "maxPositionEmbeddings", + "vramEstimateQloraGb", + "vramEstimateLoraGb", + "vramEstimateFullGb", ]); function partializePersistedState( @@ -213,22 +215,28 @@ export const useTrainingConfigStore = create()( const inferredModelType: ModelType = modelDetails.model_type ?? (isEmbedding ? "embeddings" : modelDetails.is_vision ? "vision" : modelDetails.is_audio ? "audio" : "text"); - // Auto-select training method based on model size vs GPU memory. - // If model_size * 1.5 * context_scale fits in free VRAM, use LoRA 16-bit. - // Otherwise use QLoRA 4-bit. - const modelSizeBytes = modelDetails.model_size_bytes; - if (modelSizeBytes && modelSizeBytes > 0) { - void autoSelectTrainingMethod(modelSizeBytes, patch.contextLength ?? get().contextLength) - .then((method) => { - if (get().selectedModel !== modelName) return; - if (method) { - const lrPatch = !_learningRateManuallySet && !modelConfigHasLR - ? { learningRate: method === "full" ? LR_DEFAULT_FULL : LR_DEFAULT_LORA } - : {}; - set({ trainingMethod: method, ...lrPatch }); - } - }); - } + // Store backend VRAM estimates for all three training methods. + const vramQlora = modelDetails.vram_estimate_qlora_gb ?? null; + const vramLora = modelDetails.vram_estimate_lora_gb ?? null; + const vramFull = modelDetails.vram_estimate_full_gb ?? null; + set({ + vramEstimateQloraGb: vramQlora, + vramEstimateLoraGb: vramLora, + vramEstimateFullGb: vramFull, + }); + + // Auto-select LoRA vs QLoRA using architecture-aware VRAM estimates. + // Full fine-tuning is never auto-selected; the user must opt in explicitly. + void autoSelectTrainingMethod(vramLora, vramQlora) + .then((method) => { + if (get().selectedModel !== modelName) return; + if (method) { + const lrPatch = !_learningRateManuallySet && !modelConfigHasLR + ? { learningRate: method === "full" ? LR_DEFAULT_FULL : LR_DEFAULT_LORA } + : {}; + set({ trainingMethod: method, ...lrPatch }); + } + }); set({ ...patch, @@ -377,6 +385,9 @@ export const useTrainingConfigStore = create()( isLoadingModelDefaults: false, modelDefaultsError: null, modelDefaultsAppliedFor: null, + vramEstimateQloraGb: null, + vramEstimateLoraGb: null, + vramEstimateFullGb: null, }); return; } diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 2d19dea874..953cc3d928 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -81,6 +81,9 @@ export interface TrainingConfigState { finetuneMLPModules: boolean; targetModules: string[]; maxPositionEmbeddings: number | null; + vramEstimateQloraGb: number | null; + vramEstimateLoraGb: number | null; + vramEstimateFullGb: number | null; } export interface TrainingConfigActions { diff --git a/studio/frontend/src/lib/vram.ts b/studio/frontend/src/lib/vram.ts index a5abebb043..d76a0821c7 100644 --- a/studio/frontend/src/lib/vram.ts +++ b/studio/frontend/src/lib/vram.ts @@ -2,80 +2,75 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 /** - * VRAM estimation for model loading (4-bit quantization via bitsandbytes). + * Training VRAM estimation for model fitness badges. * - * Estimates the total driver-level VRAM (what nvidia-smi reports) needed to - * load a model in 4-bit with Unsloth / bitsandbytes. This determines - * whether a model will fit on the user's GPU before any training begins. + * Estimates total GPU VRAM needed to *train* a model with Unsloth — not just + * load it. This accounts for model weights, LoRA adapters, optimizer states, + * gradients, activations, and CUDA overhead. * - * Formula: totalParams * 0.90 + 1.4 GB + * The formulas mirror the backend fallback path in hardware.py and are + * calibrated for default training params: + * batch=4, seq=2048, rank=16, adamw_8bit, unsloth gradient checkpointing * - * Calibrated against isolated Unsloth model loads on RTX 5070 Ti (2026.2): - * Qwen2.5-0.5B (0.49B) : est 1.8 vs actual 1.86 GB (-3%) - * Llama-3.2-1B (1.24B) : est 2.5 vs actual 2.54 GB (-1%) - * Llama-3.2-3B (3.21B) : est 4.3 vs actual 4.40 GB (-2%) - * Llama-3.1-8B (8.03B) : est 8.6 vs actual 8.14 GB (+6%) + * For models in the search dropdown, `totalParams` (from HF safetensors + * metadata) is used as input. For MoE models this includes all expert + * parameters, which can overestimate; once a model is selected the + * authoritative architecture-based backend estimate is used instead (stored + * in the training config store as `vramEstimate*Gb`). * - * Accuracy: within 3% for 0.5B-3B models, within 6% for 8B. + * Constants: + * QUANT_4BIT_FACTOR - fp16 → bnb 4-bit compression ratio (16/5 = 3.2×) + * LOADING_OVERHEAD_GB - CUDA driver + PyTorch runtime baseline (~1.4 GB) */ // --------------------------------------------------------------------------- // Constants (exported for testing) // --------------------------------------------------------------------------- -/** - * Effective bytes per parameter for 4-bit model weights at driver level. - * - * Raw bnb 4-bit is ~0.5 bytes/param, but embedding and lm_head layers remain - * in fp16 and bnb adds per-block quantization metadata, bringing the - * effective rate to ~0.84-0.93 across tested architectures. 0.9 is the - * calibrated middle ground. - */ -export const BNB_4BIT_LOADING_BYTES = 0.9; - -/** - * Fixed overhead (GB) for the CUDA driver context and PyTorch runtime. - * - * This is independent of model size -- it is the baseline GPU memory consumed - * before any model weights are loaded. Measured at 1.34-1.46 GB across - * tested models; we use 1.4 as the default. - */ +export const QUANT_4BIT_FACTOR = 16 / 5; // 3.2× export const LOADING_OVERHEAD_GB = 1.4; +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export type VramFitStatus = "fits" | "tight" | "exceeds"; + +export type TrainingMethod = "qlora" | "lora" | "full"; + // --------------------------------------------------------------------------- // Estimation // --------------------------------------------------------------------------- -export type VramFitStatus = "fits" | "tight" | "exceeds"; - /** - * Bytes per parameter when loading a model at fp16/bf16 (LoRA, full FT). + * Estimate total training VRAM (GB) needed for a model given its parameter + * count, using the backend's fallback formula. * - * This is the theoretical value (2 bytes = 16 bits). Not yet calibrated - * against actual measurements -- the real driver-level usage may be slightly - * higher due to buffers and metadata, similar to how 4-bit is 0.9 vs 0.5. + * Where fp16Gb = totalParams × 2 bytes / 1e9: + * QLoRA : fp16Gb / 3.2 + fp16Gb × 0.04 + fp16Gb × 0.15 + 1.4 + * LoRA : fp16Gb + fp16Gb × 0.04 + fp16Gb × 0.15 + 1.4 + * Full : fp16Gb × 3.5 + 1.4 + * + * The 0.04 term is LoRA adapter + optimizer overhead. + * The 0.15 term is activations (with unsloth gradient checkpointing, batch=4). + * The 3.5× factor for full FT covers weights + optimizer (Adam: ×3) + gradients. */ -export const FP16_LOADING_BYTES = 2.0; - -export type TrainingMethod = "qlora" | "lora" | "full"; - -/** - * Estimate VRAM (GB) needed to load a model with Unsloth. - * - * The bytes-per-param rate depends on the training method: - * - QLoRA : 4-bit quantized via bnb -> 0.90 bytes/param (calibrated) - * - LoRA : fp16 -> 2.0 bytes/param (theoretical) - * - Full : fp16 -> 2.0 bytes/param (theoretical) - * - * Formula: totalParams * bytesPerParam + 1.4 GB overhead - */ -export function estimateLoadingVram( +export function estimateTrainingVram( totalParams: number, method: TrainingMethod = "qlora", ): number { - const bytesPerParam = - method === "qlora" ? BNB_4BIT_LOADING_BYTES : FP16_LOADING_BYTES; - const gb = (totalParams / 1e9) * bytesPerParam + LOADING_OVERHEAD_GB; + const fp16Gb = (totalParams * 2) / 1e9; + + let gb: number; + if (method === "qlora") { + gb = fp16Gb / QUANT_4BIT_FACTOR + fp16Gb * 0.04 + fp16Gb * 0.15 + LOADING_OVERHEAD_GB; + } else if (method === "lora") { + gb = fp16Gb + fp16Gb * 0.04 + fp16Gb * 0.15 + LOADING_OVERHEAD_GB; + } else { + // full fine-tuning + gb = fp16Gb * 3.5 + LOADING_OVERHEAD_GB; + } + return Math.round(gb * 10) / 10; } @@ -119,7 +114,7 @@ export function buildModelVramMap( continue; } - const est = estimateLoadingVram(model.totalParams, method); + const est = estimateTrainingVram(model.totalParams, method); const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null; map.set(model.id, { est, status }); }