diff --git a/Package.swift b/Package.swift index 80f9d282..b54791ba 100644 --- a/Package.swift +++ b/Package.swift @@ -1,4 +1,4 @@ -// swift-tools-version: 5.12 +// swift-tools-version: 6.1 // The swift-tools-version declares the minimum version of Swift required to build this package. // Copyright © 2024 Apple Inc. @@ -160,6 +160,7 @@ let noCudaCmlxExcludes = [ .headerSearchPath("mlx/mlx/distributed/jaccl/lib"), .define("MLX_USE_ACCELERATE"), + .define("MLX_USE_JACCL", .when(traits: ["JACCL"])), .define("ACCELERATE_NEW_LAPACK"), .define("_METAL_"), .define("SWIFTPM_BUNDLE", to: "\"mlx-swift_Cmlx\""), @@ -229,7 +230,8 @@ let cmlx = Target.target( "mlx/mlx/backend/metal/nojit_kernels.cpp", // distributed backends: enable ring, disable MPI + NCCL. JACCL is - // built through mlx-conditional on macOS, so its sources are excluded here. + // built through mlx-conditional on macOS when the JACCL trait is + // enabled, so its sources are excluded here. "mlx/mlx/distributed/mpi/mpi.cpp", "mlx/mlx/distributed/ring/no_ring.cpp", "mlx/mlx/distributed/nccl/nccl.cpp", @@ -273,6 +275,14 @@ let package = Package( .library(name: "MLXLinalg", targets: ["MLXLinalg"]), .library(name: "MLXFast", targets: ["MLXFast"]), ], + traits: [ + .trait( + name: "JACCL", + description: + "Builds the JACCL distributed backend (RDMA over Thunderbolt 5) on macOS. Disable it to leave the backend out of an app that does not run distributed MLX." + ), + .default(enabledTraits: ["JACCL"]), + ], dependencies: [ // for Complex type .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0") @@ -391,6 +401,7 @@ let package = Package( ] ), ] + cudaTargets, + swiftLanguageModes: [.v5], cxxLanguageStandard: .gnucxx20 ) diff --git a/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp b/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp index 8e8afea5..3ff59ce0 100644 --- a/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp +++ b/Source/Cmlx/mlx-conditional/jaccl_conditional.cpp @@ -4,9 +4,10 @@ // JACCL (RDMA over Thunderbolt 5) needs , which ships with // the macOS 26.2 SDK and later. Build the backend and its library when the -// header is present, and the stub otherwise. librdma is loaded at runtime, so -// on an older macOS MLXDistributed.isAvailable(.jaccl) reports false. -#if TARGET_OS_OSX && __has_include() +// JACCL package trait is enabled and the header is present, and the stub +// otherwise. librdma is loaded at runtime, so on an older macOS +// MLXDistributed.isAvailable(.jaccl) reports false. +#if TARGET_OS_OSX && defined(MLX_USE_JACCL) && __has_include() #include "../mlx/mlx/distributed/jaccl/jaccl.cpp" #include "../mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp" #include "../mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp" diff --git a/Source/MLX/Distributed.swift b/Source/MLX/Distributed.swift index 47947636..33a12f3f 100644 --- a/Source/MLX/Distributed.swift +++ b/Source/MLX/Distributed.swift @@ -7,8 +7,9 @@ import Foundation /// /// This mirrors `mlx.core.distributed` in the Python API. /// -/// Forming a group can fail as soon as it is attempted, so -/// ``initialize(backend:strict:)`` and ``Group/split(color:key:)`` throw. +/// Forming a group, splitting it, and reading its rank or size can fail, so +/// ``initialize(backend:strict:)``, ``Group/split(color:key:)``, ``Group/rank``, +/// and ``Group/size`` throw. /// /// The collectives and the point to point operations are lazy like every /// other MLX operation. Invalid arguments -- e.g. sending in a group of size @@ -48,9 +49,28 @@ public enum MLXDistributed { case ring case mpi case nccl + /// RDMA over Thunderbolt 5, on macOS only. Built when the `JACCL` + /// package trait is enabled, which is the default. case jaccl } + /// A synchronous byte-level all-gather for JACCL initialization metadata. + /// + /// `source` contains this process's contribution. Fill `destination`, whose + /// byte count is `size * source.count`, with all contributions in rank order. + /// The contribution from rank `r` starts at byte `r * source.count`. + /// + /// Both buffers are borrowed for the duration of the call. Do not retain + /// their pointers or return before the destination is filled. The callback + /// may run outside the main actor; captured transport state must be safe to + /// use from the calling thread. A thrown error fails initialization and is + /// reported as an MLX error. + public typealias MetadataAllGather = + @Sendable ( + _ rank: Int, _ size: Int, _ source: UnsafeRawBufferPointer, + _ destination: UnsafeMutableRawBufferPointer + ) throws -> Void + /// Check if a communication backend is available. /// /// - Parameter backend: the backend to check for, defaulting to ``Backend/any`` @@ -75,17 +95,29 @@ public enum MLXDistributed { } /// The rank of this process in the group. + /// + /// Throws if the underlying group cannot provide its rank. public var rank: Int { - var result: Int32 = 0 - mlx_distributed_group_get_rank(&result, ctx) - return Int(result) + get throws { + try withError { + var result: Int32 = 0 + _ = mlx_distributed_group_get_rank(&result, ctx) + return Int(result) + } + } } /// The number of processes in the group. + /// + /// Throws if the underlying group cannot provide its size. public var size: Int { - var result: Int32 = 0 - mlx_distributed_group_get_size(&result, ctx) - return Int(result) + get throws { + try withError { + var result: Int32 = 0 + _ = mlx_distributed_group_get_size(&result, ctx) + return Int(result) + } + } } /// Split the group into subgroups based on the provided color. @@ -128,12 +160,12 @@ public enum MLXDistributed { /// backend is configured but cannot form its group, e.g. the ring backend /// with a malformed `MLX_HOSTFILE`. /// - /// MLX caches the group for each backend for the life of the process, + /// MLX caches the group for each backend until ``clearBackends()`` is called, /// including the group of size one that a non-strict call falls back to, /// and a later strict call for that backend returns the cached group /// instead of throwing. A non-strict ``Backend/any`` that finds no backend - /// caches its fallback under ``Backend/jaccl``. Initialize strictly - /// before anything else initializes the backend. + /// caches its fallback under ``Backend/jaccl``. Initialize strictly before + /// anything else initializes the backend, or clear the cache before retrying. /// /// - Parameters: /// - backend: the backend to use, defaulting to ``Backend/any`` @@ -154,6 +186,56 @@ public enum MLXDistributed { return Group(result) } + /// Initialize JACCL with a custom all-gather for initialization metadata. + /// + /// The callback runs synchronously with borrowed buffers, as described by + /// ``MetadataAllGather``. The backend validates whether it supports this + /// initialization method. + /// + /// If the backend already has a cached group, that group is returned without + /// calling `allGather`, and the new callback is released before this function + /// returns. Call ``clearBackends()`` first to initialize with a new callback. + /// Otherwise, the callback may be retained by the cache and live groups. + /// Clearing the cache does not release callbacks still held by live groups. + /// The callback is also released if initialization fails. + /// + /// - Parameters: + /// - backend: the backend to use, defaulting to ``Backend/jaccl`` + /// - strict: if `true`, throw rather than fall back to a group of size one + /// when no backend can be initialized + /// - allGather: the synchronous metadata all-gather callback + /// - Returns: the global group + @discardableResult + public static func initialize( + backend: Backend = .jaccl, strict: Bool = false, + allGather: @escaping MetadataAllGather + ) throws -> Group { + var result = mlx_distributed_group_new() + do { + try withError { + let bridge = makeMetadataAllGatherBridge(allGather) + // C takes ownership on entry, including failure and cache hits. + _ = mlx_distributed_init_with_all_gather( + &result, strict, backend.rawValue, bridge.callback, + bridge.payload, bridge.dtor) + } + } catch { + mlx_distributed_group_free(result) + throw error + } + return Group(result) + } + + /// Clear the distributed backend cache. + /// + /// Existing groups remain valid. Callbacks retained by those groups remain + /// alive until the groups release them. + public static func clearBackends() throws { + try withError { + _ = mlx_distributed_clear_backends() + } + } + // MARK: - Collectives // Note: mlx-c has no representation for an "unspecified" stream -- @@ -295,6 +377,58 @@ public enum MLXDistributed { } } +// The caller must transfer the payload to C or release it with the destructor. +func makeMetadataAllGatherBridge(_ allGather: @escaping MLXDistributed.MetadataAllGather) -> ( + callback: mlx_distributed_all_gather_fn, payload: UnsafeMutableRawPointer, + dtor: @convention(c) (UnsafeMutableRawPointer?) -> Void +) { + let payload = Unmanaged.passRetained(MetadataAllGatherCaptureState(allGather)).toOpaque() + return (metadataAllGatherTrampoline, payload, freeMetadataAllGather) +} + +private final class MetadataAllGatherCaptureState: Sendable { + let allGather: MLXDistributed.MetadataAllGather + + init(_ allGather: @escaping MLXDistributed.MetadataAllGather) { + self.allGather = allGather + } +} + +private func freeMetadataAllGather(_ payload: UnsafeMutableRawPointer?) { + if let payload { + Unmanaged.fromOpaque(payload).release() + } +} + +private func metadataAllGatherTrampoline( + rank: Int32, size: Int32, source: UnsafePointer?, + destination: UnsafeMutablePointer?, byteCount: Int, + payload: UnsafeMutableRawPointer? +) -> Int32 { + guard let payload, size > 0, rank >= 0, rank < size, byteCount >= 0 else { + return 1 + } + let (destinationCount, overflow) = Int(size).multipliedReportingOverflow(by: byteCount) + guard !overflow, + byteCount == 0 || source != nil, + destinationCount == 0 || destination != nil + else { + return 1 + } + + let state = Unmanaged.fromOpaque(payload).takeUnretainedValue() + do { + try state.allGather( + Int(rank), Int(size), + UnsafeRawBufferPointer(start: source.map { UnsafeRawPointer($0) }, count: byteCount), + UnsafeMutableRawBufferPointer( + start: destination.map { UnsafeMutableRawPointer($0) }, count: destinationCount)) + return 0 + } catch { + return 1 + } +} + extension Optional where Wrapped == MLXDistributed.Group { /// The underlying group handle, or a null handle meaning "the global group". fileprivate var groupCtx: mlx_distributed_group { diff --git a/Source/MLX/Documentation.docc/install.md b/Source/MLX/Documentation.docc/install.md index 4184b694..f8877a0c 100644 --- a/Source/MLX/Documentation.docc/install.md +++ b/Source/MLX/Documentation.docc/install.md @@ -32,6 +32,22 @@ dependencies: [.product(name: "MLX", package: "mlx-swift"), > Note that SwiftPM cannot build the metal shaders so the ultimate build has to be done via Xcode or `xcodebuild`. +## Package Traits + +`JACCL` builds the JACCL distributed backend (RDMA over Thunderbolt 5) on +macOS. It is enabled by default. iOS, tvOS and visionOS never build it. + +An app that does not run distributed MLX can leave the backend out of its +binary by disabling the default traits: + +```swift +.package(url: "https://github.com/ml-explore/mlx-swift", from: "0.33.0", traits: []) +``` + +SwiftPM enables the union of the traits that all dependents request. If another +package in the graph depends on `mlx-swift` with the default traits, `JACCL` +stays enabled. + ## Command Line Tools See (Running From Command Line) for information about running command line tools from the shell. diff --git a/Source/MLX/ErrorHandler.swift b/Source/MLX/ErrorHandler.swift index 4d0daab8..555dedd1 100644 --- a/Source/MLX/ErrorHandler.swift +++ b/Source/MLX/ErrorHandler.swift @@ -269,8 +269,8 @@ public final class ErrorBox: @unchecked Sendable { /// Throw the ``firstError`` if set, otherwise do nothing. public func check() throws { - if let _firstError { - throw _firstError + if let firstError { + throw firstError } } } diff --git a/Source/MLXNN/Distributed.swift b/Source/MLXNN/Distributed.swift index f401b09d..b000b8cc 100644 --- a/Source/MLXNN/Distributed.swift +++ b/Source/MLXNN/Distributed.swift @@ -90,9 +90,12 @@ public enum ShardingError: Error, CustomStringConvertible, Equatable { /// Returns a function that is the identity in the forward pass and sums the /// gradients across the group in the backward pass. /// +/// Throws if the group size cannot be read. The returned function is nonthrowing. +/// /// - Parameter group: the group to sum across -public func sumGradients(group: MLXDistributed.Group) -> (MLXArray) -> MLXArray { - if group.size == 1 { +public func sumGradients(group: MLXDistributed.Group) throws -> (MLXArray) -> MLXArray { + let size = try group.size + if size == 1 { return { $0 } } @@ -149,8 +152,8 @@ private func shard( _ parameters: ModuleParameters, group: MLXDistributed.Group, _ sharding: ShardingPredicate ) throws -> ModuleParameters { - let size = group.size - let rank = group.rank + let size = try group.size + let rank = try group.rank let sharded = try parameters.flattened().map { path, weight -> (String, MLXArray) in // a scalar -- the NVFP4 global scale, for example -- is the same in @@ -257,8 +260,9 @@ public func shardInPlace( let group = try group ?? MLXDistributed.initialize() let sharded: ShardingPredicate if case .shardedToAll = sharding, case .indices = segments { + let size = try group.size sharded = try shardedToAllPredicate( - segments, quantizedLayersIn: module, size: group.size) + segments, quantizedLayersIn: module, size: size) } else { sharded = predicate(for: sharding, segments: segments) } @@ -344,25 +348,26 @@ open class AllToShardedLinear: Linear { group: MLXDistributed.Group? = nil ) throws { let group = try group ?? MLXDistributed.initialize() - guard outputDimensions % group.size == 0 else { + let size = try group.size + guard outputDimensions % size == 0 else { throw ShardingError.indivisible( - dimension: "output", of: outputDimensions, across: group.size) + dimension: "output", of: outputDimensions, across: size) } self.group = group - self.aggregateGradients = sumGradients(group: group) + self.aggregateGradients = try sumGradients(group: group) let scale = sqrt(1.0 / Float(inputDimensions)) - let shardedOutput = outputDimensions / group.size + let shardedOutput = outputDimensions / size super.init( weight: MLXRandom.uniform(-scale ..< scale, [shardedOutput, inputDimensions]), bias: bias ? MLXRandom.uniform(-scale ..< scale, [shardedOutput]) : nil) } /// Hold parameters that are already this process' shard. - init(shardedWeight weight: MLXArray, bias: MLXArray?, group: MLXDistributed.Group) { + init(shardedWeight weight: MLXArray, bias: MLXArray?, group: MLXDistributed.Group) throws { self.group = group - self.aggregateGradients = sumGradients(group: group) + self.aggregateGradients = try sumGradients(group: group) super.init(weight: weight, bias: bias) } @@ -380,9 +385,10 @@ open class AllToShardedLinear: Linear { let group = try group ?? MLXDistributed.initialize() let (outputDimensions, _) = other.shape - guard outputDimensions % group.size == 0 else { + let size = try group.size + guard outputDimensions % size == 0 else { throw ShardingError.indivisible( - dimension: "output", of: outputDimensions, across: group.size) + dimension: "output", of: outputDimensions, across: size) } let parameters = try Dictionary( @@ -393,7 +399,7 @@ open class AllToShardedLinear: Linear { throw ShardingError.missingParameter("weight") } - self.init(shardedWeight: weight, bias: parameters["bias"], group: group) + try self.init(shardedWeight: weight, bias: parameters["bias"], group: group) } open override func callAsFunction(_ x: MLXArray) -> MLXArray { @@ -414,7 +420,7 @@ open class AllToShardedLinear: Linear { public override func toQuantized(groupSize: Int, bits: Int, mode: QuantizationMode) -> Module { QuantizedAllToShardedLinear( shardedWeight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode, - group: group) + group: group, aggregateGradients: aggregateGradients) } } @@ -439,9 +445,10 @@ open class ShardedToAllLinear: Linear { group: MLXDistributed.Group? = nil ) throws { let group = try group ?? MLXDistributed.initialize() - guard inputDimensions % group.size == 0 else { + let size = try group.size + guard inputDimensions % size == 0 else { throw ShardingError.indivisible( - dimension: "input", of: inputDimensions, across: group.size) + dimension: "input", of: inputDimensions, across: size) } self.group = group @@ -455,7 +462,7 @@ open class ShardedToAllLinear: Linear { let scale = sqrt(1.0 / Float(inputDimensions)) super.init( weight: MLXRandom.uniform( - -scale ..< scale, [outputDimensions, inputDimensions / group.size]), + -scale ..< scale, [outputDimensions, inputDimensions / size]), bias: bias ? MLXArray.zeros([outputDimensions]) : nil) } @@ -479,9 +486,10 @@ open class ShardedToAllLinear: Linear { let group = try group ?? MLXDistributed.initialize() let (_, inputDimensions) = other.shape - guard inputDimensions % group.size == 0 else { + let size = try group.size + guard inputDimensions % size == 0 else { throw ShardingError.indivisible( - dimension: "input", of: inputDimensions, across: group.size) + dimension: "input", of: inputDimensions, across: size) } let parameters = try Dictionary( @@ -526,12 +534,25 @@ open class QuantizedAllToShardedLinear: QuantizedLinear { private let aggregateGradients: (MLXArray) -> MLXArray /// Quantize a weight that is already this process' shard. + /// + /// Throws if the group size cannot be read. public init( shardedWeight weight: MLXArray, bias: MLXArray?, groupSize: Int = 64, bits: Int = 4, mode: QuantizationMode = .affine, group: MLXDistributed.Group + ) throws { + self.group = group + self.aggregateGradients = try sumGradients(group: group) + super.init(weight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode) + } + + /// Reuse the aggregation validated when the source layer was initialized. + fileprivate init( + shardedWeight weight: MLXArray, bias: MLXArray?, groupSize: Int, bits: Int, + mode: QuantizationMode, group: MLXDistributed.Group, + aggregateGradients: @escaping (MLXArray) -> MLXArray ) { self.group = group - self.aggregateGradients = sumGradients(group: group) + self.aggregateGradients = aggregateGradients super.init(weight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode) } @@ -540,9 +561,9 @@ open class QuantizedAllToShardedLinear: QuantizedLinear { shardedQuantizedWeight weight: MLXArray, bias: MLXArray?, scales: MLXArray, biases: MLXArray?, groupSize: Int, bits: Int, mode: QuantizationMode, globalScale: MLXArray?, group: MLXDistributed.Group - ) { + ) throws { self.group = group - self.aggregateGradients = sumGradients(group: group) + self.aggregateGradients = try sumGradients(group: group) super.init( weight: weight, bias: bias, scales: scales, biases: biases, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale) @@ -559,9 +580,10 @@ open class QuantizedAllToShardedLinear: QuantizedLinear { ) throws { let group = try group ?? MLXDistributed.initialize() let (outputDimensions, _) = other.shape - guard outputDimensions % group.size == 0 else { + let size = try group.size + guard outputDimensions % size == 0 else { throw ShardingError.indivisible( - dimension: "output", of: outputDimensions, across: group.size) + dimension: "output", of: outputDimensions, across: size) } let parameters = try Dictionary( @@ -575,7 +597,7 @@ open class QuantizedAllToShardedLinear: QuantizedLinear { throw ShardingError.missingParameter("scales") } - self.init( + try self.init( shardedQuantizedWeight: weight, bias: parameters["bias"], scales: scales, biases: parameters["biases"], groupSize: other.groupSize, bits: other.bits, mode: other.mode, globalScale: parameters["global_scale"], group: group) @@ -628,9 +650,10 @@ open class QuantizedShardedToAllLinear: QuantizedLinear { ) throws { let group = try group ?? MLXDistributed.initialize() let (_, inputDimensions) = other.shape - guard inputDimensions % group.size == 0 else { + let size = try group.size + guard inputDimensions % size == 0 else { throw ShardingError.indivisible( - dimension: "input", of: inputDimensions, across: group.size) + dimension: "input", of: inputDimensions, across: size) } let parameters = try Dictionary( @@ -638,7 +661,7 @@ open class QuantizedShardedToAllLinear: QuantizedLinear { other.parameters(), group: group, quantizedShardedToAllPredicate( segments, inputDimensions: inputDimensions, groupSize: other.groupSize, - size: group.size) + size: size) ).flattened()) guard let weight = parameters["weight"] else { throw ShardingError.missingParameter("weight") @@ -686,7 +709,7 @@ public func averageGradients( allReduceSize: Int = 32 * 1024 * 1024, stream: StreamOrDevice = .cpu ) throws -> ModuleParameters { let group = try group ?? MLXDistributed.initialize() - let size = group.size + let size = try group.size if size == 1 { return gradients } diff --git a/Tests/MLXTests/DistributedHarness.swift b/Tests/MLXTests/DistributedHarness.swift index 4b7692dc..76eca92b 100644 --- a/Tests/MLXTests/DistributedHarness.swift +++ b/Tests/MLXTests/DistributedHarness.swift @@ -117,8 +117,10 @@ enum DistributedHarness { // a group of one makes every assertion about sharding pass without // proving anything, so the size is part of the contract - print("[rank \(group.rank)] joined a group of size \(group.size)") - XCTAssertEqual(group.size, ranks, "joined a group of the wrong size") + let rank = try group.rank + let size = try group.size + print("[rank \(rank)] joined a group of size \(size)") + XCTAssertEqual(size, ranks, "joined a group of the wrong size") try withError { try body(group) diff --git a/Tests/MLXTests/DistributedNNTests.swift b/Tests/MLXTests/DistributedNNTests.swift index 861418bd..9d2ac861 100644 --- a/Tests/MLXTests/DistributedNNTests.swift +++ b/Tests/MLXTests/DistributedNNTests.swift @@ -16,10 +16,12 @@ import XCTest /// /// - Parameter world: the group to shard across func shardLinearBody(world: MLXDistributed.Group) throws { + let rank = try world.rank + let size = try world.size MLXRandom.seed(0xF0F0_F0F0) - let lower = world.rank * 1024 / world.size - let upper = (world.rank + 1) * 1024 / world.size + let lower = rank * 1024 / size + let upper = (rank + 1) * 1024 / size let x = MLXRandom.normal([4, 1024]) let linear = Linear(1024, 1024, bias: true) @@ -109,6 +111,8 @@ func shardLinearBody(world: MLXDistributed.Group) throws { /// The gradients of a sharded model must match the slice of the unsharded /// model's gradients that this rank owns. private func shardLinearBackward(world: MLXDistributed.Group) throws { + let rank = try world.rank + let size = try world.size MLXRandom.seed(0xF0F0_F0F0) let layers = (0 ..< 4).map { _ in Linear(128, 128) } @@ -130,8 +134,8 @@ private func shardLinearBackward(world: MLXDistributed.Group) throws { let (l2, g2) = valueAndGrad(model: sharded, loss)(sharded, x, y) try checkedEval(l1, g1, l2, g2) - let lower = world.rank * 128 / world.size - let upper = (world.rank + 1) * 128 / world.size + let lower = rank * 128 / size + let upper = (rank + 1) * 128 / size let full = Dictionary(uniqueKeysWithValues: g1.flattened()) let shard = Dictionary(uniqueKeysWithValues: g2.flattened()) @@ -293,6 +297,7 @@ func donationBody(world: MLXDistributed.Group) throws { /// the sharded layers have to come out with the shapes the layer they were /// built from would have, divided along the sharded axis. func quantizedShardedConstructionBody(world: MLXDistributed.Group) throws { + let size = try world.size for bits in [2, 3, 4, 5, 6, 8] { let quantized = QuantizedLinear(Linear(1536, 1024), groupSize: 64, bits: bits) let allToSharded = try QuantizedAllToShardedLinear(quantized, group: world) @@ -300,12 +305,12 @@ func quantizedShardedConstructionBody(world: MLXDistributed.Group) throws { // rows are outputs, columns are the packed inputs XCTAssertEqual( - allToSharded.weight.dim(0), quantized.weight.dim(0) / world.size, "bits \(bits)") + allToSharded.weight.dim(0), quantized.weight.dim(0) / size, "bits \(bits)") XCTAssertEqual(allToSharded.weight.dim(1), quantized.weight.dim(1), "bits \(bits)") XCTAssertEqual(shardedToAll.weight.dim(0), quantized.weight.dim(0), "bits \(bits)") XCTAssertEqual( - shardedToAll.weight.dim(1), quantized.weight.dim(1) / world.size, "bits \(bits)") + shardedToAll.weight.dim(1), quantized.weight.dim(1) / size, "bits \(bits)") // and the scales have to cover exactly the inputs the weight holds for layer in [allToSharded, shardedToAll] { @@ -318,18 +323,20 @@ func quantizedShardedConstructionBody(world: MLXDistributed.Group) throws { /// Cases the Python tests do not cover, each one a bug this port had. func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { + let rank = try world.rank + let size = try world.size // A sharded-to-all layer adds its bias after the reduction, so every // process has to hold the same bias. Seed the processes differently: // nothing keeps their generators in step -- they may have done different // work before this -- and the bias must not depend on that. - MLXRandom.seed(UInt64(world.rank + 1)) - let layer = try ShardedToAllLinear(8 * world.size, 4, group: world) + MLXRandom.seed(UInt64(rank + 1)) + let layer = try ShardedToAllLinear(8 * size, 4, group: world) let bias = try XCTUnwrap(layer.bias) let gathered = MLXDistributed.allGather(bias, group: world) try checkedEval(gathered) let width = bias.dim(0) - for rank in 1 ..< world.size { + for rank in 1 ..< size { XCTAssertTrue( gathered[0 ..< width].allClose(gathered[rank * width ..< (rank + 1) * width]) .item(Bool.self), @@ -338,15 +345,15 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { // A dimension can divide by the group size while one of its segments does // not. That has to be reported rather than trap inside split(parts:). - if world.size > 1 { + if size > 1 { XCTAssertThrowsError( try shardLinear( - Linear(8 * world.size, 8), sharding: .shardedToAll, + Linear(8 * size, 8), sharding: .shardedToAll, segments: .indices([3]), group: world) ) { error in XCTAssertEqual( error as? ShardingError, - .indivisible(dimension: "segment of weight", of: 3, across: world.size)) + .indivisible(dimension: "segment of weight", of: 3, across: size)) } } @@ -354,11 +361,11 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { // group of one as well. split(parts:) would call fatalError. XCTAssertThrowsError( try shardLinear( - Linear(8 * world.size, 8), sharding: .shardedToAll, segments: .count(3), + Linear(8 * size, 8), sharding: .shardedToAll, segments: .count(3), group: world) ) { error in XCTAssertEqual( - error as? ShardingError, .invalidSegments(.count(3), dimension: 8 * world.size)) + error as? ShardingError, .invalidSegments(.count(3), dimension: 8 * size)) } // A QuantizedLinear is a Linear, so the float layers accept one as far as @@ -385,7 +392,6 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { wide, segments: .indices([boundary]), group: world) // this rank's part of each segment - let (rank, size) = (world.rank, world.size) let first = (rank * boundary / size) ..< ((rank + 1) * boundary / size) let second = (boundary + rank * (inputs - boundary) / size) @@ -406,7 +412,7 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { ) { error in XCTAssertEqual( error as? ShardingError, - .quantizationGroup(inputDimensions: 272, groupSize: 64, across: world.size)) + .quantizationGroup(inputDimensions: 272, groupSize: 64, across: size)) } } @@ -417,17 +423,19 @@ func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { /// instead. Every batching has to produce the same average, and the average /// has to include every rank. func averageGradientsBody(world: MLXDistributed.Group) throws { + let rank = try world.rank + let size = try world.size // Ten gradients of ten float32 values, 40 bytes each as in Python, so the // limits below form different batches. Each rank contributes its own // rank + 1, so the average is (size + 1) / 2 times the gradient's own // factor. The values differ per gradient and half are matrices on // purpose: batching concatenates them, reduces once and splits the result // again, and identical or flat arrays would hide a bad split. - let mean = Float(world.size + 1) / 2 + let mean = Float(size + 1) / 2 let shapes = (0 ..< 10).map { $0.isMultiple(of: 2) ? [10] : [2, 5] } let gradients = ModuleParameters.unflattened( (0 ..< 10).map { - ("g\($0)", MLXArray.ones(shapes[$0]) * Float(($0 + 1) * (world.rank + 1))) + ("g\($0)", MLXArray.ones(shapes[$0]) * Float(($0 + 1) * (rank + 1))) }) let expected = Dictionary( uniqueKeysWithValues: (0 ..< 10).map { @@ -461,8 +469,8 @@ func averageGradientsBody(world: MLXDistributed.Group) throws { // one reduction per array and keep their dtype let mixed = try averageGradients( ModuleParameters.unflattened([ - ("wide", MLXArray.ones([4]) * Float(world.rank + 1)), - ("narrow", (MLXArray.ones([4]) * Float(world.rank + 1)).asType(.float16)), + ("wide", MLXArray.ones([4]) * Float(rank + 1)), + ("narrow", (MLXArray.ones([4]) * Float(rank + 1)).asType(.float16)), ]), group: world) // evaluate in the flattened order, which is sorted by key: a Dictionary // iterates in a per process order, and evaluating the collectives in @@ -476,13 +484,14 @@ func averageGradientsBody(world: MLXDistributed.Group) throws { let averaged = try XCTUnwrap(flatMixed[key]) XCTAssertTrue( averaged.asType(.float32) - .allClose(MLXArray.ones([4]) * (Float(world.size + 1) / 2), rtol: 1e-2) + .allClose(MLXArray.ones([4]) * (Float(size + 1) / 2), rtol: 1e-2) .item(Bool.self), key) } } /// Port of `test_clip_grad_norm_sharded`. func clipGradNormShardedBody(world: MLXDistributed.Group) throws { + let size = try world.size let value: Float = 3 let gradients = ModuleParameters.unflattened([ ("a", MLXArray.ones([4, 3]) * value), @@ -492,7 +501,7 @@ func clipGradNormShardedBody(world: MLXDistributed.Group) throws { // every rank holds the same number of elements, so the global norm counts // all of them let localCount = 4 * 3 + 5 - let expectedNorm = (Float(world.size * localCount)).squareRoot() * value + let expectedNorm = (Float(size * localCount)).squareRoot() * value // a limit far above the norm leaves the shard alone let (unclipped, norm) = try clipGradNormSharded( @@ -588,6 +597,8 @@ private func assertGradientWhole( /// layer. The positions count input features, which the packed weight and the /// scales hold at different resolutions, so each has to split at its own index. func quantizedShardInPlaceBody(world: MLXDistributed.Group) throws { + let rank = try world.rank + let size = try world.size MLXRandom.seed(0) let x = MLXRandom.normal([4, 1024]) let model = Sequential(layers: QuantizedLinear(Linear(1024, 32, bias: false))) @@ -597,8 +608,8 @@ func quantizedShardInPlaceBody(world: MLXDistributed.Group) throws { // each process holds its slice of both segments; the sum of the partial // results is the unsharded result - let step = 512 / world.size - let lower = world.rank * step + let step = 512 / size + let lower = rank * step let xPart = concatenated( [x[0..., lower ..< lower + step], x[0..., 512 + lower ..< 512 + lower + step]], axis: 1) let result = MLXDistributed.allSum(model(xPart), group: world) diff --git a/Tests/MLXTests/DistributedRingTests.swift b/Tests/MLXTests/DistributedRingTests.swift index 71f61951..7f3c5715 100644 --- a/Tests/MLXTests/DistributedRingTests.swift +++ b/Tests/MLXTests/DistributedRingTests.swift @@ -54,16 +54,18 @@ class DistributedRingTests: XCTestCase { /// Port of `test_groups`. private func groups(_ group: MLXDistributed.Group) throws { - XCTAssertEqual(group.size, Self.rankCount) - XCTAssertTrue(group.rank >= 0 && group.rank < group.size) + let rank = try group.rank + let size = try group.size + XCTAssertEqual(size, Self.rankCount) + XCTAssertTrue(rank >= 0 && rank < size) // initializing again yields the same group let again = try MLXDistributed.initialize() - XCTAssertEqual(again.size, group.size) - XCTAssertEqual(again.rank, group.rank) + XCTAssertEqual(try again.size, size) + XCTAssertEqual(try again.rank, rank) // the ring backend does not support splitting - XCTAssertThrowsError(try group.split(color: group.rank % 2)) + XCTAssertThrowsError(try group.split(color: rank % 2)) } /// Port of `test_all_reduce` and `test_all_reduce_extra`. @@ -71,13 +73,15 @@ class DistributedRingTests: XCTestCase { /// Every rank builds the same array and contributes its own row, so the /// reduction over the rows is the expected result in every process. private func reductions(_ group: MLXDistributed.Group) throws { + let rank = try group.rank + let size = try group.size let tolerances: [DType: Float] = [ .float32: 1e-6, .float16: 5e-3, .bfloat16: 1e-1, .complex64: 1e-6, ] let key = MLXRandom.key(0) var combinations = 0 defer { - if group.rank == 0 { + if rank == 0 { print("reductions: \(combinations) dtype and shape combinations") } } @@ -85,12 +89,12 @@ class DistributedRingTests: XCTestCase { for dtype in Self.dtypes { for shape in Self.shapes { let rtol = tolerances[dtype] ?? 0 - let x = (MLXRandom.uniform(0 ..< 1, [group.size] + shape, key: key) * 10) + let x = (MLXRandom.uniform(0 ..< 1, [size] + shape, key: key) * 10) .asType(dtype) let name = "\(dtype) \(shape)" combinations += 1 - let sum = MLXDistributed.allSum(x[group.rank], group: group) + let sum = MLXDistributed.allSum(x[rank], group: group) let expected = x.sum(axis: 0) var error = abs(sum - expected) if rtol > 0 { @@ -100,12 +104,12 @@ class DistributedRingTests: XCTestCase { XCTAssertLessThanOrEqual( error.max().asType(.float32).item(Float.self), rtol, "allSum \(name)") - let maximum = MLXDistributed.allMax(x[group.rank], group: group) + let maximum = MLXDistributed.allMax(x[rank], group: group) try checkedEval(maximum) XCTAssertTrue( (maximum .== x.max(axis: 0)).all().item(Bool.self), "allMax \(name)") - let minimum = MLXDistributed.allMin(x[group.rank], group: group) + let minimum = MLXDistributed.allMin(x[rank], group: group) try checkedEval(minimum) XCTAssertTrue( (minimum .== x.min(axis: 0)).all().item(Bool.self), "allMin \(name)") @@ -115,12 +119,14 @@ class DistributedRingTests: XCTestCase { /// Port of `test_all_gather` and `test_all_gather_extra`. private func allGather(_ group: MLXDistributed.Group) throws { + let rank = try group.rank + let size = try group.size for dtype in Self.dtypes { let x = MLXArray.ones([2, 2, 4]).asType(dtype) let gathered = MLXDistributed.allGather(x, group: group) try checkedEval(gathered) - XCTAssertEqual(gathered.shape, [group.size * 2, 2, 4], "allGather \(dtype)") + XCTAssertEqual(gathered.shape, [size * 2, 2, 4], "allGather \(dtype)") XCTAssertTrue( (gathered .== MLXArray(1).asType(dtype)).all().item(Bool.self), "allGather \(dtype)") @@ -128,9 +134,9 @@ class DistributedRingTests: XCTestCase { // the shards are concatenated in rank order rather than reduced let base = MLXArray([1, 2, 3]).asType(.float32) - let gathered = MLXDistributed.allGather(base * Float(group.rank + 1), group: group) + let gathered = MLXDistributed.allGather(base * Float(rank + 1), group: group) try checkedEval(gathered) - assertEqual(gathered, concatenated((0 ..< group.size).map { base * Float($0 + 1) })) + assertEqual(gathered, concatenated((0 ..< size).map { base * Float($0 + 1) })) } /// Port of `test_send_recv`. @@ -139,8 +145,8 @@ class DistributedRingTests: XCTestCase { /// Even ranks send first and odd ranks receive first, so that neighbors do /// not wait on each other. private func sendRecv(_ group: MLXDistributed.Group) throws { - let rank = group.rank - let size = group.size + let rank = try group.rank + let size = try group.size let right = (rank + 1) % size let left = (rank + size - 1) % size let key = MLXRandom.key(0) @@ -181,11 +187,12 @@ class DistributedRingTests: XCTestCase { /// The gathered array starts with rank 0's contribution, so only rank 0 /// sees a gradient. private func allGatherVJP(_ group: MLXDistributed.Group) throws { + let rank = try group.rank let gradient = grad { x in MLXDistributed.allGather(x, group: group)[0] }(MLXArray(1.0)) try checkedEval(gradient) - XCTAssertEqual(gradient.item(Float.self), group.rank == 0 ? 1.0 : 0.0) + XCTAssertEqual(gradient.item(Float.self), rank == 0 ? 1.0 : 0.0) } } diff --git a/Tests/MLXTests/DistributedTests.swift b/Tests/MLXTests/DistributedTests.swift index 4946ea6b..6229ec93 100644 --- a/Tests/MLXTests/DistributedTests.swift +++ b/Tests/MLXTests/DistributedTests.swift @@ -1,9 +1,11 @@ // Copyright © 2026 Apple Inc. +import Cmlx import Foundation -import MLX import XCTest +@testable import MLX + /// Single process tests for ``MLXDistributed``. /// /// Without `MLX_HOSTFILE` / `MLX_RANK` no backend can form a communication @@ -56,8 +58,22 @@ class DistributedTests: XCTestCase { func testDefaultGroupIsSingleton() throws { let group = try MLXDistributed.initialize() - XCTAssertEqual(group.rank, 0) - XCTAssertEqual(group.size, 1) + XCTAssertEqual(try group.rank, 0) + XCTAssertEqual(try group.size, 1) + + try MLXDistributed.clearBackends() + XCTAssertEqual(try group.size, 1) + } + + func testEmptyGroupGettersReportErrors() { + let group = MLXDistributed.Group(mlx_distributed_group_new()) + + XCTAssertThrowsError(try group.rank) { error in + XCTAssertTrue(error is MLXError) + } + XCTAssertThrowsError(try group.size) { error in + XCTAssertTrue(error is MLXError) + } } /// A strict init throws when no backend can form a group. @@ -94,6 +110,71 @@ class DistributedTests: XCTestCase { XCTAssertThrowsError(try group.split(color: 0)) } + // MARK: - Metadata all-gather initialization + + func testMetadataAllGatherBridgeForwardsBorrowedBuffers() { + let expected: [UInt8] = [0, 255, 127, 128] + let bridge = makeMetadataAllGatherBridge { rank, size, source, destination in + XCTAssertEqual(rank, 1) + XCTAssertEqual(size, 2) + XCTAssertEqual(Array(source), [0, 255]) + XCTAssertEqual(destination.count, expected.count) + guard destination.count == expected.count else { return } + destination.copyBytes(from: expected) + } + defer { bridge.dtor(bridge.payload) } + + let source: [CChar] = [0, -1] + var destination = [CChar](repeating: -2, count: 4) + let status = source.withUnsafeBufferPointer { source in + destination.withUnsafeMutableBufferPointer { destination in + bridge.callback( + 1, 2, source.baseAddress, destination.baseAddress, source.count, bridge.payload) + } + } + + XCTAssertEqual(status, 0) + XCTAssertEqual(destination.map { UInt8(bitPattern: $0) }, expected) + } + + func testMetadataAllGatherBridgeConvertsThrownErrorToFailureStatus() { + let bridge = makeMetadataAllGatherBridge { _, _, _, destination in + destination[0] = 42 + throw MetadataAllGatherTestError.callbackFailure + } + defer { bridge.dtor(bridge.payload) } + + let source: [CChar] = [42] + var destination: [CChar] = [-2, -2] + let status = source.withUnsafeBufferPointer { source in + destination.withUnsafeMutableBufferPointer { destination in + bridge.callback( + 0, 2, source.baseAddress, destination.baseAddress, source.count, bridge.payload) + } + } + + XCTAssertNotEqual(status, 0) + XCTAssertEqual(destination, [42, -2]) + } + + func testMetadataAllGatherFailureReleasesCapture() { + weak var captured: MetadataAllGatherCapture? + do { + let capture = MetadataAllGatherCapture() + captured = capture + XCTAssertThrowsError( + try MLXDistributed.initialize(backend: .ring) { [capture] _, _, _, _ in + withExtendedLifetime(capture) { + XCTFail("Callback should not run when initialization fails") + } + } + ) { error in + XCTAssertTrue(error is MLXError) + } + } + XCTAssertNil(captured) + } + // MARK: - Collectives /// In a group of size one the collectives return their input unchanged. @@ -132,3 +213,9 @@ class DistributedTests: XCTestCase { }) } } + +private enum MetadataAllGatherTestError: Error { + case callbackFailure +} + +private final class MetadataAllGatherCapture: Sendable {} diff --git a/scripts/verify-spm-manifest.py b/scripts/verify-spm-manifest.py index d3d3d0f6..2a1cf485 100755 --- a/scripts/verify-spm-manifest.py +++ b/scripts/verify-spm-manifest.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import json +import os import platform import subprocess @@ -10,14 +11,26 @@ def require(condition: bool, message: str) -> None: raise SystemExit(message) -manifest = json.loads( - subprocess.run( - ["swift", "package", "dump-package"], - check=True, - capture_output=True, - text=True, - ).stdout +# On macOS, /usr/bin/python3 is an xcrun shim that injects SDKROOT into the +# environment (possibly pointing at the CommandLineTools SDK). That SDK may +# not match the active Xcode's swift compiler, so drop it and let swift pick +# the SDK matching its own toolchain -- same as running from the shell. +env = dict(os.environ) +env.pop("SDKROOT", None) + +# Capture only stdout (the JSON); let stderr flow to the console so that +# SwiftPM errors are visible in CI logs. +result = subprocess.run( + ["swift", "package", "dump-package"], + stdout=subprocess.PIPE, + text=True, + env=env, ) +if result.returncode != 0: + raise SystemExit( + f"`swift package dump-package` failed with exit code {result.returncode}" + ) +manifest = json.loads(result.stdout) targets = {target["name"]: target for target in manifest["targets"]} dependencies = { diff --git a/xcode/xcconfig/Cmlx.xcconfig b/xcode/xcconfig/Cmlx.xcconfig index 1697339c..36ca1921 100644 --- a/xcode/xcconfig/Cmlx.xcconfig +++ b/xcode/xcconfig/Cmlx.xcconfig @@ -13,5 +13,5 @@ USER_HEADER_SEARCH_PATHS = $(SRCROOT)/../Source/Cmlx/mlx-c $(SRCROOT)/../Source/ MTL_COMPILER_FLAGS = -Wall -Wextra -fno-fast-math -Wno-c++17-extensions -Wno-c++20-extensions -GCC_PREPROCESSOR_DEFINITIONS = _METAL_=1 SWIFTPM_BUNDLE=\"com.apple.mlx.Cmlx\" MLX_VERSION=\"0.32.0\" MLX_USE_ACCELERATE=1 METAL_PATH=\"default.metallib\" ACCELERATE_NEW_LAPACK=1 +GCC_PREPROCESSOR_DEFINITIONS = _METAL_=1 SWIFTPM_BUNDLE=\"com.apple.mlx.Cmlx\" MLX_VERSION=\"0.32.0\" MLX_USE_ACCELERATE=1 METAL_PATH=\"default.metallib\" ACCELERATE_NEW_LAPACK=1 MLX_USE_JACCL=1