Conversation
Contributor
Author
|
This now depends on ml-explore/mlx#3742 , I will update this PR when that one is going to be merged. |
davidkoski
reviewed
Jun 30, 2026
davidkoski
reviewed
Jun 30, 2026
aleroot
force-pushed
the
loading_progress
branch
2 times, most recently
from
August 18, 2026 15:03
5140312 to
fdb7f7a
Compare
aleroot
added a commit
to aleroot/mlx-swift-lm
that referenced
this pull request
Aug 18, 2026
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.
aleroot
added a commit
to aleroot/mlx-swift-lm
that referenced
this pull request
Aug 18, 2026
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.
aleroot
added a commit
to aleroot/mlx-swift-lm
that referenced
this pull request
Aug 18, 2026
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.
4 tasks
Introduces LoadProgress and new loadArrays(url:stream:progressHandler:) and loadArraysAndMetadata(url:stream:progressHandler:) overloads that report bytes read as lazy arrays are evaluated. Uses a custom mlx_io_reader vtable backed by pread() so progress callbacks can be invoked from MLX worker threads. Includes a unit test verifying monotonic progress from 0 to 1.
`loadWeights()` style helpers -- including the one in mlx-swift-lm -- call the
plain `loadArrays(url:)` / `loadArraysAndMetadata(url:)`, so a per-call
`progressHandler:` argument can never reach them without changing every caller
along the way.
Add `withLoadProgressHandler(_:_:)` (sync and async), a task local scoped
handler in the style of `withErrorHandler(_:_:)`. The plain file loading
functions report byte progress to it when one is installed, so an application
can drive a precise model loading progress bar around code it does not own:
let container = try await withLoadProgressHandler({ tracker.update($0) }) {
try await factory.loadContainer(from: directory, using: tokenizerLoader)
}
`LoadProgress` gains the `url` of the file being read so progress can be
aggregated across the shards of a sharded model.
Also fix the SEEK_END case of the in-memory reader, which moved the offset
relative to the current position instead of the end of the data. mlx 0.32.1
seeks to the end of the stream to validate the tensor data offsets against the
size of the file, so `loadArrays(data:)` would fail there ("The JSON header is
N bytes long but the file is only 8 bytes").
Finally, restructure the truncated file test: a truncated file may now be
reported either eagerly, while the header is parsed, or lazily, when the arrays
are evaluated, and neither happens with the currently vendored mlx/mlx-c.
ml-explore/mlx-c#130 changed `mlx_io_vtable` so that the callbacks report whether they succeeded: int (*seek)(void*, int64_t off, int whence); size_t (*read)(void*, char* data, size_t n); size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); size_t (*write)(void*, const char* data, size_t n); and `CReader`/`CWriter` now turn a negative seek or a short read/write into a thrown `std::runtime_error`. Bump the mlx-c submodule to that commit (plus the checked-in copies of `io_types.h` and the CMake `GIT_TAG`) and adapt the Swift side: - `FileIOState.seek` returns `0`/`-1` instead of silently ignoring a bad `whence` or a negative resulting offset. - `FileIOState.read` returns the number of bytes actually read, so a file that is truncated after its header was parsed now fails instead of leaving the destination buffer uninitialized. The private implementation is renamed `readBytes` so it cannot be confused with the two public overloads. - The in-memory reader reports an out-of-bounds read as `0` bytes rather than doing nothing, and its `seek` reports failure for an unknown `whence` or a negative offset. - The file reader's `write` reports `0` bytes, so using it as a writer errors out instead of silently discarding the data. This is the error-propagation path the load progress work was waiting on (ml-explore/mlx#3742 is already in the vendored mlx v0.32.2), so the truncated file test asserts a failure rather than skipping, and a new test truncates the file *after* the header is parsed to cover the lazy read path: [mlx_io_reader] unable to read 65536 bytes (read 65504 instead) in file ... Also fix the `withLoadProgressHandler` documentation links -- the `-(_,()throws->R)` disambiguation does not resolve and `verify-docs.sh` builds with `--warnings-as-errors` -- document why the reported progress is approximate (loading is lazy, so it may stop short of the file size and it counts bytes read rather than bytes covered), and add `LoadProgress` to the MLX topics.
aleroot
force-pushed
the
loading_progress
branch
from
September 15, 2026 11:05
fdb7f7a to
5800740
Compare
Contributor
Author
This was merge and this ml-explore/mlx-c#130 one as well, so I have updated the work on top of those. |
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 byte-level progress reporting for safetensors file loading so upper layers can drive precise model-loading progress UI.
This PR depends on ml-explore/mlx#3734.
The progress API in mlx-swift can report byte-level safetensors loading progress, but the truncated/failed-read correctness path needs the upstream MLX fix first. Without ml-explore/mlx#3734, CPU lazy-load read failures can be swallowed before they reach Swift.
Once ml-explore/mlx#3734 eventaully lands, I’ll update the MLX submodule reference here and rerun the mlx-swift test suite.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes