Distributed Inference and Training with MLX

import Foundation
import MLX
import MLXDistributed

// MARK: - Distributed Group Setup
// Each node in the cluster runs this same executable.
// mlx.launch SSHes into each node and starts it with the correct rank env vars.

struct DistributedClusterDemo {
    static func run() async throws {
        // Initialize the default distributed process group.
        // MLX reads JACCL/RDMA configuration from environment variables
        // set by mlx.launch via the hostfile.
        let group = MLXDistributed.init()

        let rank = group.rank        // Index of this machine in the cluster (0-based)
        let worldSize = group.size   // Total number of machines in the cluster

        print("Node \(rank) of \(worldSize) initialized.")

        // MARK: - All-Reduce: Sum activations across all nodes
        // In tensor parallelism each node computes a partial result for every
        // layer. An all-reduce sums those partials so every node has the
        // full result before the next layer.

        // Simulate a partial activation computed locally on this node.
        // In a real LLM this would be the output of a sharded linear layer.
        var localActivation = MLXArray(
            (0..<8).map { Float($0) + Float(rank) * 8 },
            [8]
        )

        print("Node \(rank) partial activations: \(localActivation)")

        // All-reduce (sum) across the group — blocks until all nodes contribute.
        let globalActivation = group.allReduce(
            localActivation,
            operation: .sum
        )

        // Force evaluation so the result is materialized before printing.
        eval(globalActivation)

        print("Node \(rank) global activations after all-reduce: \(globalActivation)")

        // MARK: - All-Gather: Collect sharded embeddings from every node
        // Pipeline parallelism exchanges full activations at layer boundaries.
        // All-gather lets every node obtain the complete tensor.

        let gatheredActivations = group.allGather(localActivation)
        eval(gatheredActivations)

        if rank == 0 {
            // Only the rank-0 (coordinator) node logs the aggregated result.
            print("Coordinator sees gathered shape: \(gatheredActivations.shape)")
            // Shape is [worldSize * 8] — each node's shard concatenated.
        }

        // Barrier: wait for all nodes before exiting.
        group.barrier()
        print("Node \(rank): distributed run complete.")
    }
}

// Entry point — run as a Swift executable via:
// mlx.launch --hostfile cluster.json -- /path/to/DistributedClusterDemo
try await DistributedClusterDemo.run()