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 8a52747409..3d91cffb8c 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -103,9 +103,6 @@ function ModelRow({ {vramStatus === "tight" && ( TIGHT )} - {vramStatus === "fits" && ( - FIT - )} {meta ? ( {meta} ) : null} @@ -253,9 +250,6 @@ function GgufVariantExpander({ {fitStatus === "tight" && ( TIGHT )} - {fitStatus === "fits" && ( - FIT - )} {formatBytes(v.size_bytes)} diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 5cf4ebc0b5..484a11ea70 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -33,11 +33,17 @@ import { import { MODEL_TYPE_TO_HF_TASK } from "@/config/training"; import { useDebouncedValue, + useGpuInfo, useHfModelSearch, useHfTokenValidation, useInfiniteScroll, } from "@/hooks"; import { formatCompact } from "@/lib/utils"; +import { + type TrainingMethod as VramTrainingMethod, + type VramFitStatus, + buildModelVramMap, +} from "@/lib/vram"; import { useTrainingConfigStore } from "@/features/training"; import type { TrainingMethod } from "@/types/training"; import { @@ -50,6 +56,7 @@ import { useEffect, useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; export function ModelSelectionStep() { + const gpu = useGpuInfo(); const { modelType, selectedModel, @@ -93,6 +100,24 @@ export function ModelSelectionStep() { const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]); + // Match Studio behavior: only show exception signals (OOM/TIGHT) in training flows. + const vramMap = useMemo(() => { + const fitMap = buildModelVramMap( + hfResults, + trainingMethod as VramTrainingMethod, + gpu, + ); + const map = new Map(); + for (const r of hfResults) { + const fit = fitMap.get(r.id); + map.set(r.id, { + status: fit?.status ?? null, + detail: r.totalParams ? formatCompact(r.totalParams) : null, + }); + } + return map; + }, [hfResults, gpu, trainingMethod]); + const comboboxAnchorRef = useRef(null); const { scrollRef, sentinelRef } = useInfiniteScroll( fetchMore, @@ -218,19 +243,21 @@ export function ModelSelectionStep() { > {(id: string) => { - const r = hfResults.find((r) => r.id === id); - const sizeLabel = r?.totalParams - ? formatCompact(r.totalParams) - : null; + const entry = vramMap.get(id); + const sizeLabel = entry?.detail ?? null; + const fitStatus = entry?.status ?? null; + const exceeds = fitStatus === "exceeds"; return ( - + {id} @@ -241,11 +268,23 @@ export function ModelSelectionStep() { {id} - {sizeLabel ? ( - - {sizeLabel} - - ) : null} + + {fitStatus === "exceeds" && ( + + OOM + + )} + {fitStatus === "tight" && ( + + TIGHT + + )} + {sizeLabel ? ( + + {sizeLabel} + + ) : null} + ); }} diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 7a15cea02d..b2e595a4c3 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -37,8 +37,7 @@ import { formatCompact } from "@/lib/utils"; import { type TrainingMethod as VramTrainingMethod, type VramFitStatus, - checkVramFit, - estimateLoadingVram, + buildModelVramMap, } from "@/lib/vram"; import { listLocalModels, @@ -218,22 +217,23 @@ export function ModelSection() { // 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). const vramMap = useMemo(() => { - const method = trainingMethod as VramTrainingMethod; + const fitMap = buildModelVramMap( + hfResults, + trainingMethod as VramTrainingMethod, + gpu, + ); const map = new Map< string, { est: number; status: VramFitStatus | null; detail: string | null } >(); for (const r of hfResults) { const detail = r.totalParams ? formatCompact(r.totalParams) : null; - if (r.totalParams) { - const est = estimateLoadingVram(r.totalParams, method); - const status = gpu.available - ? checkVramFit(est, gpu.memoryTotalGb) - : null; - map.set(r.id, { est, status, detail }); - } else { - map.set(r.id, { est: 0, status: null, detail }); - } + const fit = fitMap.get(r.id); + map.set(r.id, { + est: fit?.est ?? 0, + status: fit?.status ?? null, + detail, + }); } return map; }, [hfResults, gpu, trainingMethod]); diff --git a/studio/frontend/src/lib/vram.ts b/studio/frontend/src/lib/vram.ts index 7aea44e850..8152104ee2 100644 --- a/studio/frontend/src/lib/vram.ts +++ b/studio/frontend/src/lib/vram.ts @@ -93,3 +93,32 @@ export function checkVramFit( 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 { + const map = new Map(); + for (const model of models) { + if (!model.totalParams) { + map.set(model.id, { est: 0, status: null }); + continue; + } + + const est = estimateLoadingVram(model.totalParams, method); + const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null; + map.set(model.id, { est, status }); + } + return map; +}