-
MLX를 사용한 분산 추론 및 학습 살펴보기
MLX를 사용하여 여러 대의 Mac에 걸쳐 머신 러닝 워크로드를 확장하세요. 인터커넥트 효율성, 대규모 모델 추론, 요청 일괄 처리, 분산 학습 관련 과제를 해결하는 방법을 알아보세요. 책상 위의 Mac 몇 대로 어떻게 고가의 클라우드 인프라를 대체하여 고사양 AI 워크로드를 처리할 수 있는지 살펴보세요.
챕터
- 0:00 - Introduction
- 2:09 - Distributed communication
- 4:32 - Setting up your cluster
- 10:33 - Distributed inference and fine-tuning
- 13:35 - Model parallelism strategies
- 15:53 - Distributed fine-tuning
- 18:34 - CLI, Python, Swift, and C++ APIs
- 20:45 - Next steps
리소스
- MLX Swift LM on GitHub
- MLX Swift Examples
- MLX Examples
- MLX Swift
- MLX LM - Python API
- MLX Explore - Python API
- MLX Framework
- MLX
관련 비디오
WWDC26
WWDC25
-
비디오 검색…
안녕하세요, 저는 Tatiana MLX 팀 연구 과학자입니다 로컬 LLM에 있어 놀라운 시대가 이어지고 있습니다 모델은 점점 더 커지고 놀라운 새 기능들을 갖추고 있습니다 더욱 똑똑해지며 더 어려운 문제들을 해결합니다 발전할수록 더 많이 활용합니다 더 긴 컨텍스트, 더 어려운 작업 더 복잡한 워크플로우까지 결국 단일 머신의 메모리, 연산 성능 또는 대역폭이 한계에 다다릅니다 WWDC 26 영상 "Run local agentic AI on the Mac using MLX"에서 Mac에서 AI 에이전트를 로컬로 실행하는 방법을 보여줍니다 하지만 여러 기기가 있다면 로컬 AI를 더욱 발전시킬 수 있습니다 더 큰 LLM을 실행하거나 가속할 수 있죠 분산 추론과 학습을 통해 오늘은 MLX로 여러 Mac에 걸쳐 확장하는 방법을 깊이 알아봅니다 바로 여러분의 책상 위 하드웨어로요 명령줄 인터페이스로 시작해서 모델을 실행하고 여러분의 머신에서 실험을 위한 Python API로 넘어가고 마지막으로 Swift를 사용해 워크플로우를 앱에 직접 임베드합니다 시작하겠습니다! 먼저 전체 하드웨어와 소프트웨어 스택을 살펴봅니다 Apple Silicon에서 분산 워크로드를 가능하게 하는 것들이죠 그런 다음 모든 것을 합쳐 네 대의 M3 Ultra로 클러스터를 구성합니다 모든 단계를 살펴봅니다 머신을 연결할 올바른 토폴로지 선택부터 빠른 통신 활성화와 분산 작업 시작까지 클러스터가 준비되면 흥미로운 부분으로 넘어갑니다 빠르고 로컬한 분산 LLM 추론과 파인튜닝입니다 MLX로 실행하고 단일 Mac과 나란히 비교하며 MLX가 클러스터에 걸쳐 모델을 분산하는 방식을 봅니다 대부분의 예제는 명령줄 인터페이스로 보여드리고 마지막으로 분산 통신이 Python, Swift, C++ API를 통해서도 제공되는 방식을 보여드립니다 Apple Silicon의 분산 통신부터 살펴보겠습니다 데이터를 빠르게 송수신하려면 머신들이 물리적 링크로 연결되어야 합니다 즉, 인터커넥트입니다 그 위에 전송 프로토콜도 필요합니다 바이트를 전달하는 메커니즘 한 머신의 메모리에서 다른 머신으로 macOS 26.2부터 RDMA, 즉 원격 직접 메모리 접근 프로토콜이 Thunderbolt 5에서 지원됩니다 RDMA는 한 머신의 메모리에서 다른 머신으로 데이터를 직접 이동합니다 대부분의 CPU 및 운영 체제 오버헤드를 줄이면서요 Thunderbolt를 통한 RDMA는 고대역폭을 제공합니다 저지연 통신 분산 워크로드에 필요한 수준의 하지만 단독으로는 두 머신 사이의 원시 데이터 이동만 제공합니다 따라서 분산 프로그램에는 더 높은 수준의 무언가가 필요합니다 통신 백엔드 데이터 전송을 위한 통신 프리미티브를 제공하는 개별 머신 간의 또는 전체 그룹 간의 조정을 위한 이 두 가지 연산은 분산 학습과 추론의 기본 구성 요소입니다 바로 이 부분에서 JACCL이 등장합니다
JACCL은 오픈 소스 집합 통신 라이브러리로 Apple이 만들었습니다 Thunderbolt를 통한 RDMA를 활용하여 집합 통신 프리미티브를 제공합니다 머신 간 데이터 전송 그룹 전체에서 결과 취합 저수준 전송을 직접 관리할 필요 없이요 머신 러닝에만 국한되지 않아서 Apple Silicon의 모든 분산 워크로드에 활용할 수 있습니다 스택의 마지막 구성 요소는 머신 러닝 프레임워크입니다 통신 백엔드를 활용해 분산 추론과 학습을 수행하는 바로 MLX입니다 MLX는 오픈 소스 머신 러닝 라이브러리로 Apple이 Apple Silicon을 위해 만들었습니다 저지연 분산 통신을 위해 JACCL을 활용하며 클러스터 전체에서 분산 작업을 조율하는 도구를 제공합니다 MLX가 처음이시라면 영상을 확인하세요 WWDC25의 "Getting Started with MLX on Apple Silicon"입니다
이제 전체 스택을 이해했습니다 이제 모든 것을 합쳐 클러스터를 구성해 봅니다 같은 작업을 함께 수행하는 머신 그룹입니다 M3 Ultra 4대를 사용합니다 클러스터를 설정하려면 Thunderbolt 5 케이블로 머신을 연결해야 합니다 연결 방법은 여러 가지이며 토폴로지는 통신 시간에 직접적인 영향을 줍니다 먼저 그 시간을 결정하는 요소를 살펴봅니다 다음으로 실제로 머신을 연결하는 방법을 봅니다 JACCL이 지원하는 토폴로지와 그 사이의 트레이드오프입니다 그 다음, 빠른 통신을 위해 머신에서 RDMA를 활성화하는 방법을 보여드립니다 마지막으로 MLX를 사용해 클러스터에서 분산 작업을 시작합니다
통신 시간에는 두 가지 구성 요소가 있습니다 지연 시간과 전송 시간입니다 지연 시간은 각 통신 작업에 지불하는 고정 비용으로 전송되는 데이터 양과는 무관합니다
전송 시간은 링크를 통해 데이터를 이동하는 비용이며 메시지 크기에 따라 커집니다 링크 대역폭에도 달려 있습니다
소규모 메시지의 경우 데이터 이동 비용은 미미해서 지연 시간이 지배합니다
대규모 메시지의 경우 트레이드오프가 반대입니다 통신이 지연 바운드인지 대역폭 바운드인지에 따라 다른 토폴로지를 선호할 수 있습니다
JACCL은 메시와 링, 두 가지를 지원합니다 풀 메시에서는 모든 머신이 다른 모든 머신에 직접 연결됩니다 따라서 그룹 통신의 지연 시간이 가능한 한 최소화됩니다 링에서는 각 노드가 인접한 두 노드에만 연결됩니다 인접하지 않은 노드 간 통신은 중간 머신을 거쳐야 합니다 지연 시간이 증가하죠 하지만 링은 머신당 필요한 케이블과 포트가 더 적어서 더 많은 노드로 확장하기 쉽습니다 각 노드는 두 개의 연결만 가지므로 추가 Thunderbolt 포트를 사용해 이웃당 두세 개의 케이블을 연결할 수 있습니다 (Mac에 따라 다름) 링크당 대역폭을 늘리고 전송 시간을 줄입니다 머신들이 메시로 연결된 경우 각 통신을 유연하게 라우팅할 수 있습니다 메시 토폴로지나 링 토폴로지를 통해서
JACCL의 장점은 자동으로 최적의 토폴로지를 선택한다는 것입니다 메시지 크기와 통신 작업에 따라 지연이 중요할 땐 메시 대역폭이 중요할 땐 링 이 유연성을 위해 모든 M3 Ultra를 메시로 연결합니다
M3 Ultra를 모두 연결했으니 이제 모든 머신에서 RDMA를 활성화해야 합니다 머신의 설정을 열고 "RDMA"를 검색하세요
"Enable RDMA over Thunderbolt"를 클릭하고
RDMA를 활성화하고 재부팅하세요
좋습니다! Mac들이 Thunderbolt 5 케이블로 연결되고 RDMA가 활성화됩니다 이제 분산 프로그램을 시작할 방법이 필요합니다
한 가지 방법은 로컬 네트워크를 통하는 것입니다 예를 들어, Wi-Fi나 이더넷을 통해 클러스터에 SSH 접근이 가능한 머신에서 예를 들어 저의 경우 MacBook에서 각 Mac에 연결해 프로그램을 시작합니다 그 시점부터는 모든 머신이 Thunderbolt 링크를 통해 직접 통신합니다 MLX가 제공하는 시작 도우미가 이 모든 것을 처리해 줍니다!
MacBook에서 mlx.launch를 실행하면 클러스터를 조율합니다 실행할 실행 파일을 지정하고 클러스터를 설명하는 JSON 호스트파일도 지정합니다 호스트파일의 호스트명을 사용해 각 노드에 SSH로 접속하고 모든 머신에서 실행 파일을 시작합니다 클러스터를 설명하는 호스트파일의 형식을 살펴봅니다 JSON 배열로, 노드당 하나의 항목입니다 "ssh"는 mlx.launch가 머신에 접근할 때 사용하는 호스트명입니다 "ips"는 로컬 네트워크에서 머신의 IP 주소로 JACCL이 노드 간 초기 조정에 사용합니다 "rdma"는 RDMA 기기 이름 목록입니다 각 Thunderbolt 피어 연결에 대한
수동으로 작성할 수도 있지만 MLX도 제공합니다 `mlx.distributed_config` 도우미 스크립트로 자동 생성할 수 있습니다 호스트명 목록과 출력 경로를 전달합니다 설정에 환경 변수를 임베드할 수도 있습니다 시작 시 모든 노드에서 자동으로 설정됩니다 여기서는 MLX_METAL_FAST_SYNCH=1을 설정합니다 더 빠른 GPU-CPU 동기화를 활성화합니다 분산 작업에 매우 중요합니다 연산은 GPU에서 실행되고 GPU에서 통신은 CPU에서 실행되기 때문입니다 --auto-setup 플래그를 사용해 Thunderbolt 네트워크를 자동으로 구성할 수도 있습니다 Communication --backend 인수는 메시인지 링인지 결정합니다 메시의 경우 이 예제처럼 --backend를 jaccl로 설정하고 링의 경우에는 jaccl-ring으로 변경합니다 이 명령을 실행해서 클러스터의 호스트파일을 생성합니다
먼저 모든 호스트가 SSH로 연결 가능한지 확인합니다 그런 다음 각 머신의 Thunderbolt 포트를 조사해서 어떤 머신이 어떤 머신과 물리적으로 연결되어 있는지 파악합니다 토폴로지 맵을 구성하는 것이죠 --auto-setup을 전달했으므로 모든 머신에서 Thunderbolt Bridge를 비활성화하고 모든 머신에서 각 Thunderbolt 링크를 RDMA용으로 구성합니다 마지막으로 mlx.launch에 필요한 모든 정보가 담긴 JSON 호스트파일을 작성합니다 --auto-setup 플래그 없이 실행하면 스크립트가 구성 명령을 출력합니다 직접 검토하고 실행할 수 있도록요
이제 클러스터가 준비됩니다 흥미로운 부분으로 넘어갑니다 분산 언어 모델 추론 및 파인튜닝입니다 가장 쉬운 시작 방법은 명령줄 인터페이스와 MLX LM입니다 MLX LM은 오픈 소스 Python 패키지로 MLX 위에 구축되었습니다 명령줄 도구를 제공합니다 Apple Silicon에서 언어 모델을 로컬로 실행하기 위한 Python API도 포함합니다 영상 "Explore large language models on Apple Silicon with MLX"를 확인하세요 WWDC25에서 단일 기기로 시작하는 방법을 다룹니다
지난 해 보여드렸듯이 단일 Mac에서 모델과 채팅하려면 mlx_lm.chat으로 명령줄 인터페이스를 사용하면 됩니다 터미널에서 실행하며 사용할 모델을 지정합니다 예를 들어, Qwen 3.6 응답에 대한 최대 토큰 수도 지정합니다 MLX LM은 내부적으로 모델을 로드하여 단일 머신에서 실행합니다
명령줄 인터페이스로 클러스터에서 같은 모델과 채팅하려면 mlx.launch로 명령을 감쌉니다 MacBook의 터미널에서 mlx.launch를 실행하고 --hostfile로 클러스터 설정을 가리킵니다 더블 대시 뒤에 동일한 mlx_lm.chat 명령을 전달합니다 각 노드의 실행 파일 원격 경로를 사용한다는 점만 다릅니다 명령은 거의 동일하고 MLX LM이 모델을 분산하고 자동으로 추론을 조율합니다 MLX 같은 필요한 라이브러리는 모든 Mac에 설치되어야 하고 실행 파일도 모든 머신에서 접근 가능해야 합니다 명령줄 인터페이스에서 한 줄로 모델을 실행할 수 있습니다 전체 클러스터에 걸쳐! 두 가지를 나란히 시험하고 Qwen 3.6으로 채팅해 봅니다 270억 개의 파라미터를 가진 모델입니다 단일 M3 Ultra와 4대의 M3 Ultra에서 이미 양쪽에서 mlx_lm.chat을 시작했습니다 왼쪽에는 단일 M3 Ultra에서 모델이 로드되었고 오른쪽에는 네 대의 머신에 분산되었습니다 양쪽에 "Implement a transformer model in MLX"를 프롬프트합니다
매우 인상적인 속도 향상입니다! 클러스터는 거의 세 배의 속도로 토큰을 생성합니다 단일 머신에 비해 Qwen 3.6 모델의 경우 여러 Mac에서 모델을 실행하면 추론 속도를 크게 높일 수 있습니다 정확한 속도 향상은 모델 크기와 아키텍처에 따라 다릅니다 하지만 분산을 선택하는 이유는 시간 개선만이 아닙니다 때로는 모델이 단순히 한 머신에 너무 큰 경우도 있습니다 Kimi 2.6은 예를 들어 총 1조 개의 파라미터를 가집니다 8비트 양자화를 적용해도 가중치만 약 1테라바이트의 메모리가 필요합니다 단일 M3 Ultra에는 들어가지 않지만 4대에 걸쳐서는 가능합니다 실제로 가중치와 연산을 머신에 어떻게 분할할까요?
MLX와 MLX LM은 두 가지 방식을 지원합니다 파이프라인과 텐서 병렬화입니다
파이프라인 병렬화는 모델을 깊이로 분할합니다 이 경우 각 머신이 레이어 그룹을 보유하며 데이터는 머신을 통해 순차적으로 이동합니다 추론 속도는 개선되지 않습니다 각 토큰은 여전히 통과해야 하기 때문입니다 레이어 그룹을 하나씩 순서대로 하지만 장점은 단순한 통신입니다 머신들은 레이어 그룹 경계에서만 활성화를 교환합니다 텐서 병렬화는 모델을 너비로 분할합니다 이 경우 각 머신이 모든 레이어의 일부를 보유하므로 모든 머신이 동일한 토큰을 동시에 처리합니다 레이어별 병렬 연산으로 추론 속도가 향상됩니다 하지만 트레이드오프로 훨씬 더 빈번한 통신이 발생합니다 모든 레이어와 모든 토큰에 대해 따라서 저지연이 중요합니다 이 경우에 메시 토폴로지가 중요한 이유입니다 모든 머신이 단일 홉으로 다른 모든 머신에 접근할 수 있기 때문입니다
텐서 병렬화는 MLX LM에서 기본 분산 전략입니다 파이프라인 병렬화로 모델을 분산하려면 명령에 --pipeline 플래그를 추가하면 됩니다 모든 모델이 파이프라인 병렬화를 지원하지는 않습니다 이제 클러스터에서 1조 개의 파라미터를 가진 Kimi 2.6과 채팅해 봅니다 우리 클러스터에서
이전처럼 MacBook에서 mlx.launch를 사용합니다 호스트파일을 가리키면서 --pipeline 플래그를 전달하지 않으므로 텐서 병렬화를 사용합니다 잠시 기다립니다 mlx.launch가 모든 머신에 연결하고 MLX LM이 모델을 로드해서 분산하고 채팅을 시작합니다
모델이 로드되었습니다! 모델에 프롬프트합니다 "Implement machine learning architecture for GPT in Python with MLX"
이렇게 됩니다. 한 명령으로 1조 개의 파라미터를 가진 대규모 모델이 여러분의 Mac에서 로컬로 실행되며 질문에 답합니다
MLX와 MLX LM으로 언어 모델 추론만 실행하는 것이 아닙니다 여러분의 하드웨어에서 모델을 파인튜닝할 수도 있습니다 빠르고, 효율적이며, 완전히 비공개입니다 데이터가 머신을 떠나지 않습니다 단일 Mac으로 시작해서 클러스터로 확장해 봅니다 단일 머신에서 파인튜닝하거나 학습할 때는 학습 데이터를 배치로 나눕니다 여러 예제로 구성된 집합입니다 각 배치에 대해 Mac이 그래디언트를 계산하고 모델 가중치를 업데이트합니다 학습 데이터셋에 대해 한 번 이상 이 과정을 반복합니다 모델이 원하는 품질에 달할 때까지 학습 데이터를 빠르게 처리할수록 파인튜닝이 더 빨리 완료됩니다 그렇다면 여러 머신을 사용해 어떻게 속도를 높일 수 있을까요? 개념은 간단합니다 모든 Mac에 모델을 복제합니다 각 머신이 다른 배치의 데이터를 받아 로컬에서 그래디언트를 계산합니다 그래디언트를 평균 내어 모델 업데이트에 모든 배치의 정보를 반영합니다 모든 배치에서 이를 데이터 병렬 학습이라고 합니다 모델은 복제되고 데이터는 머신에 걸쳐 병렬로 처리되기 때문입니다 이것이 속도 향상의 원천입니다 N대의 머신으로 데이터를 최대 N배 더 빠르게 처리할 수 있습니다 정말 놀랍습니다! MLX LM으로 데이터 병렬화를 어떻게 사용하는지 알아봅니다 이전과 마찬가지로 단일 기기와의 차이점은 mlx.launch로 작업을 시작한다는 것입니다 MacBook에서 원격 머신의 mlx_lm.lora 경로를 지정하면서 데이터 분산은 MLX LM이 처리하며 명령은 거의 동일합니다 기기 수에 맞게 --batch-size를 조정합니다 각 머신이 동일한 수의 샘플을 처리하도록 이전과 같이 단계당 같은 수의 샘플을 처리합니다 90억 개의 파라미터를 가진 Qwen 3.5를 파인튜닝합니다 단일 머신과 클러스터에서 실행해서 모델이 초당 처리하는 토큰 수를 비교합니다 왼쪽의 단일 기기에서 파인튜닝을 시작하고 오른쪽의 클러스터에서도 시작합니다 mlx.launch와 호스트파일을 사용해서 원격 머신의 mlx_lm.lora 경로를 지정합니다 먼저 데이터와 모델을 로드하고 학습이 시작됩니다 단일 M3 Ultra는 초당 약 180개의 토큰을 처리합니다 클러스터에서는 초당 약 600개의 토큰을 처리합니다 파인튜닝에서 3배 이상의 속도 향상입니다 이제 MLX로 기기들을 로컬 학습 클러스터로 전환할 수 있습니다 클라우드 없이 효율적인 파인튜닝을 위해 지금까지 명령줄 인터페이스로 MLX LM의 분산 추론과 파인튜닝을 사용했습니다 하지만 MLX는 세밀한 제어를 제공합니다 분산과 분산 연산에 대한 유연한 Python, Swift, C++ API를 통해 Python과 C++로 모델을 실험하거나 Swift로 모델을 앱에 임베드할 수 있습니다 예제를 살펴봅니다 Python API와 MLX LM으로 분산 추론을 실행하려면 먼저 통신을 위한 분산 그룹을 초기화합니다 그런 다음 원하는 병렬화 유형을 정의합니다 예를 들어, 텐서 병렬화 마지막으로 sharded_load 함수로 모델을 분산합니다 그 다음에는 단일 기기에서와 동일하게 모델을 사용합니다 MLX LM이 모든 분산 통신을 내부적으로 처리합니다
모델과 분산에 대한 더 많은 제어를 원한다면 MLX 자체의 저수준 프리미티브를 사용할 수 있습니다 예를 들어, 간단한 Linear 레이어를 정의한 후 shard_linear 함수로 텐서 병렬화를 적용할 수 있습니다 all reduce 같은 기본 분산 연산도 제어할 수 있습니다 Python, Swift, C++에서 JACCL을 통해 분산 그룹을 초기화한 후 모든 Mac에 걸쳐 텐서의 집합 분산 합산을 해당 MLX 프리미티브로 수행합니다 세션 초반에 언급했듯이 JACCL은 단독으로도 사용 가능합니다 모든 애플리케이션에서 활용할 수 있습니다 분산 통신이 필요한 경우라면 ML이 아닌 애플리케이션도 포함해서 JACCL은 MLX 없이도 빌드할 수 있으며 C++ API를 제공합니다 통신 프리미티브를 갖춘 JACCL 그룹을 초기화한 후 모든 Mac에 걸쳐 집합 분산 합산을 수행합니다 텐서를 JACCL로 직접 처리합니다 MLX 없이 이제 고수준과 저수준 API를 모두 아셨으니 MLX와 JACCL을 사용한 분산 추론과 학습에 대해 MLX로 고급 분산 워크플로우를 구축할 준비가 됩니다
이 세션에서 전체 스택을 살펴봤습니다 분산 학습과 추론을 Apple Silicon에서 가능하게 하는 Thunderbolt를 통한 RDMA부터 MLX와 MLX LM까지 모두 단일 기기에서 여러 기기로 확장하는 것이 얼마나 쉬운지 보여드렸습니다 그것이 가져다주는 이점과 함께 더 빠른 추론, 1조 개 파라미터 모델 실행 가능성 더 빠른 파인튜닝 단일 기기 코드에 최소한의 변경으로 명령줄 인터페이스, Python, Swift, C++ API를 지원합니다 분산 클러스터로 MLX 기반의 로컬 AI 에이전트를 실행할 수 있습니다 빠르고, 비공개이며 여러분이 소유한 하드웨어에서 더 알고 싶다면 WWDC 2026 영상을 확인하세요 "Run local agentic AI on the Mac using MLX" 고급 분산 기능에 대해 더 깊이 알아보려면 커스텀 병렬화 전략과 학습 루프를 포함한 문서를 확인하세요 MLX LM으로 내장 서버를 통해 모델을 분산으로 서빙할 수도 있습니다 Apple Silicon의 MLX로 여러분이 만들 것이 기대됩니다!
-
-
8:31 - Hostfile format for a 4-node MLX cluster
[ { "ssh": "m3-ultra-0", "ips": ["192.168.1.10"], "rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"] }, { "ssh": "m3-ultra-1", "ips": ["192.168.1.11"], "rdma": ["rdma_en5", null, "rdma_en4", "rdma_en3"] }, { "ssh": "m3-ultra-2", "ips": ["192.168.1.12"], "rdma": ["rdma_en5", "rdma_en4", null, "rdma_en3"] }, { "ssh": "m3-ultra-3", "ips": ["192.168.1.13"], "rdma": ["rdma_en5", "rdma_en4", "rdma_en3", null] } ] -
8:56 - Generate the cluster hostfile with mlx.distributed_config
mlx.distributed_config \ --hosts m3-ultra-0,m3-ultra-1,m3-ultra-2,m3-ultra-3 \ --output "m3-ultra-jaccl.json" \ --env MLX_METAL_FAST_SYNCH=1 \ --auto-setup \ --backend jaccl -
11:04 - Run distributed LLM inference with mlx_lm.chat
# Single-device LLM inference mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048 # Distributed LLM inference across the cluster mlx.launch --hostfile "m3-ultra-jaccl.json" -- \ /remote/path/to/mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048 -
15:03 - Run distributed inference with pipeline parallelism
# Tensor parallelism (default) mlx.launch --hostfile "m3-ultra-jaccl.json" -- \ /remote/path/to/mlx_lm.chat --model "moonshotai/Kimi-K2.6" \ --max-tokens 2048 # Pipeline parallelism — append --pipeline flag mlx.launch --hostfile "m3-ultra-jaccl.json" -- \ /remote/path/to/mlx_lm.chat --model "moonshotai/Kimi-K2.6" \ --max-tokens 2048 \ --pipeline -
17:18 - Run distributed fine-tuning with mlx_lm.lora
# Single-device fine-tuning mlx_lm.lora --model "Qwen/Qwen3.5-9B" \ --data "mlx-community/wikisql" \ --train --batch-size 4 # Distributed fine-tuning (scale --batch-size by number of devices) mlx.launch --hostfile "hostfile.json" -- \ /remote/path/to/mlx_lm.lora --model "Qwen/Qwen3.5-9B" \ --data "mlx-community/wikisql" \ --train --batch-size 16 -
19:01 - Distributed inference with the MLX LM Python API
import mlx.core as mx from mlx_lm import stream_generate from mlx_lm.utils import sharded_load # Initialise distributed backend group = mx.distributed.init(strict=True, backend="jaccl") # Define parallelism tensor_group, pipeline_group = group, None # Shard the model model, tokenizer = sharded_load("moonshotai/Kimi-K2.6", pipeline_group, tensor_group) for response in stream_generate(model, tokenizer, prompt, max_tokens=1024): if group.rank() == 0: print(response.text, end="", flush=True) -
19:31 - Shard a layer with the MLX Python API
import mlx.core as mx import mlx.nn as nn # Initialise distributed backend group = mx.distributed.init(strict=True, backend="jaccl") # Define layer and shard it column-wise layer = nn.Linear(1024, 1024) sharded_layer = nn.layers.distributed.shard_linear( layer, strategy="all-to-sharded", group=group ) data = mx.random.normal((1, 1, 1024)) output = sharded_layer(data) mx.eval(output) -
19:47 - All-reduce across devices in Python, Swift, and C++
# Python import mlx.core as mx world = mx.distributed.init(strict=True, backend="jaccl") data = mx.full((4,), float(world.rank()), dtype=mx.float32) result = mx.distributed.all_sum(data, group=world) mx.eval(result) # Swift let group = try DistributedGroup(strict: .ring) let data = rank == 0 ? MLXArray(converting: [1.0, 2.0, 3.0]) : MLXArray(converting: [5.0, 6.0, 7.0]) let result = try group.allSum(data) // C++ namespace mx = mlx::core; auto world = mx::distributed::init(/* strict */ true, "jaccl"); mx::array data = mx::full({4}, static_cast<float>(world.rank()), mx::float32); mx::array result = mx::distributed::all_sum(data, world); mx::eval(result); -
20:06 - Standalone distributed sum with the JACCL C++ API
#include <jaccl/jaccl.h> #include <iostream> int main() { // Initialize JACCL group auto group = jaccl::init(); std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl; // Perform all-reduce sum float data[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; float output[10]; group->all_sum(data, output, sizeof(data), jaccl::Float32); std::cout << "Result: " << output[0] << std::endl; return 0; }
-
-
- 0:00 - Introduction
Overview of why distributed AI becomes necessary as models grow larger, and a preview of what the session covers: CLI tools, Python API, and Swift for embedding distributed workflows in your apps.
- 2:09 - Distributed communication
A walkthrough of the full hardware and software stack enabling distributed workloads on Apple silicon: RDMA over Thunderbolt 5 for low-latency data movement, JACCL (open-source collective communication library), and MLX as the ML framework that ties them together.
- 4:32 - Setting up your cluster
How to physically connect four M3 Ultras into a cluster — understanding latency vs. bandwidth trade-offs, choosing between mesh and ring topologies, enabling RDMA in System Settings, and using mlx.distributed_config and mlx.launch to configure and orchestrate the cluster.
- 10:33 - Distributed inference and fine-tuning
How to run distributed LLM inference with MLX LM using a single CLI command — wrapping mlx_lm.chat with mlx.launch to shard a 27B-parameter Qwen model across four M3 Ultras, achieving nearly 3x the token generation rate of a single machine.
- 13:35 - Model parallelism strategies
How MLX LM splits large models across machines using tensor parallelism (splitting by width for faster inference) and pipeline parallelism (splitting by depth for simpler communication) — including a demo running the 1-trillion-parameter Kimi 2.6 model across four Macs.
- 15:53 - Distributed fine-tuning
How data-parallel training accelerates fine-tuning by replicating the model across machines, processing different data batches in parallel, and averaging gradients — demonstrated fine-tuning Qwen 3.5 (9B) at over 3x throughput on the cluster versus a single M3 Ultra.
- 18:34 - CLI, Python, Swift, and C++ APIs
How to use MLX's fine-grained Python, Swift, and C++ APIs for distributed inference — initializing a distributed group, sharding models with tensor parallelism, using low-level all_reduce primitives, and leveraging JACCL standalone for non-ML distributed workloads.
- 20:45 - Next steps
Summary of the full distributed stack — from RDMA over Thunderbolt to MLX and MLX LM — and next steps including the companion session on local agentic AI, documentation on custom parallelism strategies, and the built-in MLX LM distributed server.