85314ed162Trim and tighten code comments across studio/frontend TS/JS. Comment-only: every changed file verified code-identical to main via the TypeScript printer signature comparison.
112 lines
3.5 KiB
TypeScript
112 lines
3.5 KiB
TypeScript
// SPDX-License-Identifier: AGPL-3.0-only
|
|
// 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).
|
|
*
|
|
* Estimates the total driver-level VRAM (what nvidia-smi reports) to load a
|
|
* model in 4-bit with Unsloth / bitsandbytes, to check it fits the GPU before
|
|
* training.
|
|
*
|
|
* Formula: totalParams * 0.90 + 1.4 GB
|
|
*
|
|
* Calibrated against isolated Unsloth 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%)
|
|
*
|
|
* Accuracy: within 3% for 0.5B-3B, within 6% for 8B.
|
|
*/
|
|
|
|
// Constants (exported for testing)
|
|
|
|
/**
|
|
* Effective bytes per parameter for 4-bit weights at driver level. Raw bnb
|
|
* 4-bit is ~0.5, but embedding/lm_head stay fp16 and bnb adds per-block
|
|
* metadata, giving ~0.84-0.93 across architectures; 0.9 is the calibrated mid.
|
|
*/
|
|
export const BNB_4BIT_LOADING_BYTES = 0.9;
|
|
|
|
/**
|
|
* Fixed overhead (GB) for the CUDA driver context and PyTorch runtime,
|
|
* independent of model size. Measured at 1.34-1.46 GB; we use 1.4.
|
|
*/
|
|
export const LOADING_OVERHEAD_GB = 1.4;
|
|
|
|
export type VramFitStatus = "fits" | "tight" | "exceeds";
|
|
|
|
/**
|
|
* Bytes per parameter at fp16/bf16 (LoRA, full FT). Theoretical (2 bytes = 16
|
|
* bits); not yet calibrated, so real usage may run slightly higher (as 4-bit
|
|
* is 0.9 vs 0.5).
|
|
*/
|
|
export const FP16_LOADING_BYTES = 2.0;
|
|
|
|
export type TrainingMethod = "qlora" | "lora" | "full" | "cpt";
|
|
|
|
function usesQuantizedLoading(method: TrainingMethod, modelId?: string): boolean {
|
|
if (method === "qlora") return true;
|
|
return method === "cpt" && (modelId ?? "").toLowerCase().includes("4bit");
|
|
}
|
|
|
|
/**
|
|
* Estimate VRAM (GB) to load a model with Unsloth. Bytes/param by method:
|
|
* QLoRA: 4-bit bnb -> 0.90 (calibrated); LoRA/Full/CPT: fp16 -> 2.0
|
|
* (theoretical). Formula: totalParams * bytesPerParam + 1.4 GB overhead.
|
|
*/
|
|
export function estimateLoadingVram(
|
|
totalParams: number,
|
|
method: TrainingMethod = "qlora",
|
|
modelId?: string,
|
|
): number {
|
|
const bytesPerParam = usesQuantizedLoading(method, modelId)
|
|
? BNB_4BIT_LOADING_BYTES
|
|
: FP16_LOADING_BYTES;
|
|
const gb = (totalParams / 1e9) * bytesPerParam + LOADING_OVERHEAD_GB;
|
|
return Math.round(gb * 10) / 10;
|
|
}
|
|
|
|
/**
|
|
* Check whether a model fits in available GPU VRAM.
|
|
* fits: <= 75%; tight: 75-100%; exceeds: > 100%.
|
|
*/
|
|
export function checkVramFit(
|
|
requiredGb: number,
|
|
availableGb: number,
|
|
): VramFitStatus {
|
|
if (availableGb <= 0) return requiredGb <= 0 ? "fits" : "exceeds";
|
|
const ratio = requiredGb / availableGb;
|
|
if (ratio <= 0.75) return "fits";
|
|
if (ratio <= 1.0) return "tight";
|
|
return "exceeds";
|
|
}
|
|
|
|
export interface ModelVramMapInput {
|
|
id: string;
|
|
totalParams?: number;
|
|
}
|
|
|
|
export interface ModelVramMapEntry {
|
|
est: number;
|
|
status: VramFitStatus | null;
|
|
}
|
|
|
|
export function buildModelVramMap(
|
|
models: ModelVramMapInput[],
|
|
method: TrainingMethod,
|
|
gpu: { available: boolean; memoryTotalGb: number },
|
|
): Map<string, ModelVramMapEntry> {
|
|
const map = new Map<string, ModelVramMapEntry>();
|
|
for (const model of models) {
|
|
if (!model.totalParams) {
|
|
map.set(model.id, { est: 0, status: null });
|
|
continue;
|
|
}
|
|
|
|
const est = estimateLoadingVram(model.totalParams, method, model.id);
|
|
const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null;
|
|
map.set(model.id, { est, status });
|
|
}
|
|
return map;
|
|
}
|