Repository navigation
Sync upstream ml-explore/mlx-swift 392aa34 (mlx v0.32.4) - #29
Merged
Merged
Conversation
…re#462) * Add concurrent synchronous safetensors loading across files The safetensors loader is lazy and a single eval of a whole checkpoint visits tensors in dictionary order, serializing reads at effectively random file offsets. Add loadArrays(urls:) and loadArraysAndMetadata(urls:), which split every file into contiguous byte-balanced ranges in file-offset order and evaluate the ranges from concurrent work items, so reads are sequential and header parsing, I/O, and copies overlap. This moves the technique from ml-explore/mlx-swift-lm#575 down into mlx-swift, per review feedback there. Measured on an M4 Pro with NVMe at ~6.1 GB/s: an 18 GB checkpoint loads 20% faster cold (4.5 to 5.6 GB/s) and a 10 GB checkpoint 5-8% faster warm; parity when the serial loader is already disk-bound.
) * Add byte-progress reporting for safetensors file loading 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. * Add a scoped load-progress handler and fix memory reader SEEK_END `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. * Update for the mlx-c custom IO reader/writer API 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. * Harden progress reads and cover lazy handler lifetime
- specifically -Wno-shorten-64-to-32 - reduces build output from 148898 to 5020 lines
…-explore#509) Both took the target logits along `axis` but expanded the targets and squeezed the result on axis -1, so any other axis aborted in squeeze with "Cannot squeeze axis -1 with size N". They now expand and squeeze on `axis`, as the Python reference does. Two tests pin the values Python returns for axis 0. Fixes ml-explore#507
…ml-explore#443) The CudaBuild plugin hardwires nvcc's host compiler to the Swift toolchain's own clang++, which cannot be overridden. When that clang is outside the range the installed CUDA release accepts, nvcc refuses to compile anything: crt/host_config.h:151:2: error: -- unsupported clang version! clang version must be less than 21 and greater than 3.2 This is reachable today: Swift 6.3.3 ships clang 21.0.0, while CUDA 13.0 accepts only clang < 21. The plugin passes its resolved path to encuda as --clangpp, which encuda appends as -ccbin= *after* anything in NVCC_PREPEND_FLAGS, so that environment variable cannot win either. Resolve the host compiler from, in priority order: the MLX_CUDA_HOST_COMPILER environment variable, a new `hostCompiler` key in CudaBuild.json, then the toolchain's clang++ as before. Existing configurations are unaffected. It must still be a clang: nvcc preprocesses with the host compiler and SwiftPM compiles the resulting .cpp, so pointing -ccbin at g++ yields a translation unit full of GCC-only `_Float32` declarations that clang then rejects. Also stamp the compilation configuration beside each generated .cpp and treat a changed signature as out of date. encuda's --incremental check compared only input/output modification times, which cannot observe a changed host compiler, std, arch, or include path — so switching host compilers silently reused .cpp files preprocessed by the previous one. That produced a flood of errors from glibc headers rather than anything pointing at the real cause. Co-authored-by: Claude Opus 5 <[email protected]>
Two include roots are hardcoded or missing in the SwiftPM CUDA build, and
neither can be pointed elsewhere.
CCCL is hardcoded to the copy inside the CUDA toolkit, for both the `-I` and the
`MLX_CCCL_DIR` define that NVRTC uses at runtime. MLX pins its own CCCL version,
and when the toolkit ships a different one the mismatch surfaces as overload
ambiguity in device code rather than anything naming a version:
mlx/backend/cuda/device/binary_ops.cuh(50): error: more than one instance of
overloaded function "cuda::std::fmod" matches the argument list
Observed with CUDA 13.0 (CCCL 3.0.1) against MLX's pinned CCCL 3.1.3.
cudnn-frontend and CUTLASS are not provided at all. The CMake build fetches both
via FetchContent, but SwiftPM has no equivalent, so the build fails on headers
that are simply absent:
mlx/backend/cuda/cudnn_utils.h:9:10: fatal error: 'cudnn_frontend.h' file not found
mlx/backend/cuda/quantized/qmm/qmv.cu:9:10: fatal error: 'cute/numeric/numeric_types.hpp' file not found
Adds `MLX_CCCL_DIR` to override the CCCL root (defaulting to the toolkit copy, so
existing behaviour is unchanged) and `MLX_CUDA_INCLUDE_PATHS`, a colon-separated
list of additional roots. Both are applied to the C++ sources via Package.swift
and to `.cu` compilation via the CudaBuild plugin, so device and host code
compile against the same headers.
This keeps the roots external rather than vendoring cudnn-frontend and CUTLASS as
submodules; that remains an option if their provisioning should be automatic
rather than configured.
Co-authored-by: Claude Opus 5 <[email protected]>
* Add optimizer state save and restore * Track Adafactor step so compiled checkpoints stay current * Fix optimizer state docs so DocC builds ModuleParameters lives in MLXNN, so a symbol link in MLXOptimizers fails the documentation build. stateSuffixes() also fails once exportState's comment is inherited, so that link is absolute.
* mlx v0.32.4 New API - MLXFast.gatedDeltaUpdate - Memory.bufferSize - matrixTransposed - uniqued Changed API - gatherQuantizedMM -- adds globalScale: MLXArray? = nil, biases now has a nil default
* fix compatibility issue in update(parameters:) - with the change in ml-explore#418 the implementation changed such that the recursion was on an internal method - this restores the previous functionality
Brings in upstream mlx-swift up to 392aa34 (9 commits past the 43c90af merge base), including the mlx v0.32.4 vendor update (ml-explore#508), optimizer state, CUDA build changes, crossEntropy/nllLoss, C++ warning fixes, concurrent safetensors loading (ml-explore#462), and update(parameters:) (ml-explore#510). Conflict resolution: - Source/Cmlx/mlx and Source/Cmlx/mlx-c: re-vendored from upstream (mlx 36969c09, mlx-c a26f2d58) with the carried fork patches re-applied, including the unsigned floor_divide shortcut in mlx/ops.cpp. - Source/Cmlx/mlx-generated: regenerated with tools/update-mlx.sh; moe_stream.metal (fork-only) copied from fork main; sdpa_vector.h keeps the TurboQuant changes. - Source/Cmlx/include-framework: regenerated with tools/update-mlx-xcodeproj.sh; fork headers (mlx-fast.h, mlx-threadpool.h, mlx-c-fast.h, mlx-io-load.h) retained. - Package.swift: fork version kept with upstream's cross_entropy.cpp and gated_delta_update.cpp sources and shorten-64-to-32 warning suppression; CUDA stubbing retained. - xcode/MLX.xcodeproj: upstream project plus moe_stream.metal. - Tests/MLXTests/FP8Tests.swift: upstream version taken; the fork's NaN-aware variant papered over pre-v0.32.4 decoding and is superseded. - Source/MLX/Transforms+Compile.swift: fork's scoped default device patch kept. - .github/*, .gitmodules, MAINTENANCE.md, fmt/test/*: fork-only, kept. Not compiled or tested locally; CI is the gate. Co-Authored-By: Claude Haiku 5.5 <[email protected]>
.disableWarning requires PackageDescription 6.2, but this manifest declares swift-tools-version 5.12, so the package failed to parse and every CI job was cancelled before any step ran. Use .unsafeFlags(["-Wno-shorten-64-to-32"]) instead, which 5.12 supports. Co-Authored-By: Claude Haiku 5.5 <[email protected]>
Upstream mlx changed array::buffer() to return allocator::Buffer by
value. The fork's prefault() took a non-const reference to that
temporary via const_cast, which fails to compile
("const_cast from rvalue to reference type"). Copy the buffer into a
local and read its raw pointer from that.
This broke the MLX Tests (SwiftPM metallib), Build xcode/MLX.xcodeproj
and Downstream Integration Test jobs on dcbbf1a.
Co-Authored-By: Claude Haiku 5.5 <[email protected]>
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
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.
Summary
Manual sync of SharpAI/mlx-swift fork main (d5839e5) with ml-explore/mlx-swift main (392aa34), 9 upstream commits, including the mlx v0.32.4 vendor update (ml-explore#508).
Resolved conflicts (14 files), all staged in merge commit c114aa0:
Source/Cmlx/mlx,Source/Cmlx/mlx-c: re-vendored from upstream (mlx 36969c09, mlx-c a26f2d58) with the carried fork patches re-applied (SSD streaming, TurboKV, I/O loaders with iOS guards, unsignedfloor_divide, mlx-cfast.{h,cpp}wrappers). Three hunks did not apply cleanly and were handled by hand; see the merge commit message.Source/Cmlx/mlx-generated: regenerated withtools/update-mlx.sh. Differences from upstream:cuda_jit_sources.h(regenerated, CUDA-only),sdpa_vector.h(TurboQuant),moe_stream.metal(fork-only).Source/Cmlx/include-framework: regenerated withtools/update-mlx-xcodeproj.sh; fork headers retained.Package.swift: fork version plus upstream'scross_entropy.cpp,gated_delta_update.cpp, andshorten-64-to-32suppression; CUDA stubbing kept.xcode/MLX.xcodeproj/project.pbxproj: upstream plusmoe_stream.metal.Tests/MLXTests/FP8Tests.swift: upstream version taken (the fork's NaN-aware variant is superseded by v0.32.4).Source/MLX/Transforms+Compile.swift: fork's scoped default-device patch kept.Open points for review
cuda_jit_sources.hdiffers from upstream's committed copy (CUDA-only; this fork does not build CUDA).FloorDividemay make the unsignedfloor_divideshortcut redundant; not verified.MLXFast.swift,Transforms+Compile.swiftandForkProtectionTests/SSDStreamingTestswere reviewed by reading only; they are not compiled.Testing
Not compiled or run locally. CI is the gate. Matrix row: SL-18.
🤖 Generated with Claude Code