Compare commits

..

1 Commits

Author SHA1 Message Date
Ettore Di Giacinto
fa15af0748 fix(config): cap auto-derived context to fit VRAM
When a model is imported without an explicit context_size, the GGUF
importer defaulted the model's context to its full trained window
(n_ctx_train). For long-context models (128k / 256k / 1M) that KV cache
cannot fit a consumer GPU, so the backend aborts on load (exitCode=-1)
even though the model file is perfectly fine. Reproduced live:
gemma-4-26b-a4b-it-qat-q4_0 defaulted to context=262144 and
qwythos-9b-claude-mythos-5-1m to 1048576, both aborting on a 20 GB card.

Instead of chasing the trained max, auto-derive a conservative default:
min(trainedMax, DefaultAutoContextSize=8192). A small model keeps its
trained window; a long-context model caps at 8k and users opt into more
via context_size. This cap applies always, including CPU / unknown-VRAM
hosts, so it never regresses those paths.

Per-device VRAM is used only as a DOWNWARD safety: when a per-device
ceiling is detected (xsysinfo.MinPerGPUVRAM) and even the 8k cap would
not fit it with headroom, step down through candidate contexts to the
largest that fits, floored at DefaultContextSize. When VRAM is unknown
(0) or no GPU is detected we do NOT clamp — the bug is GPU OOM and the
8k cap is already safe, so detection gaps must not shrink the window.

The footprint estimate reuses gpustack/gguf-parser-go's
EstimateLLaMACppRun at a given context with all layers offloaded, taking
the per-device NonUMA VRAM figure. The estimate and VRAM detection are
package vars so tests inject deterministic values. Explicit context_size
always wins (guessGGUFFromFile only acts when it is nil).

Assisted-by: Claude:claude-opus-4-8 [golangci-lint go-test]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
2026-07-06 07:52:46 +00:00
4 changed files with 265 additions and 3 deletions

144
core/config/context_fit.go Normal file
View File

@@ -0,0 +1,144 @@
package config
import (
gguf "github.com/gpustack/gguf-parser-go"
"github.com/mudler/LocalAI/pkg/xsysinfo"
"github.com/mudler/xlog"
)
// contextFitHeadroomDivisor reserves a slice of per-device VRAM as headroom when
// deciding whether an auto-derived context fits. The gguf-parser footprint
// already covers weights + KV + compute buffer, but a live load also pays for
// allocator fragmentation, the CUDA/HIP context, and whatever else shares the
// card, so we require the estimate to leave at least 1/divisor of the device
// free. /5 (~20% headroom) mirrors the SWA full-cache gate's margin.
const contextFitHeadroomDivisor = 5
// contextFitCandidates is the descending set of context windows tried when the
// DefaultAutoContextSize cap itself does not fit per-device VRAM. Only the rare
// big-model-on-tiny-card case reaches this walk; it is capped at the base
// choice and floored at DefaultContextSize, and returns the first (largest)
// candidate that fits.
var contextFitCandidates = []int{8192, 6144, 4096}
// perDeviceVRAM reports the smallest per-GPU VRAM ceiling in bytes (0 = unknown
// or no GPU). It is a package var so tests can inject a deterministic value —
// detection does a live GPU probe. Per-device (not summed) is the right budget:
// with all layers offloaded to a single device the whole footprint must fit that
// one card, and a multi-GPU host is bounded by its smallest card. This mirrors
// localGPU's use of MinPerGPUVRAM in hardware_defaults.go.
var perDeviceVRAM = func() uint64 {
v, _ := xsysinfo.MinPerGPUVRAM()
return v
}
// estimateContextVRAM returns the estimated per-device VRAM footprint (bytes) of
// running f fully offloaded at ctx tokens — weights + KV cache + compute buffer.
// It returns 0 when it cannot produce an estimate (nil file, no tensors, or a
// parser panic), which the caller treats as "cannot confirm a smaller fit" and
// so keeps the conservative cap rather than clamping on a bogus number. It is a
// package var so tests can stub it (a fabricated GGUF carries no tensors and
// estimates to ~0).
var estimateContextVRAM = func(f *gguf.GGUFFile, ctx int) (footprint uint64) {
if f == nil {
return 0
}
if ctx <= 0 {
ctx = DefaultContextSize
}
// The gguf-parser estimator panics on degenerate / partially-parsed GGUFs;
// treat any failure as "unknown" so config loading never crashes on a model
// the parser mis-handles.
defer func() {
if r := recover(); r != nil {
xlog.Debug("[context_fit] per-device VRAM estimate failed; treating as unknown", "error", r)
footprint = 0
}
}()
// Offload all layers (LocalAI's DefaultNGPULayers default; the estimator
// clamps to the model's block count) so the estimate reflects a fully
// GPU-resident model. NonUMA is the discrete-GPU figure (larger than the UMA
// one), which keeps the fit check conservative on unified-memory hosts — they
// have ample memory to clear it anyway.
est := f.EstimateLLaMACppRun(
gguf.WithLLaMACppContextSize(int32(ctx)),
gguf.WithLLaMACppOffloadLayers(uint64(DefaultNGPULayers)),
)
sum := est.Summarize(true, 0, 0)
if len(sum.Items) == 0 {
return 0
}
var total uint64
for _, v := range sum.Items[0].VRAMs {
total += uint64(v.NonUMA)
}
return total
}
// contextFitsVRAM reports whether an estimated footprint fits a per-device VRAM
// ceiling with headroom (VRAM must exceed the footprint by ~1/divisor). Unknown
// inputs (0) are treated as "cannot confirm" so a detection or estimate gap does
// not clamp the context.
func contextFitsVRAM(footprint, vram uint64) bool {
if footprint == 0 || vram == 0 {
return false
}
return vram >= footprint+footprint/contextFitHeadroomDivisor
}
// autoContextSize picks the default context to use for f when the user did not
// set context_size. The choice is deliberately conservative, NOT
// VRAM-maximizing:
//
// 1. Base cap: min(trainedMax, DefaultAutoContextSize). A small model keeps its
// trained window; a long-context model (128k / 256k / 1M) is capped so its
// KV cache does not default to a size no consumer GPU can hold. This applies
// always, including CPU / unknown-VRAM hosts.
// 2. VRAM is only a downward safety: when a per-device VRAM ceiling IS detected
// and even the base cap would not fit it (with headroom), step down through
// contextFitCandidates to the largest window that fits, floored at
// DefaultContextSize. When VRAM is unknown we skip this — the base cap is
// already safe and we must not regress CPU / detection-gap hosts.
//
// trainedMax <= 0 means the estimate yielded nothing usable; the caller keeps
// its existing DefaultContextSize fallback in that case, so this is only called
// with a positive trainedMax.
func autoContextSize(f *gguf.GGUFFile, trainedMax int) int {
chosen := trainedMax
if chosen > DefaultAutoContextSize {
chosen = DefaultAutoContextSize
}
vram := perDeviceVRAM()
if vram == 0 {
// No per-device VRAM detected (CPU-only, unified memory reporting nothing,
// or a detection gap). The bug is GPU OOM-on-load, so with no GPU budget to
// reason about we must not clamp — the base cap already bounds long-context
// models.
return chosen
}
if contextFitsVRAM(estimateContextVRAM(f, chosen), vram) {
return chosen
}
// The base cap does not fit this card. Walk candidates downward and take the
// largest that fits, never below DefaultContextSize.
for _, cand := range contextFitCandidates {
if cand > chosen || cand < DefaultContextSize {
continue
}
if contextFitsVRAM(estimateContextVRAM(f, cand), vram) {
xlog.Debug("[context_fit] capped auto context to fit per-device VRAM",
"context", cand, "base_cap", chosen, "vram_gib", vram>>30)
return cand
}
}
// Nothing fit (an unusually large model on a tiny card): fall back to the
// floor. The backend still clamps n_gpu_layers to what fits, so a partial
// offload can keep the model loadable rather than aborting outright.
xlog.Debug("[context_fit] no candidate context fit per-device VRAM; using floor",
"context", DefaultContextSize, "base_cap", chosen, "vram_gib", vram>>30)
return DefaultContextSize
}

View File

@@ -0,0 +1,101 @@
package config
import (
gguf "github.com/gpustack/gguf-parser-go"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
// These specs exercise the auto-derived default context. The detection seams
// (perDeviceVRAM, estimateContextVRAM) are package vars so a deterministic VRAM
// ceiling and footprint can be injected without a real GPU or model file — the
// same pattern hardware_defaults_internal_test.go uses for localGPU.
var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
const gib = uint64(1) << 30
var (
origVRAM func() uint64
origEstimate func(f *gguf.GGUFFile, ctx int) uint64
)
BeforeEach(func() {
origVRAM = perDeviceVRAM
origEstimate = estimateContextVRAM
})
AfterEach(func() {
perDeviceVRAM = origVRAM
estimateContextVRAM = origEstimate
})
Context("autoContextSize", func() {
It("caps a long-context model at DefaultAutoContextSize when VRAM is ample", func() {
// 1M-context model on an 80 GiB card: we do NOT chase the trained max,
// we keep the conservative 8k cap (users opt into more via context_size).
perDeviceVRAM = func() uint64 { return 80 * gib }
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return gib } // trivially fits
Expect(autoContextSize(nil, 1048576)).To(Equal(DefaultAutoContextSize))
})
It("keeps a small model's trained window instead of inflating it", func() {
// trained 4096 < 8192: min() keeps 4096, it is not raised to the cap.
perDeviceVRAM = func() uint64 { return 80 * gib }
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return gib }
Expect(autoContextSize(nil, 4096)).To(Equal(4096))
})
It("steps below the cap when even 8k would not fit a tiny card", func() {
// A large model on a 2 GiB card where the 8k footprint overflows but a
// smaller context fits: choose the largest that fits, never below the
// floor. Footprint grows with context so the walk finds a fit.
perDeviceVRAM = func() uint64 { return 2 * gib }
estimateContextVRAM = func(_ *gguf.GGUFFile, ctx int) uint64 {
return gib + uint64(ctx)*100000
}
chosen := autoContextSize(nil, 1048576)
Expect(chosen).To(BeNumerically("<", DefaultAutoContextSize))
Expect(chosen).To(BeNumerically(">=", DefaultContextSize))
// The chosen context's footprint must actually fit the card with headroom.
Expect(contextFitsVRAM(estimateContextVRAM(nil, chosen), 2*gib)).To(BeTrue())
})
It("falls back to the floor when nothing fits", func() {
// Even DefaultContextSize does not fit: return the floor and let the
// backend clamp n_gpu_layers to what it can (partial offload) rather
// than defaulting to a window guaranteed to abort.
perDeviceVRAM = func() uint64 { return 1 * gib }
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return 100 * gib }
Expect(autoContextSize(nil, 1048576)).To(Equal(DefaultContextSize))
})
It("does not clamp when per-device VRAM is unknown", func() {
// CPU-only / detection gap: no GPU budget to reason about, so we must
// not regress — keep the conservative base cap regardless of estimate.
perDeviceVRAM = func() uint64 { return 0 }
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return 999 * gib }
Expect(autoContextSize(nil, 1048576)).To(Equal(DefaultAutoContextSize))
})
})
Context("guessGGUFFromFile", func() {
It("never overrides an explicitly configured context_size", func() {
// A fabricated GGUF is enough: the context branch is skipped entirely
// when the user pinned context_size, so the estimate is never consulted.
explicit := 262144
cfg := &ModelConfig{LLMConfig: LLMConfig{ContextSize: &explicit}}
f := &gguf.GGUFFile{
Header: gguf.GGUFHeader{
MetadataKV: gguf.GGUFMetadataKVs{
{
Key: "general.architecture",
ValueType: gguf.GGUFMetadataValueTypeString,
Value: "llama",
},
},
},
}
guessGGUFFromFile(cfg, f, 0)
Expect(cfg.ContextSize).ToNot(BeNil())
Expect(*cfg.ContextSize).To(Equal(262144))
})
})
})

View File

@@ -18,6 +18,18 @@ const (
// safe default beats a tiny, surprising window that truncates real prompts.
DefaultContextSize = 4096
// DefaultAutoContextSize caps the context we auto-derive from a GGUF when the
// user did not set context_size. The GGUF importer used to default a model's
// context to its full trained window (n_ctx_train). For long-context models
// (128k / 256k / 1M) that KV cache cannot fit a consumer GPU and the backend
// aborts on load (exitCode=-1) even though the model file is fine. So instead
// of shooting for the trained max, we keep a modest default: a small model
// (trained < this) keeps its trained window, while a long-context model caps
// here. Users who want the full window raise context_size explicitly. This is
// a conservative default, not a VRAM-maximizing one — VRAM is only used to
// step further DOWN when even this cap would not fit (see context_fit.go).
DefaultAutoContextSize = 8192
// DefaultNGPULayers means "offload all layers"; the backend (fit_params)
// clamps to what actually fits in device memory.
DefaultNGPULayers = 99999999

View File

@@ -28,9 +28,14 @@ func reservedNonChatModel(cfg *ModelConfig) bool {
func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
if defaultCtx == 0 && cfg.ContextSize == nil {
ctxSize := f.EstimateLLaMACppRun().ContextSize
if ctxSize > 0 {
cSize := int(ctxSize)
// trainedMax is the model's full trained context window (n_ctx_train).
// Defaulting a model to it unbounded is what OOMs long-context models at
// load: a 128k / 256k / 1M KV cache cannot fit a consumer GPU and the
// backend aborts (exitCode=-1). autoContextSize instead caps to a modest
// default and only steps below it when detected per-device VRAM demands.
trainedMax := int(f.EstimateLLaMACppRun().ContextSize)
if trainedMax > 0 {
cSize := autoContextSize(f, trainedMax)
cfg.ContextSize = &cSize
} else {
defaultCtx = DefaultContextSize