zhangshuai/mlx-swift-examples
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix float types in Phi (use float16) (#25)

Browse Source 
- per suggestions in #23 ensure that the values that go into the cache are float16
This commit is contained in:
David Koski
2024-03-14 13:18:40 -07:00
committed by GitHub
parent a1431e7155
commit ac273a14ea
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
Libraries/LLM/Phi.swift
View File

@@ -73,14 +73,12 @@ private class PhiAttention: Module {
keys = rope(keys) keys = rope(keys)
} }
queries = queries.asType(Float.self)
keys = keys.asType(Float.self)
// Finally perform the attention computation // Finally perform the attention computation
let scale = sqrt(1 / Float(queries.dim(-1))) let scale = sqrt(1 / Float(queries.dim(-1)))
let output = MLXFast.scaledDotProductAttention( let output = MLXFast.scaledDotProductAttention(
queries: queries, keys: keys, values: values, scale: scale, mask: mask queries: queries.asType(.float32), keys: keys, values: values, scale: scale, mask: mask
) )
.asType(values.dtype)
.transposed(0, 2, 1, 3) .transposed(0, 2, 1, 3)
.reshaped(B, L, -1) .reshaped(B, L, -1)