diff --git a/core/config/context_fit.go b/core/config/context_fit.go new file mode 100644 index 000000000..7ee98540e --- /dev/null +++ b/core/config/context_fit.go @@ -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 +} diff --git a/core/config/context_fit_internal_test.go b/core/config/context_fit_internal_test.go new file mode 100644 index 000000000..0e0f9fa8a --- /dev/null +++ b/core/config/context_fit_internal_test.go @@ -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)) + }) + }) +}) diff --git a/core/config/defaults.go b/core/config/defaults.go index bb993d075..98045ed48 100644 --- a/core/config/defaults.go +++ b/core/config/defaults.go @@ -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 diff --git a/core/config/gguf.go b/core/config/gguf.go index 177e68749..6e6c796e1 100644 --- a/core/config/gguf.go +++ b/core/config/gguf.go @@ -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