Inferencia Distribuida de LLMs en Mac Studio con MLX y RDMA sobre Thunderbolt

La ejecución eficiente de modelos de lenguaje grandes (LLMs) en hardware de consumo presenta un desafío fundamental en la computación distribuida: cómo gestionar la memoria y el cómputo de modelos que exceden la capacidad de una sola máquina, manteniendo una latencia aceptable. Este artículo aborda este problema demostrando la viabilidad de la inferencia de LLMs de escala de billones de parámetros en un clúster de Mac Studios, aprovechando la arquitectura de memoria unificada de Apple Silicon y una interconexión de baja latencia basada en RDMA sobre Thunderbolt.

La explosión de LLMs ha desplazado gran parte de la investigación y el desarrollo hacia infraestructuras de GPU de alta gama en la nube. Sin embargo, la capacidad de ejecutar estos modelos localmente, en entornos privados y con hardware accesible, es crucial para la privacidad, la iteración rápida y la democratización de la IA. Este trabajo demuestra que, con las optimizaciones de software y hardware adecuadas, incluso modelos de frontera pueden ser operados fuera de los centros de datos tradicionales, abriendo nuevas posibilidades para aplicaciones de IA "agentic" en el edge o en estaciones de trabajo personales.

Históricamente, la distribución de cargas de trabajo de cómputo intensivo ha dependido de redes de alta velocidad como InfiniBand o Ethernet con RoCE. La novedad aquí es la adaptación de principios de paralelización de modelos a una interconexión de consumo (Thunderbolt) con capacidades de RDMA, lo que permite un rendimiento comparable al de infraestructuras más especializadas para ciertos tipos de cargas de trabajo de inferencia.

Arquitectura del Sistema

El sistema se basa en el framework MLX de Apple, diseñado para aprovechar la arquitectura de memoria unificada de Apple Silicon. MLX abstrae la gestión de la memoria entre CPU y GPU, permitiendo que los modelos de LLM residan en una única piscina de memoria accesible por ambos procesadores, eliminando costosas transferencias de datos. Para la inferencia distribuida, MLX se integra con un driver RDMA sobre Thunderbolt 5, proporcionado por macOS 26.2, que permite comunicación de baja latencia entre múltiples Mac Studios. Esta interconexión es gestionada por la librería de colectivos jaccl de MLX.

Se exploran dos esquemas de paralelización de modelos: pipeline parallelization y tensor parallelization. En la paralelización de pipeline, el modelo se divide por capas (profundidad), donde cada Mac Studio ejecuta un subconjunto secuencial de capas y pasa las activaciones intermedias al siguiente nodo. Esto es efectivo para modelos muy grandes que no caben en una sola máquina, pero el throughput de decodificación es inherentemente secuencial. Para mejorar el prefill en contextos grandes, MLX segmenta el contexto en chunks y los procesa en pipeline.

La paralelización de tensor divide cada capa del modelo (ancho) entre múltiples máquinas, requiriendo operaciones de comunicación colectiva (broadcast y aggregation) en cada paso de inferencia. Este esquema es más complejo de implementar pero ofrece mejoras significativas en la velocidad de inferencia (throughput) al paralelizar el cómputo dentro de cada capa. La baja latencia de RDMA sobre Thunderbolt es crítica para que esta estrategia sea efectiva, minimizando el overhead de comunicación. La configuración de red utiliza un "Thunderbolt mesh" con dos cables por enlace para duplicar el ancho de banda. El flag MLX_METAL_FAST_SYNCH=1 optimiza la sincronización CPU-GPU para reducir la latencia de comunicación.

mlx.launch --hostfile hostfile.json --backend jaccl --env MLX_METAL_FAST_SYNCH=1 -- \
mlx_lm.benchmark --model mistralai/Mistral-Medium-3.5-128B -p 8192 -g 128 -b 1 --prefill-step-size 8192 --pipeline

mlx.launch --hostfile hostfile.json --backend jaccl --env MLX_METAL_FAST_SYNCH=1 -- \
mlx_lm.server --model meta-models/Muse-Spark-Open-816B-A42B-NVFP4 --port 8080

{
  "$schema": "https://opencode.ai/config.json",
  "model": "meta-models/Muse-Spark-Open-816B-A42B-NVFP4",
  "provider": {
    "mlx": {
      "npm": "@ai-sdk/openai-compatible",
      "name": "Local MLX",
      "options": {
        "baseURL": "http://127.0.0.1:8080/v1"
      },
      "models": {
        "meta-models/Muse-Spark-Open-816B-A42B-NVFP4": {
          "name": "Muse Spark Open 816B (MLX)"
        }
      }
    }
  }
}

Fundamentos Teóricos

La paralelización de modelos, tanto pipeline como tensor, tiene sus raíces en la investigación de computación de alto rendimiento y redes neuronales distribuidas de las últimas décadas. El concepto de dividir un modelo por capas o por tensores para su ejecución en múltiples dispositivos es fundamental para escalar modelos que exceden la capacidad de memoria o cómputo de una sola unidad. Trabajos seminales como "GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism" (Huang et al., 2018) formalizaron la paralelización de pipeline, mientras que la paralelización de tensor se ha desarrollado en paralelo con la evolución de frameworks de deep learning como TensorFlow y PyTorch, que ofrecen primitivas para operaciones colectivas distribuidas.

La eficiencia de estos esquemas depende críticamente de la latencia y el ancho de banda de la interconexión. La utilización de RDMA (Remote Direct Memory Access) es un principio bien establecido en la computación de alto rendimiento (HPC) para reducir la sobrecarga de la CPU y la latencia en la comunicación entre nodos, permitiendo que las tarjetas de red accedan directamente a la memoria de otras máquinas. La novedad aquí es la aplicación de RDMA sobre una interfaz de consumo como Thunderbolt, lo que extiende las capacidades de HPC a un ecosistema de hardware más accesible, uniendo la teoría de la computación distribuida con la practicidad del hardware de escritorio.