add phi2 mem

1eb382da · Michael Yang · bb6fd022 · 1eb382da
Commit 1eb382da authored May 10, 2024 by Michael Yang
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

llm/ggml.go llm/ggml.go +4 -1

No files found.
--- a/llm/ggml.go
+++ b/llm/ggml.go
@@ -329,7 +329,10 @@ func (llm GGML) GraphSize(context, batch uint64) (partialOffload, fullOffload ui
 			4*batch*(1+4*embedding+context+context*heads),
 		)

-		partialOffload = 4*batch*(2*embedding+vocab) + embedding*vocab*105/128
+		partialOffload = max(
+			4*batch*(2*embedding+vocab)+embedding*vocab*105/128,
+			4*batch*(2+3*embedding+context+context*heads),
+		)
 	case "stablelm":
 		fullOffload = 4 * batch * (context*(1+heads) + 3*embedding + 2)
 		partialOffload = max(