Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
57 commits
Select commit Hold shift + click to select a range
c3c6ba9
Add mission infrastructure for MLX distributed port
ronaldmannak Mar 14, 2026
14b92b7
Enable distributed C/C++ compilation in Package.swift
ronaldmannak Mar 14, 2026
58d78a4
Add scrutiny synthesis for distributed compilation
ronaldmannak Mar 14, 2026
5b849c1
Add user-testing synthesis for distributed compilation
ronaldmannak Mar 14, 2026
46f8ab2
Add Swift bindings for MLX distributed primitives
ronaldmannak Mar 14, 2026
1584469
Add comprehensive single-process distributed tests
ronaldmannak Mar 14, 2026
6777eab
Add multi-process distributed tests with ring backend
ronaldmannak Mar 14, 2026
066e0e2
Apply swift-format fixes to distributed files and update architecture…
ronaldmannak Mar 14, 2026
01c442a
Add scrutiny synthesis for swift bindings
ronaldmannak Mar 14, 2026
12b0365
Fix scrutiny review issues for swift-bindings milestone
ronaldmannak Mar 14, 2026
269ef95
Update swift-bindings scrutiny rerun synthesis
ronaldmannak Mar 14, 2026
d4453ee
Override scrutiny for swift-bindings: split unsupported upstream
ronaldmannak Mar 14, 2026
636f1ed
Add user-testing synthesis for swift bindings
ronaldmannak Mar 14, 2026
f0dad1d
Add AllToShardedLinear, ShardedToAllLinear, and sumGradients helper
ronaldmannak Mar 14, 2026
1b594ab
Add QuantizedAllToShardedLinear and QuantizedShardedToAllLinear
ronaldmannak Mar 14, 2026
0ccca43
Add ShardingType, shardLinear, shardInPlace, and averageGradients uti…
ronaldmannak Mar 14, 2026
384dc17
Add comprehensive distributed NN tests and fix shardLinear type dispatch
ronaldmannak Mar 14, 2026
467b1f2
Apply swift-format fixes to DistributedNNTests.swift
ronaldmannak Mar 14, 2026
4f17279
Add scrutiny synthesis for distributed nn layers
ronaldmannak Mar 14, 2026
d6e6c4a
Improve testNonDivisibleDimensionError to verify validation logic acr…
ronaldmannak Mar 14, 2026
4ad8bc7
Add round 2 distributed NN scrutiny results
ronaldmannak Mar 14, 2026
0ac11d4
Override scrutiny for distributed-nn-layers: precondition is MLXNN co…
ronaldmannak Mar 14, 2026
dfc9fc1
Record distributed NN user-testing validation
ronaldmannak Mar 14, 2026
574cb2f
Add multi-process collective ops tests for Python test parity
ronaldmannak Mar 14, 2026
dd8857f
Add communicationType parameter to averageGradients for Python parity
ronaldmannak Mar 14, 2026
f8ca770
Add multi-process NN parity tests for shard_linear forward and backward
ronaldmannak Mar 14, 2026
3e79dee
Add JACCL availability check test and document JACCL testing limitations
ronaldmannak Mar 14, 2026
4bf5ea9
Fix multi-process test flakiness with deterministic ports, retry logi…
ronaldmannak Mar 14, 2026
521b2f3
Fix multi-process test timeouts: use _exit(0) and handle valid JSON o…
ronaldmannak Mar 14, 2026
0062951
Record test-parity scrutiny findings
ronaldmannak Mar 14, 2026
2ed44f8
Override scrutiny for test-parity: resolved/upstream issues only
ronaldmannak Mar 14, 2026
36725e2
Record test-parity user-testing validation
ronaldmannak Mar 14, 2026
bb67b90
Remove protocol conformance tests
ronaldmannak Mar 15, 2026
308490e
Fix warnings
ronaldmannak Mar 15, 2026
155c726
Added todos
ronaldmannak Mar 15, 2026
4664fc1
better unit test coverage
ronaldmannak Mar 15, 2026
3e06b02
swift lint
ronaldmannak Mar 15, 2026
cd22fca
Add separate mlx-distributed skill
ronaldmannak Mar 15, 2026
e0d4059
Add .factory/ and .claude/ to .gitignore and remove from tracking
ronaldmannak Mar 15, 2026
dbad44e
Remove BoolBox
ronaldmannak Mar 15, 2026
4f71170
Fix errorbox
ronaldmannak Mar 15, 2026
a1988a7
Fix mixed-dtype averageGradients multi-process
ronaldmannak Mar 15, 2026
c80c1aa
Swift lint
ronaldmannak Mar 24, 2026
24eedb1
Integrate MLX-c 0.31.1
ronaldmannak Mar 24, 2026
3297f20
Remove DistributedWorker tests from CI
ronaldmannak Apr 3, 2026
3a32139
Remove Distributed worker and depricated setDefault(device:)'
ronaldmannak Apr 4, 2026
2d9da3c
Remove old tests
ronaldmannak Apr 4, 2026
ea76ccc
Refactor for more “swiftiness”
ronaldmannak Apr 4, 2026
5a893c5
Add default initializer
ronaldmannak Apr 4, 2026
d076d1c
Update comments
ronaldmannak Apr 4, 2026
1b5a35a
Update skills
ronaldmannak Apr 4, 2026
5e42346
Update skills
ronaldmannak Apr 4, 2026
1280a6a
Swift lint
ronaldmannak Apr 4, 2026
677f5d5
Remove unchecked sendable
ronaldmannak Apr 5, 2026
e8b930f
remove testing references in mulit-process.md
ronaldmannak Apr 6, 2026
cb4adcb
Improve API by moving withError to DistributedGroup
ronaldmannak Apr 6, 2026
ba68dfa
swift lint
ronaldmannak Apr 6, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,7 @@ iOSInjectionProject/

# VS Code
.vscode/

# AI agent working directories
.factory/
.claude/
945 changes: 945 additions & 0 deletions DISTRIBUTED-LM-INTEGRATION.md

Large diffs are not rendered by default.

25 changes: 15 additions & 10 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -111,10 +111,8 @@ let cmlx = Target.target(
// vendor docs
"vendor-README.md",

// example code + mlx-c distributed
// example code
"mlx-c/examples",
"mlx-c/mlx/c/distributed.cpp",
"mlx-c/mlx/c/distributed_group.cpp",

// vendored library, include header only
"json",
Expand Down Expand Up @@ -190,15 +188,12 @@ let cmlx = Target.target(
"mlx/mlx/backend/metal/kernels",
"mlx/mlx/backend/metal/nojit_kernels.cpp",

// do not build distributed support (yet)
// distributed backends: enable ring + JACCL, disable MPI + NCCL
"mlx/mlx/distributed/mpi/mpi.cpp",
"mlx/mlx/distributed/ring/ring.cpp",
"mlx/mlx/distributed/nccl/nccl.cpp",
"mlx/mlx/distributed/nccl/nccl_stub",
"mlx/mlx/distributed/jaccl/jaccl.cpp",
"mlx/mlx/distributed/jaccl/mesh.cpp",
"mlx/mlx/distributed/jaccl/ring.cpp",
"mlx/mlx/distributed/jaccl/utils.cpp",
"mlx/mlx/distributed/ring/no_ring.cpp",
"mlx/mlx/distributed/jaccl/no_jaccl.cpp",
],
cSettings: [
.headerSearchPath("mlx"),
Expand Down Expand Up @@ -302,10 +297,20 @@ let package = Package(
.testTarget(
name: "MLXTests",
dependencies: [
"MLX", "MLXNN", "MLXOptimizers",
"MLX", "MLXNN", "MLXOptimizers", "DistributedWorker",
]
),

// ------
// Test support executables

.executableTarget(
name: "DistributedWorker",
dependencies: ["MLX", "MLXNN"],
path: "Tests/DistributedTestSupport",
sources: ["DistributedWorkerMain.swift", "DistributedWorkerOperations.swift"]
),

// ------
// Example programs

Expand Down
Loading