Están agregando soporte para CUDA en Apple MLX

• MLX es un framework de arreglos para investigadores de aprendizaje automático sobre Apple Silicon, con una experiencia de uso similar a NumPy y PyTorch
• Se está trabajando en agregar un backend de CUDA a MLX
• El objetivo principal es ofrecer soporte para memoria unificada (unified memory) y un aprovechamiento más amplio del hardware de NVIDIA
• Por ahora solo funcionan los ejemplos del tutorial, y la compilación y las pruebas se confirmaron en Ubuntu 22.04 + CUDA 11.6
• Mientras se corrigen los problemas iniciales de rendimiento y los cuellos de botella, continúan de forma iterativa las tareas de optimización y refactorización
• El backend de CUDA de MLX se está desarrollando con el respaldo de Apple, y su objetivo es ofrecer una experiencia consistente al desarrollar en Mac y desplegar en entornos de cómputo de gran escala mediante memoria unificada y soporte para hardware de NVIDIA

---

Introducción al framework MLX

• MLX es un framework de arreglos para aprendizaje automático desarrollado por el equipo de investigación de machine learning de Apple
• Fue creado para que ingenieros e investigadores puedan experimentar de forma efectiva y validar ideas rápidamente
• Su rasgo distintivo ha sido el soporte para Apple Silicon, pero ahora busca ampliar su alcance para funcionar también en entornos con GPU de NVIDIA mediante la incorporación de un backend de CUDA

Características principales

• API familiar: la API de Python es similar a NumPy, y los paquetes de nivel superior mlx.nn y mlx.optimizers ofrecen una interfaz parecida a PyTorch
  • También ofrece APIs de C++, C y Swift, manteniendo una experiencia de uso consistente con la API de Python
• Transformaciones de funciones componibles: incluye funciones como autodiferenciación, vectorización automática y optimización de grafos de cómputo
• Cálculo diferido: varias operaciones se agrupan y el resultado se calcula solo cuando realmente se necesita, lo que permite un uso más eficiente de los recursos
• Generación de grafos dinámicos: aunque cambie la shape de los datos de entrada, se ejecuta de inmediato sin compilaciones lentas, lo que facilita el desarrollo y la depuración
• Modelo de memoria unificada y multidispositivo:
  • Antes solo soportaba CPU y GPU de Apple, pero con el backend de CUDA ahora busca dar soporte también a entornos con GPU de NVIDIA
  • Al adoptar una arquitectura de memoria unificada, es posible operar sobre el mismo objeto en todos los dispositivos compatibles sin copiar memoria
• Diseño simple y extensible: permite a los investigadores extender y mejorar el framework con facilidad

Diferencias frente a otros frameworks e inspiración

• Toma inspiración de diseño de NumPy, PyTorch, Jax y ArrayFire, entre otros
• En particular, pone énfasis en el modelo de memoria unificada, una interfaz simple y un entorno de ejecución inmediata/depuración

Ejemplos representativos y áreas de uso

• El repositorio MLX Examples contiene diversos códigos de práctica
  • Entrenamiento de modelos de lenguaje Transformer
  • Generación de texto a gran escala basada en LLaMA y fine-tuning con LoRA
  • Generación de imágenes con Stable Diffusion
  • Casos de implementación de modelos recientes, como reconocimiento de voz basado en OpenAI Whisper

El significado de MLX y el soporte para CUDA

• MLX, que era un framework centrado en Apple Silicon, está ampliando su potencial de uso general en entornos con GPU de NVIDIA mediante la adopción de un backend de CUDA
• Con soporte actualizado para CUDA, se perfila como un framework adecuado tanto para investigación de punta como para uso industrial en hardware de Apple y NVIDIA

[gosarinamu]

si tan solo hicieran posible usar únicamente GPUs de NVIDIA en macOS... je, je.

[yangeok]

Es una buena noticia jaja. ¡Ojalá agreguen compatibilidad con CUDA pronto para que también se pueda entrenar a alta velocidad en Mac!

[GN⁺]

Opiniones en Hacker News

• Quiero confirmar si entendí bien: 1, un programa compilado con MLX puede ejecutarse en chips con soporte para CUDA. Pero 2, un programa CUDA no puede ejecutarse en Apple Silicon. ¿Es correcto que la razón por la que 2 no es posible es que violaría derechos de autor (en particular, la famosa barrera de entrada de Nvidia)?
  • La 1 es correcta. Permite que los desarrolladores trabajen en equipos Apple relativamente modestos, incluidos los de UMA, y luego desplieguen el código terminado en sistemas Nvidia relativamente más potentes. Es útil por varias razones.
  • La 2 no es una violación de derechos de autor. Se puede reimplementar una API.
  • Sobre la 2, no creo que eso sea realmente correcto. HIP de AMD también está haciendo básicamente lo mismo desde que abandonó OpenCL por ahí de 2017 o 2018.
  • Yo quiero la 3. Sería genial poder conectar una GPU de NVIDIA a Apple Silicon y usar CUDA. Me imagino usando Apple Silicon, memoria unificada, GPU y CUDA junto con PyTorch, JAX o TensorFlow, aunque todavía no he probado MLX a fondo.
  • La razón principal por la que la 2 no se hace es que sería mucho más difícil.
• Si pensaste: "Pero la plataforma de Apple no tiene soporte oficial para CUDA", también vale la pena tener en cuenta que este conjunto de parches además soporta "plataformas Linux con CUDA 12 y SM 7.0 (Volta) o superior"
  Guía de instalación
• Este proyecto está liderado por zcbenz, conocido por crear Electron y otras cosas
  Presentación de zcbenz
• Una de las características principales de MLX es el uso de una arquitectura de memoria unificada (UMA), y tengo curiosidad por cómo funciona esto
  Ver el punto correspondiente en el readme del repositorio: GitHub de MLX
  Creo que sería interesante extenderlo a todos los APU con UMA, sin importar el fabricante, pero me pregunto si en las GPU discretas el enfoque tendría que ser diferente
  Leyendo los comentarios del PR, parece que CUDA también soporta directamente la API de UMA y maneja las copias de forma transparente cuando hace falta
  • En mi experiencia, la falta de capacidades de prefetch provoca cuellos de botella graves de memoria por la espera de copias. Si todo el dataset cabe en VRAM, hacer prefetch manual funciona bien, pero en mi aplicación, entrenamiento de ML, la caída de rendimiento era tan grande que terminé cambiando a carga por streaming.
• Cambiando un poco de tema, parece que muchos de quienes trabajan en MLX no están oficialmente afiliados a Apple. Por ejemplo, véase el Twitter de prince_canuma
  Esto me hace preguntarme si Apple está apoyando el proyecto por detrás sin hacerlo oficial. Incluso hubo rumores de que el equipo se iba a mover.
• En la descripción del PR dice "ongoing effort to add a CUDA backend"
  O sea, significa que el código de MLX puede compilarse y ejecutarse en hardware x86 + GeForce, no al revés.
• Si la idea es que "puedes escribir y probar código en una Mac y desplegarlo en una supercomputadora para una buena experiencia de desarrollo", entonces surge la duda de si ya se puede usar MLX en Linux
  Lo probé directamente y solo la versión para Python 3.12 está publicada en PyPI
  Página de MLX-CUDA en PYPI
• Me pregunto si esto tendrá que ver con Strix Halo. Como me sobraba presupuesto para la computadora del trabajo, compré una EVO-X2, y está al nivel de máquinas EPYC znver5 de gama media o incluso mejor. Supera por mucho a la mayoría de las instancias de EC2 o GCE, incluso considerando NVMe. Este hardware cuesta 1800
  Últimamente también están saliendo reseñas en YouTube diciendo que DGX Spark ofrece peor relación precio-rendimiento que Strix Halo. El punto débil es que ROCm todavía está verde, aunque parece que mejorará pronto. Si CUDA funcionara perfectamente en equipos Apple, sería una opción para pensar seriamente, por más barato y bueno que sea Strix
  • Strix Halo es básicamente lo mismo que AMD Ryzen AI Max+ 395. Lo están usando en Framework Desktop y en varias mini PC. El ancho de banda de memoria de este chip es de 200 GB/s, muy bueno para una plataforma x86 común, pero se queda muy atrás frente a una GPU Nvidia, por ejemplo una 5090 con 1792 GB/s, o un Apple M3 Ultra con 800 GB/s. La relación precio-rendimiento es excelente, pero para tareas intensivas en memoria como los LLM siento que apenas entra en el umbral mínimo de rendimiento.
  • Por la descripción del PR, queda claro que el objetivo es el entrenamiento en clústeres en la nube.
  • Llevamos 10 años escuchando que el cómputo para redes neuronales en GPU de AMD pronto competirá con Nvidia, y sigue sin hacerse realidad.
  • Me pregunto cómo se compara con una Mac mini M4.
• Los modelos de MLX me impresionan muchísimo. Puedes dejar modelos locales abiertos para toda la familia sin preocuparte por el riesgo de incendio como con una computadora Nvidia. Ojalá Apple Silicon se convierta en un competidor serio del ecosistema de chips de Nvidia. También me pregunto si el soporte para CUDA podría ser una estrategia de embrace, extend, extinguish (EEE).
• Apple planea en el futuro construir centros de datos basados en chips de la serie M para usarlos en desarrollo de apps, pruebas y hospedaje de servicios externos.