Skip to content
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

MLX model support #124

Open
wants to merge 29 commits into
base: main
Choose a base branch
from
Open

MLX model support #124

wants to merge 29 commits into from

Conversation

ZachNagengast
Copy link
Contributor

@ZachNagengast ZachNagengast commented Apr 23, 2024

Draft PR for the early stages of supporting MLX based whisper models directly in WhisperKit. (To be updated)

Initial TODOs:

  • Setup swift package structure
  • MLXFeatureExtractor using FeatureExtracting protocol
  • MLXAudioEncoder inference using AudioEncoding protocol
  • MLXTextDecoder inference using TextDecoding protocol or TextDecoder subclass (research needed)
  • Downloading MLX models (research needed to mix and match model downloads between CoreML and MLX)
  • Loading MLX models with a new model protocol
  • Pipeline allowing mixing CoreML and MLX models
  • Compiling whisperkit-cli optionally with/without MLX Crossing this out for now: whisperkit-cli does not have the same bundle size constraints as the core library
  • CI unit tests for iOS platform using xcodebuild instead of iOS simulator (simulator is not supported for MLX)
  • Review, cleanup, and documentation

@ZachNagengast ZachNagengast marked this pull request as draft April 23, 2024 04:57
ZachNagengast and others added 5 commits April 22, 2024 21:57
* Added MLX feature extractor implementation

* CI fix

* added better multiarray conversion

* CI fix

* CI fix

* fixed `asMLMultiArray` implementation, fixed CI

* update xcode, trigger pr when targeting not main branch

* check if vision os builds

* update watch os version

* conditional watchos compilation

* conditional package.swift

* conditional package.swift

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* ci fix

* add other tests targest back

* package.swift cleanup

* general cleanup

* revert to xcode 15.2
* added mlx audio encoder

* fixed model protocols

* removed not needed
@ZachNagengast
Copy link
Contributor Author

@jkrukowski relevant: https://ml-explore.github.io/mlx/build/html/install.html#binary-size-minimization

We'll want to minimize the binary size as much as reasonable for iOS deployment

@ZachNagengast ZachNagengast marked this pull request as ready for review June 12, 2024 11:56
* update mlx-swift

* - reverted mlx version
- updated readme
- updated makefile

* reversed

* fixed tests

* updated mlx-swift

* updated makefile

* remove device change, fft can run on gpu now

* updated readme, added tests

* updated readme

* review changes

* review changes

* CI model cache path fix

* tests failed

* Update package.swift

* Keep setupModels with adjustments

* Test CI skip cache

* Test CI package name change

* Test CI optional CLI settings

* Use correct logits size for MLX

---------

Co-authored-by: ZachNagengast <[email protected]>
}

private func qkvAttention(_ q: MLXArray, _ k: MLXArray, _ v: MLXArray, _ mask: MLXArray?) -> (MLXArray, MLXArray) {
let (nBatch, nCtx, nState) = (q.shape[0], q.shape[1], q.shape[2])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
Status: In Progress
Development

Successfully merging this pull request may close these issues.

3 participants