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