|
| 1 | +package config |
| 2 | + |
| 3 | +import ( |
| 4 | + gguf "github.com/gpustack/gguf-parser-go" |
| 5 | + "github.com/mudler/LocalAI/pkg/xsysinfo" |
| 6 | + "github.com/mudler/xlog" |
| 7 | +) |
| 8 | + |
| 9 | +// contextFitHeadroomDivisor reserves a slice of per-device VRAM as headroom when |
| 10 | +// deciding whether an auto-derived context fits. The gguf-parser footprint |
| 11 | +// already covers weights + KV + compute buffer, but a live load also pays for |
| 12 | +// allocator fragmentation, the CUDA/HIP context, and whatever else shares the |
| 13 | +// card, so we require the estimate to leave at least 1/divisor of the device |
| 14 | +// free. /5 (~20% headroom) mirrors the SWA full-cache gate's margin. |
| 15 | +const contextFitHeadroomDivisor = 5 |
| 16 | + |
| 17 | +// contextFitCandidates is the descending set of context windows tried when the |
| 18 | +// DefaultAutoContextSize cap itself does not fit per-device VRAM. Only the rare |
| 19 | +// big-model-on-tiny-card case reaches this walk; it is capped at the base |
| 20 | +// choice and floored at DefaultContextSize, and returns the first (largest) |
| 21 | +// candidate that fits. |
| 22 | +var contextFitCandidates = []int{8192, 6144, 4096} |
| 23 | + |
| 24 | +// perDeviceVRAM reports the smallest per-GPU VRAM ceiling in bytes (0 = unknown |
| 25 | +// or no GPU). It is a package var so tests can inject a deterministic value — |
| 26 | +// detection does a live GPU probe. Per-device (not summed) is the right budget: |
| 27 | +// with all layers offloaded to a single device the whole footprint must fit that |
| 28 | +// one card, and a multi-GPU host is bounded by its smallest card. This mirrors |
| 29 | +// localGPU's use of MinPerGPUVRAM in hardware_defaults.go. |
| 30 | +var perDeviceVRAM = func() uint64 { |
| 31 | + v, _ := xsysinfo.MinPerGPUVRAM() |
| 32 | + return v |
| 33 | +} |
| 34 | + |
| 35 | +// estimateContextVRAM returns the estimated per-device VRAM footprint (bytes) of |
| 36 | +// running f fully offloaded at ctx tokens — weights + KV cache + compute buffer. |
| 37 | +// It returns 0 when it cannot produce an estimate (nil file, no tensors, or a |
| 38 | +// parser panic), which the caller treats as "cannot confirm a smaller fit" and |
| 39 | +// so keeps the conservative cap rather than clamping on a bogus number. It is a |
| 40 | +// package var so tests can stub it (a fabricated GGUF carries no tensors and |
| 41 | +// estimates to ~0). |
| 42 | +var estimateContextVRAM = func(f *gguf.GGUFFile, ctx int) (footprint uint64) { |
| 43 | + if f == nil { |
| 44 | + return 0 |
| 45 | + } |
| 46 | + if ctx <= 0 { |
| 47 | + ctx = DefaultContextSize |
| 48 | + } |
| 49 | + // The gguf-parser estimator panics on degenerate / partially-parsed GGUFs; |
| 50 | + // treat any failure as "unknown" so config loading never crashes on a model |
| 51 | + // the parser mis-handles. |
| 52 | + defer func() { |
| 53 | + if r := recover(); r != nil { |
| 54 | + xlog.Debug("[context_fit] per-device VRAM estimate failed; treating as unknown", "error", r) |
| 55 | + footprint = 0 |
| 56 | + } |
| 57 | + }() |
| 58 | + // Offload all layers (LocalAI's DefaultNGPULayers default; the estimator |
| 59 | + // clamps to the model's block count) so the estimate reflects a fully |
| 60 | + // GPU-resident model. NonUMA is the discrete-GPU figure (larger than the UMA |
| 61 | + // one), which keeps the fit check conservative on unified-memory hosts — they |
| 62 | + // have ample memory to clear it anyway. |
| 63 | + est := f.EstimateLLaMACppRun( |
| 64 | + gguf.WithLLaMACppContextSize(int32(ctx)), |
| 65 | + gguf.WithLLaMACppOffloadLayers(uint64(DefaultNGPULayers)), |
| 66 | + ) |
| 67 | + sum := est.Summarize(true, 0, 0) |
| 68 | + if len(sum.Items) == 0 { |
| 69 | + return 0 |
| 70 | + } |
| 71 | + var total uint64 |
| 72 | + for _, v := range sum.Items[0].VRAMs { |
| 73 | + total += uint64(v.NonUMA) |
| 74 | + } |
| 75 | + return total |
| 76 | +} |
| 77 | + |
| 78 | +// contextFitsVRAM reports whether an estimated footprint fits a per-device VRAM |
| 79 | +// ceiling with headroom (VRAM must exceed the footprint by ~1/divisor). Unknown |
| 80 | +// inputs (0) are treated as "cannot confirm" so a detection or estimate gap does |
| 81 | +// not clamp the context. |
| 82 | +func contextFitsVRAM(footprint, vram uint64) bool { |
| 83 | + if footprint == 0 || vram == 0 { |
| 84 | + return false |
| 85 | + } |
| 86 | + return vram >= footprint+footprint/contextFitHeadroomDivisor |
| 87 | +} |
| 88 | + |
| 89 | +// autoContextSize picks the default context to use for f when the user did not |
| 90 | +// set context_size. The choice is deliberately conservative, NOT |
| 91 | +// VRAM-maximizing: |
| 92 | +// |
| 93 | +// 1. Base cap: min(trainedMax, DefaultAutoContextSize). A small model keeps its |
| 94 | +// trained window; a long-context model (128k / 256k / 1M) is capped so its |
| 95 | +// KV cache does not default to a size no consumer GPU can hold. This applies |
| 96 | +// always, including CPU / unknown-VRAM hosts. |
| 97 | +// 2. VRAM is only a downward safety: when a per-device VRAM ceiling IS detected |
| 98 | +// and even the base cap would not fit it (with headroom), step down through |
| 99 | +// contextFitCandidates to the largest window that fits, floored at |
| 100 | +// DefaultContextSize. When VRAM is unknown we skip this — the base cap is |
| 101 | +// already safe and we must not regress CPU / detection-gap hosts. |
| 102 | +// |
| 103 | +// trainedMax <= 0 means the estimate yielded nothing usable; the caller keeps |
| 104 | +// its existing DefaultContextSize fallback in that case, so this is only called |
| 105 | +// with a positive trainedMax. |
| 106 | +func autoContextSize(f *gguf.GGUFFile, trainedMax int) int { |
| 107 | + chosen := trainedMax |
| 108 | + if chosen > DefaultAutoContextSize { |
| 109 | + chosen = DefaultAutoContextSize |
| 110 | + } |
| 111 | + |
| 112 | + vram := perDeviceVRAM() |
| 113 | + if vram == 0 { |
| 114 | + // No per-device VRAM detected (CPU-only, unified memory reporting nothing, |
| 115 | + // or a detection gap). The bug is GPU OOM-on-load, so with no GPU budget to |
| 116 | + // reason about we must not clamp — the base cap already bounds long-context |
| 117 | + // models. |
| 118 | + return chosen |
| 119 | + } |
| 120 | + |
| 121 | + if contextFitsVRAM(estimateContextVRAM(f, chosen), vram) { |
| 122 | + return chosen |
| 123 | + } |
| 124 | + |
| 125 | + // The base cap does not fit this card. Walk candidates downward and take the |
| 126 | + // largest that fits, never below DefaultContextSize. |
| 127 | + for _, cand := range contextFitCandidates { |
| 128 | + if cand > chosen || cand < DefaultContextSize { |
| 129 | + continue |
| 130 | + } |
| 131 | + if contextFitsVRAM(estimateContextVRAM(f, cand), vram) { |
| 132 | + xlog.Debug("[context_fit] capped auto context to fit per-device VRAM", |
| 133 | + "context", cand, "base_cap", chosen, "vram_gib", vram>>30) |
| 134 | + return cand |
| 135 | + } |
| 136 | + } |
| 137 | + |
| 138 | + // Nothing fit (an unusually large model on a tiny card): fall back to the |
| 139 | + // floor. The backend still clamps n_gpu_layers to what fits, so a partial |
| 140 | + // offload can keep the model loadable rather than aborting outright. |
| 141 | + xlog.Debug("[context_fit] no candidate context fit per-device VRAM; using floor", |
| 142 | + "context", DefaultContextSize, "base_cap", chosen, "vram_gib", vram>>30) |
| 143 | + return DefaultContextSize |
| 144 | +} |
0 commit comments