Add distributed communication framework for multi-device tensor parallelism

[ronaldmannak]

Proposed changes

Ports MLX's distributed communication framework to Swift, enabling multi-device model inference and training across Apple Silicon nodes connected via Ethernet (ring/TCP) or Thunderbolt 5 (JACCL/RDMA). All C/C++ code was already vendored in this repository but excluded from compilation.

What's included

Package.swift

• Un-excludes vendored distributed C/C++ sources (ring and JACCL backends)
• Adds DistributedWorker helper executable target for multi-process testing

Swift Bindings (Source/MLX/Distributed.swift)

• DistributedGroup class wrapping mlx_distributed_group (rank, size, split)
• MLXDistributed enum with 8 collective operations: allSum, allGather, allMax, allMin, sumScatter, send, recv, recvLike
• Follows established namespace pattern (MLXRandom, MLXFFT)

Distributed NN Layers (Source/MLXNN/Distributed.swift)

• AllToShardedLinear / ShardedToAllLinear for column-parallel and row-parallel tensor sharding
• QuantizedAllToShardedLinear / QuantizedShardedToAllLinear with full Quantized protocol conformance
• shardLinear / shardInPlace utilities with segments parameter for fused QKV weights
• averageGradients with batched allReduce, communicationType for bandwidth-reducing cast-on-wire, and mixed-dtype fallback
• sumGradients helper using CustomFunction for identity-forward / allSum-backward VJP

Skill documentation (skills/mlx-distributed/)

• Complete SKILL.md with architecture overview, quick start, 4 prioritized workflows, and best practices
• 5 reference docs: primitives, NN layers, sharding, gradient averaging, multi-process setup

Known upstream limitations

Limitation	Impact
MLX-C doesn't expose backend selection parameter	Cannot programmatically choose ring vs JACCL; priority order used. See ml-explore/mlx-c#108
mlx_distributed_group_free() not in public C API	Group deallocation relies on C++ shared_ptr ref counting
group.split() unsupported by ring and JACCL	Subgroup creation requires MPI backend (not available on macOS)
reduceScatter not implemented in ring backend	sumScatter only testable for graceful error handling
All distributed ops CPU-only	Must use Device.withDefaultDevice(.cpu) in distributed code paths

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[davidkoski]

@ronaldmannak it looks like a lint issue. 0.31.1 is merged to main now so the API you were looking for on the mlx-c side is available

[ronaldmannak]

@davidkoski fixed and updated for 0.31.1

[davidkoski]

I haven't gone over the code in detail yet. CI is failing with this:

Test Case '-[MLXTests.DistributedNNTests testMultiProcessAverageGradients]' started.
/Users/runner/actions-runner/_work/mlx-swift/mlx-swift/Tests/MLXTests/DistributedNNTests.swift:1497: error: -[MLXTests.DistributedNNTests testMultiProcessAverageGradients] : failed - DistributedWorker binary not found. Build with: xcodebuild build -scheme mlx-swift-Package

Do we need to adjust how the tests are run? Are these tests appropriate for CI, or are they more integration tests that should be run on developer machines?

[ronaldmannak]

@davidkoski These tests use the DistributedWorker helper executable to simulate another distributed rank. That works locally, but not in the default GitHub Actions lane. I can skip the DistributedWorker-backed multi-process tests when GITHUB_ACTIONS=true. Local runs are unchanged. Would that work for you?

[davidkoski]

@davidkoski These tests use the DistributedWorker helper executable to simulate another distributed rank. That works locally, but not in the default GitHub Actions lane. I can skip the DistributedWorker-backed multi-process tests when GITHUB_ACTIONS=true. Local runs are unchanged. Would that work for you?

That could work, but if it runs locally I wonder what is missing from the CI worker? So this other program is just a local (but not-the-same-process) worker that it should find?

I messed around with it a bit locally and found if I did a "project clean" followed by a running e.g. testMultiProcessSendRecv it failed in the same way. I think the problem is that the tests do not have a dependency on DistributedWorker, e.g. something like this:

        .testTarget(
            name: "MLXTests",
            dependencies: [
                "MLX", "MLXNN", "MLXOptimizers",
                "DistributedWorker",
            ]
        ),

I wonder it should be flipped around: these run by default on CI but not locally? There are a lot of tests and they are pretty slow since they have to fire up subprocesses and wait for them (well, not all of them).

Also, the .xcodeproj version doesn't pass these tests or build the helper tool. I think it would be fine to skip the test in the xcodeproj variant (as long as we have coverage from the Package.swift side).

[davidkoski]

This matches the tests, but is this how real applications would work? I wonder if this is too specific to the tests.

[davidkoski]

again this may be too test specific

[davidkoski]

Perhaps this is generally true: distributed compute can have latency and failures, that is the nature of distributed computing. This particular warning is good for the limitations of these backends, but perhaps the bigger item is explaining that callers must use error handling and deal with errors.

[davidkoski]

A bit surprising to see fputs here, but it looks like printing to stderr isn't a common use case (there is no idiomatic way to do so that I see).

[davidkoski]

Really? This seems surprising -- the python variant supports GPU. I wonder if parts of the pipeline only support CPU?

[davidkoski]

This looks like a critical missing piece.

[ronaldmannak]

Yes, MLX-C does not currently expose a public mlx_distributed_group deinit/destroy.

I can open an issue in mlx-c if you want to

[davidkoski]

I talked to Ronan.

[davidkoski]

This function is awkward to call:

        let group = MLXDistributed.`init`()!

Why not just use DistributedGroup directly? It can have a no-arg failable init that does this. If not, I think this function should be renamed to be easier to call.

[ronaldmannak]

I agree. I tried to keep as close to the Python code as possible, but that made the Swift code ugly and harder to use than necessary. I'm refactoring the code now.

[davidkoski]

This pattern is repeated over and over -- might be good to refactor into a function.

[davidkoski]

This is also a lot of repeated code. Maybe use a variant on this:

func assertEqual(
    _ array1: MLXArray, _ array2: MLXArray, rtol: Double = 1e-5, atol: Double = 1e-8,
    file: StaticString = #filePath, line: UInt = #line
) {
    XCTAssertEqual(array1.shape, array2.shape, "shapes differ: \(array1.shape) != \(array2.shape)")
    XCTAssertTrue(
        array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self),
        "contents differ:\n\(array1)\n\(array2)")
}

[davidkoski]

A thought on the error handling -- previously I commented that it is important for distributed operations to use error handling because they are all failable. I wonder if this should be throws and we do the withError internally for all of these?

e.g. the IO routines that can also fail:

public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [String: MLXArray] {
    precondition(url.isFileURL)
    let path = url.path(percentEncoded: false)

    switch url.pathExtension {
    case "safetensors":
        var r0 = mlx_map_string_to_array_new()
        var r1 = mlx_map_string_to_string_new()
        defer { mlx_map_string_to_array_free(r0) }
        defer { mlx_map_string_to_string_free(r1) }

        _ = try withError {
            mlx_load_safetensors(&r0, &r1, path.cString(using: .utf8), stream.ctx)
        }

        return mlx_map_array_values(r0)

Making these throw forces callers to handle the error or use try! if they don't care.

[ronaldmannak]

@davidkoski Great suggestion. I applied that pattern to the APIs with a meaningful synchronous failure boundary, but I left allSum non-throwing because it is still a lazy op.
For allSum, the important distributed/runtime failures show up when the returned MLXArray is evaluated, not when allSum is called. There is a narrow synchronous failure surface during lazy-op construction (invalid/empty MLXArray), but it didn’t feel strong enough to justify a throws on the primary lazy collective API, especially since callers still need withError { ... } / checkedEval(...) at the evaluation boundary. Lmk if you think we should make allSum throw.

[davidkoski]

Ah, interesting point -- yes, lazy distributed ops are a bit mind blowing :-)

I guess checkedEval() might be a good fit then.

[davidkoski]

Perhaps use:

class Cache<Key: Hashable, Element>: @unchecked (Sendable) {

It is only used in one place, but it does at least have an LRU(-ish) replacement policy like the python one has.