ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Aggregating sharded models, showing fit/oom for quantizations

Browse source
This commit is contained in:
Manan17 2026-02-27 08:23:15 +00:00
commit 168957a87a
10 changed files with 105 additions and 48 deletions
Download patch file Download diff file Expand all files Collapse all files

4
studio/backend/core/inference/llama_cpp.py
View file

@ -62,6 +62,10 @@ class LlamaCppBackend:
def is_vision(self) -> bool:
return self._is_vision
@property
def hf_variant(self) -> Optional[str]:
return self._hf_variant
# ── Binary discovery ──────────────────────────────────────────
@staticmethod

1
studio/backend/models/inference.py
View file

@ -59,6 +59,7 @@ class InferenceStatusResponse(BaseModel):
active_model: Optional[str] = Field(None, description="Currently active model identifier")
is_vision: bool = Field(False, description="Whether the active model is a vision model")
is_gguf: bool = Field(False, description="Whether the active model is a GGUF model (llama.cpp)")
gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. Q4_K_M)")
loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
loaded: List[str] = Field(default_factory=list, description="Models currently loaded")

1
studio/backend/routes/inference.py
View file

@ -309,6 +309,7 @@ async def get_status():
active_model=llama_backend.model_identifier,
is_vision=llama_backend.is_vision,
is_gguf=True,
gguf_variant=llama_backend.hf_variant,
loading=[],
loaded=[llama_backend.model_identifier],
)

41
studio/backend/utils/models/model_config.py
View file

@ -491,21 +491,28 @@ def _extract_quant_label(filename: str) -> str:
Extract quantization label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename.
Examples:
"gemma-3-4b-it-Q4_K_M.gguf" → "Q4_K_M"
"model-IQ4_NL.gguf" → "IQ4_NL"
"model-BF16.gguf" → "BF16"
"model-UD-IQ1_S.gguf" → "UD-IQ1_S"
"gemma-3-4b-it-Q4_K_M.gguf" → "Q4_K_M"
"model-IQ4_NL.gguf" → "IQ4_NL"
"model-BF16.gguf" → "BF16"
"model-UD-IQ1_S.gguf" → "UD-IQ1_S"
"model-UD-TQ1_0.gguf" → "UD-TQ1_0"
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE"
"""
import re
stem = filename.rsplit(".", 1)[0] # Remove .gguf
# Match known quantization patterns (UD- prefix, IQ, Q, BF/F variants)
# Use only the basename (rfilename may include directory)
basename = filename.rsplit("/", 1)[-1]
# Strip .gguf and any shard suffix (-00001-of-00010)
stem = re.sub(r'-\d{3,}-of-\d{3,}', '', basename.rsplit(".", 1)[0])
# Match known quantization patterns
match = re.search(
r'(UD-)?' # Optional UD- prefix (Ultra Discrete)
r'(IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?' # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r'|Q[0-9]+_K_[A-Z]+' # K-quant: Q4_K_M, Q3_K_S
r'|Q[0-9]+_[0-9]+' # Standard: Q8_0, Q5_1
r'|Q[0-9]+_K' # Short K-quant: Q6_K
r'|BF16|F16|F32)', # Full precision
r'(MXFP[0-9]+(?:_[A-Z0-9]+)*' # MXFP variants: MXFP4, MXFP4_MOE
r'|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?' # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r'|TQ[0-9]+_[0-9]+' # Ternary quant: TQ1_0, TQ2_0
r'|Q[0-9]+_K_[A-Z]+' # K-quant: Q4_K_M, Q3_K_S
r'|Q[0-9]+_[0-9]+' # Standard: Q8_0, Q5_1
r'|Q[0-9]+_K' # Short K-quant: Q6_K
r'|BF16|F16|F32)', # Full precision
stem, re.IGNORECASE,
)
if match:
@ -534,6 +541,9 @@ def list_gguf_variants(
variants: list[GgufVariantInfo] = []
has_vision = False
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_first_file: dict[str, str] = {} # quant -> first filename (for display)
for sibling in info.siblings:
fname = sibling.rfilename
if not fname.endswith(".gguf"):
@ -546,10 +556,15 @@ def list_gguf_variants(
continue
quant = _extract_quant_label(fname)
quant_totals[quant] = quant_totals.get(quant, 0) + size
if quant not in quant_first_file:
quant_first_file[quant] = fname
for quant, total_size in quant_totals.items():
variants.append(GgufVariantInfo(
filename=fname,
filename=quant_first_file[quant],
quant=quant,
size_bytes=size,
size_bytes=total_size,
))
return variants, has_vision

14
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -27,6 +27,7 @@ interface ModelSelectorProps {
loraModels?: LoraModelOption[];
value?: string;
defaultValue?: string;
activeGgufVariant?: string | null;
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
onEject?: () => void;
variant?: "outline" | "ghost" | "muted";
@ -158,6 +159,7 @@ export function ModelSelector({
loraModels = [],
value,
defaultValue,
activeGgufVariant,
onValueChange,
onEject,
variant = "outline",
@ -202,9 +204,15 @@ export function ModelSelector({
return all;
}, [loraModels, models]);
const currentModel = selected
? optionById.get(selected) ?? { id: selected, name: selected }
: undefined;
const currentModel = useMemo(() => {
if (!selected) return undefined;
const found = optionById.get(selected);
if (activeGgufVariant) {
const desc = `GGUF · ${activeGgufVariant}`;
return found ? { ...found, description: desc } : { id: selected, name: selected, description: desc };
}
return found ?? { id: selected, name: selected };
}, [selected, optionById, activeGgufVariant]);
function handleSelect(id: string, meta: ModelSelectorChangeMeta) {
if (onValueChange) {

65
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

@ -130,9 +130,11 @@ function ModelRow({
function GgufVariantExpander({
repoId,
onSelect,
gpuGb,
}: {
repoId: string;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
gpuGb?: number;
}) {
const [variants, setVariants] = useState<GgufVariantDetail[] | null>(null);
const [defaultVariant, setDefaultVariant] = useState<string | null>(null);
@ -209,28 +211,45 @@ function GgufVariantExpander({
<span className="text-[9px] font-medium text-blue-400">Vision</span>
)}
</div>
{variants.map((v) => (
<button
key={v.filename}
type="button"
onClick={() => handleVariantClick(v.quant)}
className={cn(
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1 text-left text-sm transition-colors hover:bg-accent",
)}
>
<span className="min-w-0 flex-1 truncate font-mono text-xs">
{v.quant}
{v.quant === defaultVariant && (
<span className="ml-1.5 text-[9px] font-sans font-medium text-primary/70">
recommended
</span>
{variants.map((v) => {
const sizeGb = v.size_bytes / (1024 ** 3);
const fitStatus = gpuGb != null && gpuGb > 0 && sizeGb > 0
? checkVramFit(sizeGb, gpuGb)
: null;
return (
<button
key={v.filename}
type="button"
onClick={() => handleVariantClick(v.quant)}
className={cn(
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1 text-left text-sm transition-colors hover:bg-accent",
)}
</span>
<span className="text-[10px] text-muted-foreground shrink-0">
{formatBytes(v.size_bytes)}
</span>
</button>
))}
>
<span className="min-w-0 flex-1 truncate font-mono text-xs">
{v.quant}
{v.quant === defaultVariant && (
<span className="ml-1.5 text-[9px] font-sans font-medium text-primary/70">
recommended
</span>
)}
</span>
<span className="flex items-center gap-1.5 shrink-0">
{fitStatus === "exceeds" && (
<span className="text-[9px] font-medium text-red-400">OOM</span>
)}
{fitStatus === "tight" && (
<span className="text-[9px] font-medium text-amber-400">TIGHT</span>
)}
{fitStatus === "fits" && (
<span className="text-[9px] font-medium text-emerald-500/90">FIT</span>
)}
<span className="text-[10px] text-muted-foreground">
{formatBytes(v.size_bytes)}
</span>
</span>
</button>
);
})}
</div>
);
}
@ -390,7 +409,7 @@ export function HubModelPicker({
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
{expandedGguf === id && (
<GgufVariantExpander repoId={id} onSelect={onSelect} />
<GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} />
)}
</div>
);
@ -425,7 +444,7 @@ export function HubModelPicker({
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
{expandedGguf === id && (
<GgufVariantExpander repoId={id} onSelect={onSelect} />
<GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} />
)}
</div>
);

9
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -314,6 +314,7 @@ export function ChatPage(): ReactElement {
);
const inferenceParams = useChatRuntimeStore((state) => state.params);
const setInferenceParams = useChatRuntimeStore((state) => state.setParams);
const activeGgufVariant = useChatRuntimeStore((state) => state.activeGgufVariant);
const autoTitle = useChatRuntimeStore((state) => state.autoTitle);
const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle);
const modelsFromStore = useChatRuntimeStore((state) => state.models);
@ -336,9 +337,10 @@ export function ChatPage(): ReactElement {
const handleCheckpointChange = useCallback(
(value: string, meta?: { isLora: boolean; ggufVariant?: string }) => {
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
if (!value || value === currentCheckpoint) return;
const store = useChatRuntimeStore.getState();
const currentCheckpoint = store.params.checkpoint;
const currentVariant = store.activeGgufVariant;
if (!value || (value === currentCheckpoint && (meta?.ggufVariant ?? null) === (currentVariant ?? null))) return;
void (async () => {
let switchNote: string | undefined;
const activeThreadId = await resolveActiveSingleThreadId(view);
@ -591,6 +593,7 @@ export function ChatPage(): ReactElement {
models={models}
loraModels={loraModels}
value={inferenceParams.checkpoint}
activeGgufVariant={activeGgufVariant}
onValueChange={handleCheckpointChange}
onEject={handleEject}
variant="ghost"

9
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -144,7 +144,7 @@ export function useChatModelRuntime() {
setLoras(lorasRes.loras.map(toLoraSummary));
if (statusRes.active_model) {
setCheckpoint(statusRes.active_model);
setCheckpoint(statusRes.active_model, statusRes.gguf_variant);
}
} catch (error) {
const message =
@ -159,14 +159,15 @@ export function useChatModelRuntime() {
const selectModel = useCallback(
async (selection: string | SelectedModelInput) => {
const modelId = typeof selection === "string" ? selection : selection.id;
if (!modelId || params.checkpoint === modelId) {
const ggufVariant =
typeof selection === "string" ? undefined : selection.ggufVariant;
const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
if (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null))) {
return;
}
const explicitIsLora =
typeof selection === "string" ? undefined : selection.isLora;
const ggufVariant =
typeof selection === "string" ? undefined : selection.ggufVariant;
const extraLoadingDescription =
typeof selection === "string" ? undefined : selection.loadingDescription;
const model = models.find((entry) => entry.id === modelId);

8
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -39,13 +39,14 @@ type ChatRuntimeStore = {
runningByThreadId: Record<string, boolean>;
autoTitle: boolean;
modelsError: string | null;
activeGgufVariant: string | null;
setParams: (params: InferenceParams) => void;
setModels: (models: ChatModelSummary[]) => void;
setLoras: (loras: ChatLoraSummary[]) => void;
setThreadRunning: (threadId: string, running: boolean) => void;
setAutoTitle: (enabled: boolean) => void;
setModelsError: (error: string | null) => void;
setCheckpoint: (modelId: string) => void;
setCheckpoint: (modelId: string, ggufVariant?: string | null) => void;
clearCheckpoint: () => void;
};
@ -56,6 +57,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
runningByThreadId: {},
autoTitle: loadBool(AUTO_TITLE_KEY, false),
modelsError: null,
activeGgufVariant: null,
setParams: (params) => set({ params }),
setModels: (models) => set({ models }),
setLoras: (loras) => set({ loras }),
@ -75,12 +77,13 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
return { autoTitle };
}),
setModelsError: (modelsError) => set({ modelsError }),
setCheckpoint: (modelId) =>
setCheckpoint: (modelId, ggufVariant) =>
set((state) => ({
params: {
...state.params,
checkpoint: modelId,
},
activeGgufVariant: ggufVariant ?? null,
})),
clearCheckpoint: () =>
set((state) => ({
@ -88,5 +91,6 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
...state.params,
checkpoint: "",
},
activeGgufVariant: null,
})),
}));

1
studio/frontend/src/features/chat/types/api.ts
View file

@ -69,6 +69,7 @@ export interface InferenceStatusResponse {
active_model: string | null;
is_vision: boolean;
is_gguf?: boolean;
gguf_variant?: string | null;
loading: string[];
loaded: string[];
}