graph: skip the LM head for pooled embeddings (#30445)
* graph: skip the LM head for pooled embeddings A context that extracts pooled embeddings outputs one vector per sequence, but every model still expanded its LM head, so the graph computed n_vocab logits per token and the context kept a host buffer for them. The models now only declare t_logits, and set_outputs() expands it unless the context extracts embeddings with a pooling type other than NONE. Since the toggle changes the graph shape for a pooled context, set_embeddings() schedules a re-reserve in that case. * graph: drop the redundant t_embd expand Address review from @CISC: with embeddings on, build_pooling() already expands its output for every pooling type, and a decoder graph without embeddings does not read t_embd, so set_outputs() only expands t_logits. * graph: add build_logits() for the LM head output Address review from @ggerganov: the LM head is expanded by a dedicated build_logits(), called from llama_model::build_graph() like build_pooling(), and set_outputs() only marks the outputs again. * graph: call build_logits() before build_sampling() Address review from @ggerganov: build_sampling() reads t_logits, so the LM head is built first. A pooled context then skips the sampling layers too, since it has no logits to sample from. * nit
This commit is contained in:
1 parent
a17e94d229
commit
cb55ffe6e2
136 files changed
+27
-282
No files matched your search
@@ -145,6 +145,4 @@ llama_model_maple::graph::graph(const llama_model & model, const llm_graph_param
|
||||
cur = build_lora_mm(model.output, cur, model.output_s);
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
Reference in new issue
Block a user