Skip to content

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

Merged
merged 2 commits into from
Feb 27, 2025
Merged

Support Phi-4-mini #216

merged 2 commits into from
Feb 27, 2025

Conversation

DePasqualeOrg
Copy link
Contributor

@DePasqualeOrg DePasqualeOrg commented Feb 27, 2025

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.

@DePasqualeOrg DePasqualeOrg marked this pull request as ready for review February 27, 2025 09:49
Comment on lines 59 to 66
// 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())
}
}
Copy link
Member

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?

Copy link
Contributor Author

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.

Copy link
Member

@awni awni left a 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)
}
Copy link
Collaborator

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)

Copy link
Contributor Author

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.

@davidkoski davidkoski merged commit 0c50f71 into ml-explore:main Feb 27, 2025
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants