-
Notifications
You must be signed in to change notification settings - Fork 205
Support Phi-4-mini #216
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Support Phi-4-mini #216
Conversation
85a0501
to
00da9ff
Compare
// Extension to Embedding to support using it as a linear layer | ||
extension Embedding { | ||
public func asLinear(_ x: MLXArray) -> MLXArray { | ||
// Use the embedding weights as a linear layer | ||
// This is equivalent to x @ weight.T | ||
return matmul(x, weight.transposed()) | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The embedding layer already has that: https://github.com/ml-explore/mlx-swift/blob/main/Source/MLXNN/Embedding.swift#L39-L45
Can you remove that code?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, I've removed it.
00da9ff
to
0ac6107
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Very nice, thanks!!
return model.embedTokens.asLinear(out) | ||
} else { | ||
return lmHead!(out) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I suggest using this structure -- what is here matches the python but I think this is more idiomatic swift:
if let lmHead {
out = lmHead(out)
} else {
out = model.embedTokens.asLinear(out)
}
(fine with the returns rather than the out =
, but the if let
is better I think)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks. I've used your suggestion and added a fatal error for the fall-through case, which shouldn't normally happen.
Based on ml-explore/mlx-examples#1305
I've verified that this works in llm-tool with these arguments:
--model microsoft/Phi-4-mini-instruct --prompt "Write a story about Einstein" -m 2048
I've also verified that it works with Phi 3.5 mini and Phi 4.
This is ready for review.