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

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

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:

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;

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<HTMLDivElement>(null);
   const localComboboxAnchorRef = useRef<HTMLDivElement>(null);

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 {

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<TrainingMethod | null> {
+  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<keyof TrainingConfigState> = new Set
   "isDatasetAudio",
   "trainOnCompletions",
   "maxPositionEmbeddings",
+  "vramEstimateQloraGb",
+  "vramEstimateLoraGb",
+  "vramEstimateFullGb",
 ]);
 
 function partializePersistedState(
@@ -213,22 +215,28 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             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<TrainingConfigStore>()(
               isLoadingModelDefaults: false,
               modelDefaultsError: null,
               modelDefaultsAppliedFor: null,
+              vramEstimateQloraGb: null,
+              vramEstimateLoraGb: null,
+              vramEstimateFullGb: null,
             });
             return;
           }

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 {

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