mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-05 22:09:02 -04:00
fix(llama-cpp): cap single-pass embedding batch to fit VRAM
Embedding/score/rerank all decode or pool the whole input in one physical batch, so EffectiveBatchSize sized the batch to the full context window. For a large context that makes n_ubatch huge, and the per-device CUDA compute buffer (forward-graph scratch, ~n_ubatch * n_ctx, NOT split across GPUs) balloons into multi-GiB: a large-context embedding model then aborts on load (exitCode=-1) even with plenty of free VRAM. Reproduced with qwen3-embedding-4b (context 40960 -> n_batch 40960 -> abort) and qwen3-embedding-0.6b (n_batch 8192); pinning batch:512 avoided it. This is the same root cause as issue #10485 (a large context turns the batch into multi-GiB of scratch that must fit on a SINGLE card), but the single-pass path bypassed the VRAM headroom guard the config layer already had — it returned the unbounded context as the batch with no GPU awareness. Make the single-pass batch VRAM-aware: cap it to the largest batch whose compute buffer fits the per-device VRAM headroom, clamped to [DefaultPhysicalBatch, ctx], reusing the existing computeBufferBytesPerCell and headroom-divisor math (no duplication). Unknown per-device VRAM (0) stays conservative (DefaultPhysicalBatch, not the context) so a detection gap can't OOM. The GPU is resolved through an injectable package var (config.LocalGPU, backed by sync.Once-cached xsysinfo detection) so the per-request router call stays cheap and tests inject a deterministic device. Explicit batch: still wins. An input longer than the cap can no longer be pooled in one pass — the accepted tradeoff, since a batch that OOMs the device processes nothing. Assisted-by: Claude:claude-opus-4-8 golangci-lint go-test
This commit is contained in:
@@ -223,13 +223,24 @@ func EffectiveContextSize(c config.ModelConfig) int {
|
||||
return DefaultContextSize
|
||||
}
|
||||
|
||||
// localGPU resolves the device that will run the model, for single-pass batch
|
||||
// sizing. It is a package var so tests inject a deterministic device; production
|
||||
// reads config.LocalGPU, whose detection is sync.Once-cached in xsysinfo — so the
|
||||
// per-request call from the router's prompt trimmer (modelTokenTrim) stays cheap.
|
||||
var localGPU = config.LocalGPU
|
||||
|
||||
// EffectiveBatchSize is the single-decode batch the backend will run with.
|
||||
// Score, embedding and rerank all process the whole input in one pass: score
|
||||
// decodes prompt+candidate (asserts n_tokens <= n_batch), and embedding/rerank
|
||||
// pool over the full sequence in one physical batch (n_ubatch). So the batch
|
||||
// is sized to the context — anything that fits the context fits one pass,
|
||||
// pool over the full sequence in one physical batch (n_ubatch). Ideally the batch
|
||||
// covers the whole context so any input that fits the context fits one pass,
|
||||
// avoiding both the GGML_ASSERT crash and the "input is too large to process"
|
||||
// error. Explicit `batch:` always wins.
|
||||
// error — BUT a full ctx-sized n_ubatch makes the per-device CUDA compute buffer
|
||||
// multi-GiB (it scales ~ n_ubatch * n_ctx and can't be split across GPUs), so a
|
||||
// large-context embedding model aborts on load with free VRAM to spare (#10485).
|
||||
// So we cap the batch to the largest that fits the per-device VRAM headroom; an
|
||||
// input longer than that cap is the accepted tradeoff (it can't be pooled in one
|
||||
// pass, but the load no longer OOMs). Explicit `batch:` always wins.
|
||||
func EffectiveBatchSize(c config.ModelConfig) int {
|
||||
if c.Batch != 0 {
|
||||
return c.Batch
|
||||
@@ -238,7 +249,7 @@ func EffectiveBatchSize(c config.ModelConfig) int {
|
||||
c.HasUsecases(config.FLAG_EMBEDDINGS) ||
|
||||
c.HasUsecases(config.FLAG_RERANK)
|
||||
if ctx := EffectiveContextSize(c); singlePass && ctx > DefaultBatchSize {
|
||||
return ctx
|
||||
return config.SinglePassBatchForContext(localGPU(), ctx)
|
||||
}
|
||||
return DefaultBatchSize
|
||||
}
|
||||
|
||||
@@ -103,6 +103,19 @@ var _ = Describe("grpcModelOpts NBatch", func() {
|
||||
threads := 1
|
||||
ctx := 4096
|
||||
|
||||
// The single-pass batch is now VRAM-aware, so inject a deterministic GPU with
|
||||
// ample per-device VRAM: at these small contexts the compute buffer fits
|
||||
// easily, so EffectiveBatchSize returns the full context (the pre-#10485
|
||||
// behaviour these cases assert). Without injection the value would depend on
|
||||
// the CI host's real (often unknown) VRAM.
|
||||
const gib = uint64(1) << 30
|
||||
var origLocalGPU func() config.GPU
|
||||
BeforeEach(func() {
|
||||
origLocalGPU = localGPU
|
||||
localGPU = func() config.GPU { return config.GPU{VRAM: 119 * gib} }
|
||||
})
|
||||
AfterEach(func() { localGPU = origLocalGPU })
|
||||
|
||||
It("defaults to 512 for an ordinary model", func() {
|
||||
cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
|
||||
opts := grpcModelOpts(cfg, "/tmp/models")
|
||||
@@ -162,6 +175,58 @@ var _ = Describe("grpcModelOpts NBatch", func() {
|
||||
})
|
||||
})
|
||||
|
||||
// Guards the VRAM-aware cap on the single-pass (embedding/score/rerank) batch:
|
||||
// a large context must not turn n_ubatch into a multi-GiB compute buffer that
|
||||
// aborts the load on a device with free VRAM (issue #10485). The GPU is injected
|
||||
// via the localGPU package var so the cap is deterministic without a real device.
|
||||
var _ = Describe("EffectiveBatchSize VRAM cap", func() {
|
||||
const gib = uint64(1) << 30
|
||||
embeddings := config.FLAG_EMBEDDINGS
|
||||
threads := 1
|
||||
|
||||
var origLocalGPU func() config.GPU
|
||||
BeforeEach(func() { origLocalGPU = localGPU })
|
||||
AfterEach(func() { localGPU = origLocalGPU })
|
||||
|
||||
singlePassCfg := func(ctx int) config.ModelConfig {
|
||||
return config.ModelConfig{
|
||||
Threads: &threads,
|
||||
LLMConfig: config.LLMConfig{ContextSize: &ctx},
|
||||
KnownUsecases: &embeddings,
|
||||
}
|
||||
}
|
||||
|
||||
It("caps a large embedding context to a batch below the context but at least the default", func() {
|
||||
// Reproduces qwen3-embedding-4b: context 40960 on a modest 20 GiB card.
|
||||
// Full-context n_ubatch=40960 aborts; the cap must fit the VRAM headroom.
|
||||
localGPU = func() config.GPU { return config.GPU{VRAM: 20 * gib} }
|
||||
batch := EffectiveBatchSize(singlePassCfg(40960))
|
||||
Expect(batch).To(BeNumerically(">=", DefaultBatchSize))
|
||||
Expect(batch).To(BeNumerically("<", 40960))
|
||||
})
|
||||
|
||||
It("keeps an explicit batch even with a large context and small VRAM", func() {
|
||||
localGPU = func() config.GPU { return config.GPU{VRAM: 20 * gib} }
|
||||
cfg := singlePassCfg(40960)
|
||||
cfg.Batch = 512
|
||||
Expect(EffectiveBatchSize(cfg)).To(Equal(512))
|
||||
})
|
||||
|
||||
It("stays conservative (default, not the context) when per-device VRAM is unknown", func() {
|
||||
// A detection gap (VRAM 0) must not fall back to the unbounded context —
|
||||
// that's exactly what would OOM the load.
|
||||
localGPU = func() config.GPU { return config.GPU{VRAM: 0} }
|
||||
Expect(EffectiveBatchSize(singlePassCfg(40960))).To(Equal(DefaultBatchSize))
|
||||
})
|
||||
|
||||
It("returns the default batch for a non-single-pass model regardless of VRAM", func() {
|
||||
localGPU = func() config.GPU { return config.GPU{VRAM: 20 * gib} }
|
||||
ctx := 40960
|
||||
cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
|
||||
Expect(EffectiveBatchSize(cfg)).To(Equal(DefaultBatchSize))
|
||||
})
|
||||
})
|
||||
|
||||
// Guards the generic chat_template_kwargs forwarding: the model config map plus any
|
||||
// per-request metadata overrides are merged, coerced, and serialised into the
|
||||
// backend metadata blob that llama.cpp reads. Client metadata also overrides the
|
||||
|
||||
@@ -149,6 +149,47 @@ func largeContextForDevice(g GPU, ctx int) bool {
|
||||
return extra > g.VRAM/blackwellBatchHeadroomDivisor
|
||||
}
|
||||
|
||||
// SinglePassBatchForContext caps the physical batch (n_batch / n_ubatch) for a
|
||||
// single-pass load — embedding, score and rerank all decode/pool the whole input
|
||||
// in ONE physical batch, so they want a batch >= the input length to avoid the
|
||||
// GGML_ASSERT(n_tokens <= n_batch) abort and the "input is too large to process"
|
||||
// error. The naive choice is batch == context, but n_ubatch == context turns the
|
||||
// per-device CUDA compute buffer (which scales ~ n_ubatch * n_ctx and is NOT
|
||||
// split across GPUs) into multi-GiB of scratch that must fit on a SINGLE card, so
|
||||
// a large-context embedding model aborts on load (exitCode=-1) even with plenty
|
||||
// of free VRAM — the same #10485 root cause the Blackwell batch boost guards
|
||||
// against, which the single-pass path previously bypassed entirely.
|
||||
//
|
||||
// So instead of the full context we return the LARGEST batch whose compute buffer
|
||||
// fits the per-device VRAM headroom (VRAM / blackwellBatchHeadroomDivisor),
|
||||
// clamped to [DefaultPhysicalBatch, ctx]. The tradeoff: an input longer than the
|
||||
// returned cap can no longer be pooled in a single pass — but a batch that OOMs
|
||||
// the device processes nothing at all.
|
||||
//
|
||||
// g.VRAM must be the PER-DEVICE ceiling (smallest device on a multi-GPU host).
|
||||
// VRAM 0 (unknown) stays conservative and returns DefaultPhysicalBatch rather
|
||||
// than the unbounded context, so a detection gap can't OOM the load.
|
||||
func SinglePassBatchForContext(g GPU, ctx int) int {
|
||||
if ctx <= DefaultPhysicalBatch {
|
||||
return DefaultPhysicalBatch
|
||||
}
|
||||
if g.VRAM == 0 {
|
||||
return DefaultPhysicalBatch
|
||||
}
|
||||
perBatchCell := uint64(ctx) * computeBufferBytesPerCell
|
||||
if perBatchCell == 0 {
|
||||
return DefaultPhysicalBatch
|
||||
}
|
||||
batchCap := int(g.VRAM / blackwellBatchHeadroomDivisor / perBatchCell)
|
||||
if batchCap < DefaultPhysicalBatch {
|
||||
return DefaultPhysicalBatch
|
||||
}
|
||||
if batchCap > ctx {
|
||||
return ctx
|
||||
}
|
||||
return batchCap
|
||||
}
|
||||
|
||||
// IsManagedPhysicalBatch reports whether n is a value PhysicalBatch assigns.
|
||||
// Callers that re-tune a value chosen by an upstream host (the distributed
|
||||
// router correcting the frontend's guess) use this to avoid clobbering an
|
||||
@@ -254,6 +295,14 @@ var localGPU = func() GPU {
|
||||
}
|
||||
}
|
||||
|
||||
// LocalGPU exposes the locally-detected device descriptor to other packages
|
||||
// (e.g. core/backend's single-pass batch sizing) so they resolve the same
|
||||
// per-device VRAM this package's heuristics reason about. It goes through the
|
||||
// injectable localGPU var, so a config-package test seam also affects callers.
|
||||
func LocalGPU() GPU {
|
||||
return localGPU()
|
||||
}
|
||||
|
||||
// ApplyHardwareDefaults fills ModelConfig values that depend on the target GPU
|
||||
// and were left unset by the user. Currently: a larger physical batch on
|
||||
// Blackwell. Explicit config always wins (we only touch zero values).
|
||||
|
||||
@@ -46,3 +46,38 @@ var _ = Describe("SetDefaults hardware defaults (single-instance)", func() {
|
||||
Expect(cfg.Batch).To(Equal(1024))
|
||||
})
|
||||
})
|
||||
|
||||
// SinglePassBatchForContext is the VRAM-aware cap for the single-pass
|
||||
// (embedding/score/rerank) batch — the compute buffer scales ~ n_ubatch * n_ctx
|
||||
// and must fit a single device, so a large context can't take the full context
|
||||
// as its batch (issue #10485).
|
||||
var _ = Describe("SinglePassBatchForContext", func() {
|
||||
const gib = uint64(1) << 30
|
||||
|
||||
It("returns the default when the context is at or below the default batch", func() {
|
||||
Expect(SinglePassBatchForContext(GPU{VRAM: 119 * gib}, DefaultPhysicalBatch)).To(Equal(DefaultPhysicalBatch))
|
||||
Expect(SinglePassBatchForContext(GPU{VRAM: 119 * gib}, 256)).To(Equal(DefaultPhysicalBatch))
|
||||
})
|
||||
|
||||
It("returns the full context when the compute buffer fits ample VRAM", func() {
|
||||
// 4096 ctx on 119 GiB: the compute buffer is tiny, so the batch covers
|
||||
// the whole context (single-pass pooling in one physical batch).
|
||||
Expect(SinglePassBatchForContext(GPU{VRAM: 119 * gib}, 4096)).To(Equal(4096))
|
||||
})
|
||||
|
||||
It("caps below the context when a large context would overflow the VRAM headroom", func() {
|
||||
batch := SinglePassBatchForContext(GPU{VRAM: 20 * gib}, 40960)
|
||||
Expect(batch).To(BeNumerically(">=", DefaultPhysicalBatch))
|
||||
Expect(batch).To(BeNumerically("<", 40960))
|
||||
// The compute buffer for the capped batch must fit VRAM/headroom.
|
||||
Expect(uint64(batch) * 40960 * computeBufferBytesPerCell).To(BeNumerically("<=", (20*gib)/blackwellBatchHeadroomDivisor))
|
||||
})
|
||||
|
||||
It("never caps below the default batch even when VRAM is very tight", func() {
|
||||
Expect(SinglePassBatchForContext(GPU{VRAM: 1 * gib}, 40960)).To(Equal(DefaultPhysicalBatch))
|
||||
})
|
||||
|
||||
It("stays conservative (default) when per-device VRAM is unknown", func() {
|
||||
Expect(SinglePassBatchForContext(GPU{VRAM: 0}, 40960)).To(Equal(DefaultPhysicalBatch))
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user