diff --git a/CHANGELOG.md b/CHANGELOG.md index b5ba9dc3..92db48d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,12 @@ APIs change between minors without a deprecation period. ### Fixed +- `llmgw ctl model serve` served an embedding model as a chat model: no + `--embedding`, four slots, a 512-token batch, `embeddings` unset and a + "say hello" smoke test it could not pass. A model whose GGUF says it + pools (or `--embeddings`) is served for embedding: 32 slots of 1024 + tokens unless sized otherwise, a batch that holds a whole passage, and + one vector asked for as its smoke test. - `llmgw ctl upgrade` checks every node's version once the waves are done and fails naming any not on the release, and says when the control plane has no nodes registered; `llmgw upgrade` on a control plane says how many diff --git a/internal/cli/model_serve.go b/internal/cli/model_serve.go index 47fdb303..e993b323 100644 --- a/internal/cli/model_serve.go +++ b/internal/cli/model_serve.go @@ -60,6 +60,8 @@ func newModelServeCmd(client func() *ctlClient) *cobra.Command { f.StringVar(&req.Quant, "quant", "", "the quantisation to serve (default: the one that suits the cluster best)") f.StringVar(&req.File, "file", "", "the file to serve, instead of choosing by --quant") f.StringVar(&req.Projector, "projector", "", "a vision model's mmproj file, fetched and loaded beside it") + f.BoolVar(&req.Embeddings, "embeddings", false, "serve it as an embedding model (its GGUF header says so for most): "+ + "--embedding, 32 slots of 1024 tokens unless --ctx-pool or --context-length says otherwise") f.StringVar(&req.Format, "format", "gguf", "gguf, run by llama.cpp, or safetensors: the original weights, run by vLLM on GPUs") f.BoolVar(&req.CPU, "cpu", false, "run it on the CPU, whatever the plan says") f.StringVar(&req.GPUClass, "gpu-class", "", "run it on GPUs of this class, whatever the plan says") diff --git a/internal/control/modelserve/canary.go b/internal/control/modelserve/canary.go index aff78e33..b7a134d3 100644 --- a/internal/control/modelserve/canary.go +++ b/internal/control/modelserve/canary.go @@ -92,7 +92,7 @@ func (s *Server) addCandidate(ctx context.Context, j *Job, p prepared, cn catalo s.step(j, "smoke test", Running, "asking "+r.ID+" to say hello") var reply string var tps float64 - if reply, tps, err = s.smoke(ctx, r, cand); err == nil { + if reply, tps, err = s.smoke(ctx, r, cand, s.embeds(cand)); err == nil { s.step(j, "smoke test", Done, fmt.Sprintf("answered at %.1f tokens/s: %q", tps, reply)) } else { s.step(j, "smoke test", Failed, err.Error()) diff --git a/internal/control/modelserve/modelserve.go b/internal/control/modelserve/modelserve.go index 53105945..1bb9236a 100644 --- a/internal/control/modelserve/modelserve.go +++ b/internal/control/modelserve/modelserve.go @@ -153,6 +153,13 @@ type Request struct { // is written into the catalog, where clients and editors read it. Check // a running replica with llmgw ctl tool-check before saying so. Tools bool `json:"tools,omitempty"` + // Embeddings serves it as an embedding model, which its GGUF header + // says by itself for most: --embedding, many slots, a batch that holds + // a whole passage, /v1/embeddings. + Embeddings bool `json:"embeddings,omitempty"` + // sized is the caller having said --ctx-pool or --context-length, which + // an embedding model's defaults then leave alone. + sized bool // CPUCores is what a replica computes on; zero takes half the smallest // node it fits on. CPUCores int `json:"cpu_cores,omitempty"` @@ -268,6 +275,7 @@ func (s *Server) Start(req Request) (Job, error) { if req.GuardFormat != "" && req.Replicas == 0 { return Job{}, errors.New("a judge model must always run: give it --replicas 1 or more, not --on-demand") } + req.sized = req.CtxPool > 0 || req.ContextLength > 0 if req.CtxPool <= 0 { req.CtxPool = 16384 } @@ -473,6 +481,16 @@ func (s *Server) run(j *Job) { } header = v.Header } + if header.Embedding { + req.Embeddings = true + } + if req.Embeddings && !req.sized { + // Many passages at once rather than four: an import of a hundred + // thousand chunks went at the pace of four slots with the card a + // third busy. + req.ContextLength, req.CtxPool = embedSlotTokens, embedSlotTokens*embedSlots + } + j.Request = req memGB := int((header.MemoryBytes(set.Bytes, req.CtxPool) + 1<<30 - 1) >> 30) basis := fmt.Sprintf("weights %.1f GB + KV cache %.1f GB for %d tokens + engine", gb(set.Bytes), gb(int64(req.CtxPool)*header.KVBytesPerToken()), req.CtxPool) @@ -587,10 +605,24 @@ func profilesFor(layout []modelplan.Option, req Request, memGB, cores int) []cat // slots of 8192, and the gateway, told 32768, sent requests the engine then // refused. func llamaArgs(req Request) []string { - return []string{"--ctx-size", strconv.Itoa(req.CtxPool), + args := []string{"--ctx-size", strconv.Itoa(req.CtxPool), "--parallel", strconv.Itoa(max(1, req.CtxPool/req.ContextLength))} + if req.Embeddings { + // A passage is embedded in one batch: with the default 512, one of + // 516 tokens was refused. + n := strconv.Itoa(req.ContextLength) + args = append([]string{"--embedding"}, append(args, "--batch-size", n, "--ubatch-size", n)...) + } + return args } +// An embedding model's slots, when its caller did not size them: passages +// of up to a thousand tokens, thirty-two at a time. +const ( + embedSlotTokens = 1024 + embedSlots = 32 +) + // maxOutput is the longest answer the catalog allows: a quarter of the // context, from 1024 to 8192. 1024 everywhere cut an editor's agent off in // the middle of the file it was writing. @@ -620,7 +652,7 @@ func (s *Server) addAndStart(ctx context.Context, j *Job, repo hf.Repo, set hf.F RoutingProfile: "chat", ContextLength: req.ContextLength, MaxOutput: maxOutput(req.ContextLength), Profiles: profs, Scaling: catalog.Scaling{Min: req.Replicas, Max: max(req.Replicas, 1)}, - Capabilities: catalog.Capabilities{Vision: req.Projector != "", Tools: req.Tools}, + Capabilities: catalog.Capabilities{Vision: req.Projector != "", Tools: req.Tools, Embeddings: req.Embeddings}, Origin: &catalog.Origin{Repo: repo.ID, Revision: repo.Revision, Files: originFiles(set)}, } if req.GuardFormat != "" { @@ -655,7 +687,7 @@ func (s *Server) addAndStart(ctx context.Context, j *Job, repo hf.Repo, set hf.F // 6. Does it answer. s.step(j, "smoke test", Running, "asking "+rep.ID+" to say hello") - reply, tps, err := s.smoke(ctx, rep, name) + reply, tps, err := s.smoke(ctx, rep, name, req.Embeddings) if err != nil { s.fail(j, "smoke test", err) return @@ -859,7 +891,10 @@ func (s *Server) waitReady(ctx context.Context, j *Job, name string, hardware bo } // smoke sends one short request straight to the replica, with its engine key. -func (s *Server) smoke(ctx context.Context, r reconciler.Replica, name string) (string, float64, error) { +func (s *Server) smoke(ctx context.Context, r reconciler.Replica, name string, embeddings bool) (string, float64, error) { + if embeddings { + return s.smokeEmbeddings(ctx, r, name) + } body, _ := json.Marshal(map[string]any{ "model": name, "max_tokens": 32, @@ -989,3 +1024,50 @@ func pickSafetensors(repo hf.Repo, _ Request) (hf.FileSet, error) { } return set, nil } + +// smokeEmbeddings asks an embedding model for one vector, and says its +// length: an embedding model does not chat. +func (s *Server) smokeEmbeddings(ctx context.Context, r reconciler.Replica, name string) (string, float64, error) { + body, _ := json.Marshal(map[string]any{"model": name, "input": "hello"}) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + strings.TrimRight(r.Endpoint, "/")+"/v1/embeddings", bytes.NewReader(body)) + if err != nil { + return "", 0, err + } + req.Header.Set("Content-Type", "application/json") + if k := s.deps.Fleet.EngineKey(r.ID); k != "" { + req.Header.Set("Authorization", "Bearer "+k) + } + resp, err := s.deps.HTTP.Do(req) + if err != nil { + return "", 0, err + } + defer resp.Body.Close() + b, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + if resp.StatusCode != http.StatusOK { + return "", 0, fmt.Errorf("the replica answered %s: %s", resp.Status, strings.TrimSpace(string(b))) + } + var out struct { + Data []struct { + Embedding []float64 `json:"embedding"` + } `json:"data"` + } + if err := json.Unmarshal(b, &out); err != nil || len(out.Data) == 0 || len(out.Data[0].Embedding) == 0 { + return "", 0, fmt.Errorf("the replica answered no vector: %s", strings.TrimSpace(string(b))) + } + return fmt.Sprintf("a vector of %d dimensions", len(out.Data[0].Embedding)), 0, nil +} + +// embeds is whether the catalog serves name as an embedding model, for a +// smoke test of one already there (a canary, an update). +func (s *Server) embeds(name string) bool { + if s.deps.Current == nil { + return false + } + for _, m := range s.deps.Current() { + if m.Name == name { + return m.Capabilities.Embeddings + } + } + return false +} diff --git a/internal/control/modelserve/modelserve_test.go b/internal/control/modelserve/modelserve_test.go index 1606b1b8..c186a968 100644 --- a/internal/control/modelserve/modelserve_test.go +++ b/internal/control/modelserve/modelserve_test.go @@ -397,3 +397,18 @@ func TestOnlyTheLastFinishedJobsAreRemembered(t *testing.T) { t.Error("the oldest finished job is still remembered") } } + +func TestAnEmbeddingModelIsServedForEmbedding(t *testing.T) { + // Four slots and a 512-token batch: an import went at four passages a + // pass, and one passage of 516 tokens was refused. + args := strings.Join(llamaArgs(Request{Embeddings: true, ContextLength: embedSlotTokens, + CtxPool: embedSlotTokens * embedSlots}), " ") + for _, want := range []string{"--embedding", "--ctx-size 32768", "--parallel 32", "--batch-size 1024", "--ubatch-size 1024"} { + if !strings.Contains(args, want) { + t.Errorf("%q not in %q", want, args) + } + } + if args := strings.Join(llamaArgs(Request{ContextLength: 4096, CtxPool: 16384}), " "); strings.Contains(args, "embedding") { + t.Errorf("a chat model got %q", args) + } +} diff --git a/internal/control/modelserve/updates.go b/internal/control/modelserve/updates.go index 0fb32a25..57e03633 100644 --- a/internal/control/modelserve/updates.go +++ b/internal/control/modelserve/updates.go @@ -689,7 +689,7 @@ func (s *Server) replace(ctx context.Context, j *Job, name string) error { } s.step(j, "replicas", Done, fmt.Sprintf("%d replaced", len(old))) s.step(j, "smoke test", Running, "asking "+ready[0].ID+" to say hello") - reply, tps, err := s.smoke(ctx, ready[0], name) + reply, tps, err := s.smoke(ctx, ready[0], name, s.embeds(name)) if err != nil { s.step(j, "smoke test", Failed, err.Error()) return err diff --git a/internal/gguf/gguf.go b/internal/gguf/gguf.go index 558d16d1..b9842f14 100644 --- a/internal/gguf/gguf.go +++ b/internal/gguf/gguf.go @@ -41,6 +41,10 @@ type Header struct { // them; otherwise a head is EmbeddingLength / HeadCount. KeyLength uint64 ValueLength uint64 + // Embedding is whether the model pools its tokens into one vector + // (.pooling_type set and not none): an embedding model, which + // llama.cpp serves with --embedding and does not chat. + Embedding bool } // KVBytesPerToken is what one token of context costs in llama.cpp's default @@ -265,6 +269,8 @@ func Read(r io.Reader) (Header, error) { h.KeyLength = num case "attention.value_length": h.ValueLength = num + case "pooling_type": + h.Embedding = num > 0 } } // The tokenizer's arrays — megabytes of vocabulary — come after the diff --git a/internal/gguf/gguf_test.go b/internal/gguf/gguf_test.go index 81fbabec..ff687706 100644 --- a/internal/gguf/gguf_test.go +++ b/internal/gguf/gguf_test.go @@ -163,3 +163,28 @@ func TestArraysNestedDeepAreRefused(t *testing.T) { t.Errorf("ten thousand nested arrays: %v", err) } } + +func TestAnEmbeddingModelIsKnownByItsPooling(t *testing.T) { + build := func(pooling uint32) []byte { + var b bytes.Buffer + w := func(v any) { _ = binary.Write(&b, binary.LittleEndian, v) } + str := func(s string) { w(uint64(len(s))); b.WriteString(s) } + u32 := func(k string, v uint32) { str(k); w(uint32(tUint32)); w(v) } + b.WriteString("GGUF") + w(uint32(3)) + w(uint64(0)) + w(uint64(3)) + str("general.architecture") + w(uint32(tString)) + str("bert") + u32("bert.block_count", 24) + u32("bert.pooling_type", pooling) + return b.Bytes() + } + if h, err := Read(bytes.NewReader(build(2))); err != nil || !h.Embedding { + t.Errorf("CLS pooling: %+v %v", h, err) + } + if h, err := Read(bytes.NewReader(build(0))); err != nil || h.Embedding { + t.Errorf("no pooling: %+v %v", h, err) + } +}