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
This commit is contained in:
Roland Tannous
parent
6984e118eb
commit
d96e3a7096
8 changed files with 200 additions and 100 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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 });
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue