Conversation
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.
Group the progress callbacks of a load in one value, `LoadProgressHandlers`:
let container = try await LLMModelFactory.shared.loadContainer(
from: directory, using: tokenizerLoader,
progress: .weights { progress in
print(progress.fractionCompleted)
})
It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.
Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.
The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.
Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
Contributor
Author
|
This one cannot build correctly before, the dependencies are ready. Sorry I should have created it as a draft probably ... |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
Adds a model weights progress handler.
Close the circle of ml-explore/mlx-swift#427 -> ml-explore/mlx-c#126 -> ml-explore/mlx#3742
Requires the first
mlx-swiftrelease containing ml-explore/mlx-swift#427, which in turn requires mlx v0.32.1Note
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes