diff --git a/libs/mlx-swift b/libs/mlx-swift index fa6a4e8b4..a83e602e0 160000 --- a/libs/mlx-swift +++ b/libs/mlx-swift @@ -1 +1 @@ -Subproject commit fa6a4e8b4ed98dd08da7f0384f05003541523d87 +Subproject commit a83e602e0f431043b30fadbf0b510ad405e245e3 diff --git a/libs/mlx-swift-lm b/libs/mlx-swift-lm index c2fbbdc71..735c28ac0 160000 --- a/libs/mlx-swift-lm +++ b/libs/mlx-swift-lm @@ -1 +1 @@ -Subproject commit c2fbbdc71283eece0a1dd2f3b7b73f08fb9f6c0c +Subproject commit 735c28ac0e5e915b3d0e09e470efbd72a8e3472e diff --git a/provider-swift/Sources/ProviderCore/P2P/Parallelism.swift b/provider-swift/Sources/ProviderCore/P2P/Parallelism.swift new file mode 100644 index 000000000..5f3c53ae3 --- /dev/null +++ b/provider-swift/Sources/ProviderCore/P2P/Parallelism.swift @@ -0,0 +1,157 @@ +import Foundation + +// MARK: - Parallelism +// +// Selects the model-split strategy for a cluster session. Operator-visible +// via the `--parallelism` flag on `darkbloom serve`. The default is `.auto`, +// which picks `.tp` when both cluster peers can support tensor parallelism +// for the loaded model, falls back to `.pp` when they can't, and falls back +// to `.single` when no cluster peer is available. +// +// Why `.tp` wins by default on 2-Mac Thunderbolt 5: both Macs run all +// transformer layers in parallel rather than taking turns, halving +// single-stream decode latency. The per-layer allreduce cost is negligible +// on TB5 (~80 Gbps, sub-ms per hidden-dim vector). PP stays as a fallback +// for non-Llama models (until they get *TP variants) and for clusters +// where the link bandwidth makes allreduces too expensive. + +public enum Parallelism: String, Sendable, CustomStringConvertible { + /// Tensor parallelism: both ranks run all layers in parallel; per-layer + /// allreduce synchronizes activations. Requires a *TP model variant + /// (e.g. `LlamaModelTP`). + case tp + + /// Pipeline parallelism: rank 0 runs layers 0..N/2, rank 1 runs + /// layers N/2..N. One activation transfer per token. Works for any + /// model whose architecture exposes a layer-range entry point + /// (i.e. `callPartial` on Llama). + case pp + + /// No cluster: rank 0 runs the full model alone. The fallback when no + /// peer is connected. + case single + + /// Auto-select per `decide(...)`. + case auto + + public var description: String { rawValue } +} + +// MARK: - Decision + +extension Parallelism { + + /// Inputs the dispatcher considers when `--parallelism auto` is set. + public struct DecisionInputs: Sendable { + public let operatorChoice: Parallelism + /// Number of peers in the cluster (including self). 1 → single-rank. + public let worldSize: Int + /// True iff the loaded model has a published TP variant (e.g. Llama + /// → true via `LlamaModelTP`). Other architectures fall back to PP + /// until they get their own *TP variant. + public let modelHasTPVariant: Bool + /// `attentionHeads` of the loaded model. Used to check divisibility. + public let attentionHeads: Int + /// `kvHeads` of the loaded model. Used to check divisibility. + public let kvHeads: Int + /// True iff the underlying RDMA / jaccl backend is initialized and + /// `DistributedGroup` can be created. False on macOS < 26.2, on + /// pre-M5 hardware, or when `rdma_ctl` reports disabled. + public let distributedGroupAvailable: Bool + + public init( + operatorChoice: Parallelism, + worldSize: Int, + modelHasTPVariant: Bool, + attentionHeads: Int, + kvHeads: Int, + distributedGroupAvailable: Bool + ) { + self.operatorChoice = operatorChoice + self.worldSize = worldSize + self.modelHasTPVariant = modelHasTPVariant + self.attentionHeads = attentionHeads + self.kvHeads = kvHeads + self.distributedGroupAvailable = distributedGroupAvailable + } + } + + /// Pick the actual parallelism strategy from the operator's choice and + /// the runtime capabilities. Honors explicit operator overrides as long + /// as they're achievable; falls back with a reason when they aren't. + /// + /// Returns the chosen strategy and a short, operator-readable reason for + /// the decision (suitable for logging at startup). + public static func decide(_ inputs: DecisionInputs) -> (Parallelism, reason: String) { + // worldSize=1 short-circuits everything: no cluster peer, no + // parallelism. Operator can pass `--parallelism tp` and we'll still + // end up here because there's no peer to be tensor-parallel with. + if inputs.worldSize < 2 { + return (.single, reason: "no cluster peer connected (worldSize=\(inputs.worldSize))") + } + + switch inputs.operatorChoice { + case .single: + return (.single, reason: "operator selected --parallelism single") + + case .pp: + return (.pp, reason: "operator selected --parallelism pp") + + case .tp: + // Honor explicit TP if achievable. If not, fail closed rather than + // silently downgrade — operator asked for TP for a reason. + if !inputs.distributedGroupAvailable { + return ( + .single, + reason: "operator selected --parallelism tp but DistributedGroup unavailable (RDMA / M5 capability missing) — refusing to silently downgrade to PP, falling back to single-rank" + ) + } + if !inputs.modelHasTPVariant { + return ( + .single, + reason: "operator selected --parallelism tp but the loaded model has no TP variant — refusing to silently downgrade to PP, falling back to single-rank" + ) + } + if !canShard(heads: inputs.attentionHeads, worldSize: inputs.worldSize) + || !canShard(heads: inputs.kvHeads, worldSize: inputs.worldSize) + { + return ( + .single, + reason: "operator selected --parallelism tp but model heads don't divide evenly across worldSize=\(inputs.worldSize)" + ) + } + return (.tp, reason: "operator selected --parallelism tp; capabilities OK") + + case .auto: + return autoDecide(inputs) + } + } + + private static func autoDecide(_ inputs: DecisionInputs) -> (Parallelism, reason: String) { + // Auto: prefer TP if every capability lines up; else PP. + if !inputs.distributedGroupAvailable { + return ( + .pp, + reason: "auto → pp: DistributedGroup unavailable (RDMA / M5 capability missing)" + ) + } + if !inputs.modelHasTPVariant { + return (.pp, reason: "auto → pp: loaded model has no TP variant") + } + if !canShard(heads: inputs.attentionHeads, worldSize: inputs.worldSize) + || !canShard(heads: inputs.kvHeads, worldSize: inputs.worldSize) + { + return ( + .pp, + reason: + "auto → pp: model heads don't divide evenly across worldSize=\(inputs.worldSize) (attentionHeads=\(inputs.attentionHeads), kvHeads=\(inputs.kvHeads))" + ) + } + return (.tp, reason: "auto → tp: all capabilities OK") + } + + /// True iff `heads` is positive and divides evenly across `worldSize`. + public static func canShard(heads: Int, worldSize: Int) -> Bool { + worldSize > 0 && heads > 0 && heads % worldSize == 0 + } +} diff --git a/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift b/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift new file mode 100644 index 000000000..4f89ce06b --- /dev/null +++ b/provider-swift/Sources/ProviderCore/P2P/TensorParallelInference.swift @@ -0,0 +1,165 @@ +import CryptoKit +import Foundation +import MLX +import MLXLLM +import MLXLMCommon +#if canImport(os) +import os +#endif + +// MARK: - TensorParallelInference +// +// TP path counterpart to EncryptedPipelineInference. Both engines run on +// rank 0; their rank-1 counterparts (TensorParallelServer / EncryptedPipelineServer) +// run on the other Mac. The dispatcher (Parallelism.decide) picks which pair +// to instantiate at cluster-session-ready time. +// +// Key TP-vs-PP difference: in PP, only rank 0's first half of the model +// runs at a time, then rank 1's second half runs, with an activation +// transfer in between. In TP, BOTH ranks run all N layers in parallel, +// with allreduce per layer (handled inside `LlamaModelTP`'s sharded +// linear layers via the underlying `DistributedGroup`). One activation +// transfer per token is replaced by 2N small allreduces per token; the +// per-Mac compute drops by ~½ because each Mac runs half the heads. +// +// On Thunderbolt 5 the allreduce cost is sub-ms and the latency win is +// roughly 2× over PP for single-stream decode. +// +// STATUS: scaffold only. The cluster runtime is currently a stub +// (see ClusterCommand.swift line 431 — "Integrate ... with the coordinator +// request queue"). Wiring the engine into the live inference loop, jaccl +// DistributedGroup initialization, and the rank-1 serve loop are tracked +// for a follow-up. This file defines the API shape so the dispatcher and +// tests have something concrete to target. + +// MARK: - TensorParallelConfig + +public struct TensorParallelConfig: Sendable { + public let numLayers: Int + public let hiddenDim: Int + public let vocabSize: Int + public let worldSize: Int + + public init(numLayers: Int, hiddenDim: Int, vocabSize: Int, worldSize: Int) { + self.numLayers = numLayers + self.hiddenDim = hiddenDim + self.vocabSize = vocabSize + self.worldSize = worldSize + } +} + +// MARK: - TensorParallelEngine (rank 0) + +/// Drives tensor-parallel inference from rank 0's perspective. +/// +/// Both ranks load `LlamaModelTP` and call `model.callAsFunction(input)` in +/// lockstep; the sharded linear layers internally allreduce activations via +/// the `DistributedGroup`. Both ranks produce identical logits; rank 0 +/// samples the next token and streams it to the consumer, then signals +/// rank 1 to advance. +/// +/// On a singleton group (worldSize=1), this degenerates to ordinary +/// single-rank inference — the sharded layers' allreduce is a no-op. +/// +/// On link failure, each step is retried up to 3 times with a 3-second +/// delay (mirrors EncryptedPipelineEngine semantics). After 3 consecutive +/// failures, `ClusterError.serviceUnavailable` is thrown. +public actor TensorParallelEngine { + private let config: TensorParallelConfig + /// Held as `any LLMModel` so callers can pass either the fp16 variant + /// (`LlamaModelTP`) or the quantized variant (`LlamaModelTPQ`) without + /// the engine needing to be templated on the concrete type. Callers + /// (typically the dispatcher in `ClusterDiscovery`) are responsible + /// for passing a TP-capable model — passing a non-TP `LlamaModel` + /// would type-check but produce single-rank semantics. + private let model: any LLMModel + private let tokenizer: any Tokenizer + private var cache: [any KVCache] + private let session: ClusterSession + + private let logger = Logger( + subsystem: "io.darkbloom.provider", category: "TensorParallelEngine") + + public init( + config: TensorParallelConfig, + model: any LLMModel, + tokenizer: any Tokenizer, + session: ClusterSession + ) { + self.config = config + self.model = model + self.tokenizer = tokenizer + self.session = session + self.cache = [] + logger.info( + """ + TensorParallelEngine initialized: layers=\(config.numLayers), \ + hidden=\(config.hiddenDim), vocab=\(config.vocabSize), \ + worldSize=\(config.worldSize) + """) + } + + /// Reset the KV cache to a fresh state. Called between distinct requests. + public func resetCache() { + cache = model.newCache(parameters: nil) + logger.info("TP KV cache reset (per-rank shard of \(self.cache.count) layers)") + } + + // The actual decode loop is intentionally NOT implemented here yet — + // it depends on the cluster runtime integration that's still a stub + // in ClusterCommand.swift. The shape will be: + // + // func generate(prompt: [Int], maxTokens: Int) -> AsyncStream { + // AsyncStream { continuation in + // Task { + // // 1. Send prompt token IDs to rank 1 over ThunderboltLink. + // // 2. Both ranks call model.callAsFunction (synced via allreduce). + // // 3. Rank 0 samples token from logits, yields to consumer. + // // 4. Rank 0 signals rank 1 with the chosen token + "continue". + // // 5. Repeat until EOS or maxTokens. + // } + // } + // } +} + +// MARK: - TensorParallelServer (rank 1) + +/// Drives the rank-1 side of tensor-parallel inference. Symmetric to rank 0 +/// — both ranks run `model.callAsFunction(input)` in lockstep with allreduce +/// synchronization. Rank 1 doesn't sample; it just provides compute for the +/// allreduces and waits for rank 0's "next input" / "session end" signals +/// over the ThunderboltLink control channel. +public actor TensorParallelServer { + private let config: TensorParallelConfig + /// See `TensorParallelEngine.model` — same polymorphism rationale. + private let model: any LLMModel + private let peer: ClusterPeer + + private let logger = Logger( + subsystem: "io.darkbloom.provider", category: "TensorParallelServer") + + public init( + config: TensorParallelConfig, + model: any LLMModel, + peer: ClusterPeer + ) { + self.config = config + self.model = model + self.peer = peer + logger.info( + "TensorParallelServer initialized: layers=\(config.numLayers), worldSize=\(config.worldSize)" + ) + } + + // Same stub status as TensorParallelEngine — the serve loop will be: + // + // func serve() async throws { + // try await peer.serve(modelState: ..., inferenceHandler: { conn, _, _ in + // // Loop: + // // Receive input token IDs from rank 0. + // // model.callAsFunction(input, cache) — synced via allreduce. + // // Receive sampled token from rank 0; advance cache. + // // On sessionEnd, exit loop. + // }) + // } +} diff --git a/provider-swift/Sources/darkbloom/StartCommand.swift b/provider-swift/Sources/darkbloom/StartCommand.swift index bd254d454..61252b933 100644 --- a/provider-swift/Sources/darkbloom/StartCommand.swift +++ b/provider-swift/Sources/darkbloom/StartCommand.swift @@ -2,6 +2,8 @@ import Foundation import ArgumentParser import ProviderCore +extension Parallelism: ExpressibleByArgument {} + struct Start: AsyncParsableCommand { static let configuration = CommandConfiguration( abstract: "Start the provider as a background service.", @@ -37,6 +39,12 @@ struct Start: AsyncParsableCommand { @Flag(help: "Register as RDMA-capable and auto-discover Thunderbolt-connected peers for pipeline inference. Requires a Secure Enclave and a logged-in account.") var rdmaEnabled = false + @Option( + help: + "Cluster parallelism strategy: auto (default — prefer TP, fall back to PP), tp (tensor-parallel — both Macs run all layers in parallel; requires Llama-class model + M5 + rdma_ctl enabled), pp (pipeline-parallel — split layers across Macs), single (no clustering)." + ) + var parallelism: Parallelism = .auto + mutating func run() async throws { // GPU is required. Reject CPU fallback up-front so we never // come up reporting healthy and then silently churn at 0.5 tok/s. @@ -252,6 +260,7 @@ struct Start: AsyncParsableCommand { ) Task { await discovery.start() } print("RDMA cluster discovery enabled — watching for Thunderbolt peers") + print("Cluster parallelism preference: \(parallelism) (resolved at session handshake)") } } catch { printError("Warning: --rdma-enabled requires Secure Enclave support: \(error.localizedDescription)") diff --git a/provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift b/provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift new file mode 100644 index 000000000..d8546c7b8 --- /dev/null +++ b/provider-swift/Tests/ProviderCoreTests/ParallelismTests.swift @@ -0,0 +1,138 @@ +import Foundation +import Testing +@testable import ProviderCore + +// MARK: - Helpers + +private func defaultInputs( + operatorChoice: Parallelism = .auto, + worldSize: Int = 2, + modelHasTPVariant: Bool = true, + attentionHeads: Int = 32, + kvHeads: Int = 8, + distributedGroupAvailable: Bool = true +) -> Parallelism.DecisionInputs { + Parallelism.DecisionInputs( + operatorChoice: operatorChoice, + worldSize: worldSize, + modelHasTPVariant: modelHasTPVariant, + attentionHeads: attentionHeads, + kvHeads: kvHeads, + distributedGroupAvailable: distributedGroupAvailable + ) +} + +// MARK: - Single-rank short-circuit + +@Test func parallelismShortCircuitsToSingleWhenNoPeer() { + let (chosen, reason) = Parallelism.decide(defaultInputs(worldSize: 1)) + #expect(chosen == .single) + #expect(reason.contains("worldSize=1")) +} + +@Test func parallelismShortCircuitsToSingleEvenIfOperatorAsksForTP() { + // worldSize < 2 always wins; operator choice can't conjure a peer. + let (chosen, _) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, worldSize: 1)) + #expect(chosen == .single) +} + +// MARK: - Auto-decide path (the default) + +@Test func parallelismAutoPicksTPWhenAllCapabilitiesAlign() { + let (chosen, reason) = Parallelism.decide(defaultInputs()) + #expect(chosen == .tp) + #expect(reason.contains("auto")) +} + +@Test func parallelismAutoFallsBackToPPWhenDistributedGroupUnavailable() { + // Capability gate failed (non-M5, or rdma_ctl disabled). PP can still + // run because callPartial only needs ThunderboltLink, not jaccl. + let (chosen, reason) = Parallelism.decide( + defaultInputs(distributedGroupAvailable: false)) + #expect(chosen == .pp) + #expect(reason.contains("DistributedGroup unavailable")) +} + +@Test func parallelismAutoFallsBackToPPWhenModelHasNoTPVariant() { + // Non-Llama model (e.g. Mistral, Qwen) until they get their own *TP + // variants. PP works via callPartial. + let (chosen, reason) = Parallelism.decide( + defaultInputs(modelHasTPVariant: false)) + #expect(chosen == .pp) + #expect(reason.contains("no TP variant")) +} + +@Test func parallelismAutoFallsBackToPPWhenHeadsDontDivide() { + // Odd head count can't shard across worldSize=2. + let (chosen, reason) = Parallelism.decide( + defaultInputs(attentionHeads: 33)) + #expect(chosen == .pp) + #expect(reason.contains("divide evenly")) +} + +@Test func parallelismAutoFallsBackToPPWhenKVHeadsDontDivide() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(kvHeads: 7)) + #expect(chosen == .pp) + #expect(reason.contains("divide evenly")) +} + +// MARK: - Explicit operator overrides + +@Test func parallelismHonorsExplicitPP() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .pp)) + #expect(chosen == .pp) + #expect(reason.contains("operator selected")) +} + +@Test func parallelismHonorsExplicitSingle() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .single)) + #expect(chosen == .single) + #expect(reason.contains("operator selected")) +} + +@Test func parallelismHonorsExplicitTPWhenAchievable() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp)) + #expect(chosen == .tp) + #expect(reason.contains("operator selected --parallelism tp")) +} + +// MARK: - Explicit TP refuses to silently downgrade + +@Test func parallelismExplicitTPFailsClosedWhenDistributedGroupUnavailable() { + // Operator asked for TP. Capability gate failed. We refuse to silently + // give them PP — that would mask a misconfiguration. Fall back to single + // and surface the reason. + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, distributedGroupAvailable: false)) + #expect(chosen == .single) + #expect(reason.contains("refusing to silently downgrade")) +} + +@Test func parallelismExplicitTPFailsClosedWhenModelHasNoTPVariant() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, modelHasTPVariant: false)) + #expect(chosen == .single) + #expect(reason.contains("refusing to silently downgrade")) +} + +@Test func parallelismExplicitTPFailsClosedWhenHeadsDontDivide() { + let (chosen, reason) = Parallelism.decide( + defaultInputs(operatorChoice: .tp, attentionHeads: 33)) + #expect(chosen == .single) + #expect(reason.contains("divide evenly")) +} + +// MARK: - Divisibility helper + +@Test func parallelismCanShardRequiresEvenDivision() { + #expect(Parallelism.canShard(heads: 32, worldSize: 2) == true) + #expect(Parallelism.canShard(heads: 33, worldSize: 2) == false) + #expect(Parallelism.canShard(heads: 64, worldSize: 4) == true) + #expect(Parallelism.canShard(heads: 0, worldSize: 2) == false) + #expect(Parallelism.canShard(heads: 32, worldSize: 0) == false) +}