Explora grandes modelos de lenguaje en el chip de Apple con MLX

Hola, soy Angelos, ingeniero del equipo MLX. Hoy te mostraré por qué MLX es perfecto para grandes modelos de lenguaje en el chip de Apple. Con él, puedes realizar inferencias y ajustar modelos masivos desde tu Mac. Todo esto con apps CLI o desde Python o Swift. Si hace poco usas MLX, es una biblioteca de código abierto diseñada para el aprendizaje automático en el chip de Apple. Usa Metal para la aceleración en el GPU y aprovecha la memoria unificada para que las operaciones en el CPU y GPU funcionen con los mismos datos simultáneamente. Puedes usar MLX en tu lenguaje preferido, ya que proporciona API en Python, Swift, C++ e incluso C. Para más información, consulta nuestra sesión sobre MLX para el chip de Apple. Para ejecutar grandes modelos de lenguaje en el chip de Apple, MLX ofrece nuevas y potentes capacidades que te permiten ejecutar modelos de vanguardia en la Mac con un comando de una sola línea. Carguemos el último modelo de IA de DeepSeek, con sus impresionantes 670.000 millones de parámetros.

Incluso al cuantificarlos a 4.5 bits por peso, los pesos del modelo por sí solos necesitan unos 380 GB de memoria. Para manejar eso, usamos un chip M3 Ultra con sus 512 GB de memoria unificada, ningún otro dispositivo de consumo se le acerca. El modelo está cargado, así que podemos empezar a interactuar con él. Podemos preguntarle: “¿Cuál es el lago más profundo de Estados Unidos?”,

o pedirle que escriba código.

MLX permite interacción y generación fluidas en tiempo real a velocidades superiores a las de lectura, incluso con modelos con cientos de miles de millones de parámetros ejecutándose en el escritorio de tu Mac. Ya ya viste lo que es posible, veamos ahora cómo puedes usar MLX para ejecutar estos potentes modelos en tu propia Mac. Presentaremos MLX LM, una biblioteca de Python y un conjunto de apps de línea de comandos que satisfacen todos los requisitos de tu gran modelo de lenguaje con una solución robusta y versátil para diversas apps.

Luego, profundizaremos en la generación de texto con MLX LM y mostraremos lo fácil que es generar texto desde Python o desde la terminal. Además, analizaremos la descarga de modelos de Hugging Face y los cuantificaremos para una inferencia más rápida en el dispositivo.

MLX puede hacer mucho más que una mera inferencia. Luego, usaremos MLX LM para ajustar un modelo con nuestros propios datos. Entrenaremos un adaptador de bajo rango, que podremos fusionar en el modelo para un uso más sencillo y una inferencia más rápida. Finalmente, analizaremos el uso de MLX de Swift para integrar un gran modelo de lenguaje en tu app de Swift con solo unas pocas líneas de código.

La forma más sencilla de empezar a usar modelos de lenguaje en MLX es con MLX LM: un paquete de Python creado sobre MLX, diseñado para ejecutar y experimentar con grandes modelos de lenguaje. Proporciona un conjunto de herramientas de línea de comandos para generar texto o ajustar modelos sin escribir código. Para más control, también proporciona una API de Python para personalizar la generación o el proceso de entrenamiento tanto como quieras. También está integrado con Hugging Face. Así que puedes descargar rápidamente miles de modelos de internet e incluso cargar los tuyos para compartirlos con la comunidad.

Empezar es fácil: basta con ejecutar pip install mlx-lm.

Veamos los detalles del caso de uso más común de los modelos de lenguaje: generar de texto.

Esta herramienta de línea de comandos genera texto con un modelo de lenguaje desde tu terminal, sin necesidad de código. Así funciona: le das un modelo de Hugging Face o una ruta local, una solicitud de texto y se encarga del resto. Descarga el modelo si es necesario, ejecuta el símbolo del sistema e imprime la respuesta generada. Así que, ejecutemos este comando.

Obtenemos una implementación de Swift Quick Sort en segundos.

Puedes modificar el comportamiento del modelo con indicadores como temperatura, TopP o tokens máximos, igual que con cualquier generador de texto estándar. Para conocer todas las opciones disponibles, puedes ejecutar mlx_lm.generate --help. Ya sea que estés creando prototipos, generando código o explorando el modelo, este es el punto de partida más sencillo. Ya vimos lo fácil que es generar texto desde la línea de comandos con mlx_lm.generate. Pero la potencia de MLX LM no se limita a las herramientas de terminal. También proporciona una API de Python clara y flexible, perfecta para un control más preciso o para la generación en un flujo de trabajo más amplio. Veamos cómo hacer lo mismo, generar texto, con solo unas pocas líneas de Python.

Primero, importamos cargar y generar. Cargar, maneja todo lo relacionado con la carga del modelo. Obtiene el modelo solicitado desde el disco local o desde Hugging Face y configura el objeto del modelo junto con el identificador. Luego invocamos “generate”. La cual realiza un bucle de generación de tokens y devuelve el texto de salida para procesar en Python, registrar o alimentar en otros sistemas.

Así, con solo cargar y generar, obtenemos la misma funcionalidad que la CLI, pero con control y flexibilidad totales en Python. Veamos otra de sus características importantes. El modelo que obtienes no es un objeto opaco con el que solo puedes interactuar a través de una interfaz fija. Es una red neuronal MLX estructurada que puedes inspeccionar, explorar su arquitectura e incluso modificarla. Te haré una demostración rápida.

Podemos empezar con la lista de capas que componen el modelo.

Esto nos da un desglose del transformador, capa por capa. También podemos ver los parámetros, básicamente las ponderaciones y los sesgos que el modelo aprendió.

También podemos profundizar en una parte específica de la red, como el módulo de autoatención en la primera capa.

Este nivel de transparencia es muy útil, no solo para la depuración o el aprendizaje, sino para el intercambio de capas, las rutinas de ajuste personalizado o la operación de modelos de bajo nivel.

Ya vimos cómo generar texto a partir de una sola solicitud. Pero ¿qué pasa si quieres mantener una conversación o generar respuestas por turnos donde cada nuevo mensaje se basa en el anterior? Ahí entra en juego la caché de valores clave o KV. Los modelos de lenguaje usan mecanismos de atención para procesar por identificadores de entrada y calculan la atención sobre los identificadores previos. Esto puede ser costoso para solicitudes largas o con varios turnos. La caché de KV lo soluciona al almacenar los resultados intermedios previos, en concreto claves y valores.

En lugar de volver a calcular todo, el modelo reutiliza esta caché y ahorra tiempo y cálculos. En MLX LM, usar una caché de KV es sencillo.

Actualicemos el ejemplo anterior de Python con una caché KV creada para reutilizarse en varias generaciones.

Primero creamos el objeto de caché con la función make_prompt_cache. Podemos usarlo para editar el historial, guardarlo para un uso posterior cambiar de conversación sin problemas.

Luego, lo pasamos a la función generar. Al generarse nuevos identificadores, la caché se actualiza. Cada llamada continúa desde donde se interrumpió la última, manteniendo el contexto. Esto es muy útil al crear chatbots, asistentes virtuales o apps interactivas donde el seguimiento del historial sea importante.

Ahora, hablemos de la cuantificación de los modelos. Vimos cómo generar texto y trabajar con modelos de forma interactiva. Al implementarlo en el mundo real, la eficiencia importa tanto como la funcionalidad. Los modelos se publican con la misma precisión con la que se entrenaron, como float32 o float16. Eso los hace grandes y lentos, especialmente en dispositivos más pequeños. Ahí entra en juego la cuantificación.

Reduce la precisión del modelo a Int8 o incluso a 4 bits, lo cual usa menos memoria y acelera la inferencia con poco impacto en la calidad. Esto implica herramientas adicionales, scripts de conversión y problemas de compatibilidad. En MLX es más sencillo. La cuantificación está integrada. Puedes comprimir modelos en varios niveles y usarlos para inferencias o entrenamiento sin configuración adicional. Veamos cómo funciona.

Para cuantificar o convertir un modelo con MLX, se usa el comando mlx_lm.convert. Esta herramienta descarga un modelo de Hugging Face, lo convierte a una precisión diferente y lo guarda localmente, todo en un paso. En este ejemplo, obtenemos el modelo Mistral original de 16 bits y lo cuantificamos a unos 4 bits por peso.

El resultado es mucho más pequeño, más rápido de ejecutar y requiere menos memoria. Una vez convertido, el modelo se guarda y se puede usar inmediatamente para inferencias o entrenamiento con las mismas herramientas MLX LM.

Y si quieres compartir tu modelo de cuantificación, puedes volver a cargarlo fácilmente en Hugging Face con un nombre de repositorio. Ya sea que estés optimizando la velocidad, ahorrando espacio o contribuyendo a la comunidad, solo necesitas este comando.

Como en la generación de texto, usar la API de Python para convertir y cuantificar modelos ofrece más flexibilidad sin más complejidad. MLX LM facilita la aplicación de diferentes ajustes de cuantificación a diferentes partes del modelo o desde Python.

Por ejemplo, es habitual mantener las capas de incrustación y proyección final con mayor precisión, ya que suelen ser más sensibles a la cuantificación. En este ejemplo, cuantificamos esas capas a 6 bits y el resto del modelo usa 4 bits, logrando un excelente equilibrio entre calidad y eficiencia. Esto se hace pasando una función de predicado de cuantificación que recibe cada capa y devuelve los parámetros que se deben usar para ella. Lo demás funciona exactamente igual. Invocamos “convert”, pasamos la ruta de Hugging Face y el directorio de salida, MLX se encarga del resto, incluida la descarga del modelo y el almacenamiento del resultado. Este control detallado es muy útil cuando experimentas con la compresión de modelos o intentas encontrar el mejor equilibrio entre rendimiento y precisión.

Ya vimos cómo generar texto con grandes modelos y cómo cuantificarlos para una inferencia más rápida y una implementación más ligera. Pero MLX puede hacer más, sobre todo en entrenamiento. Puedes ajustar un gran modelo de lenguaje con tus propios datos en la Mac y, lo que es más importante, sin que esos datos salgan del dispositivo. Lo mejor es que puedes hacerlo sin escribir una sola línea de código. Veamos cómo funciona el ajuste fino.

Los modelos suelen entrenarse en conjuntos de datos masivos de uso general procedentes de Internet. Eso les da un conocimiento amplio, pero pueden carecer de profundidad en dominios especializados o no captar detalles de una tarea concreta. Con el ajuste fino adaptamos los modelos a nuevos contextos. Al entrenarlos en un conjunto de datos más específico, les damos nuevas capacidades o adaptamos las respuestas a necesidades particulares. Este proceso suele realizarse en la nube, lo que puede ser costoso y no es ideal cuando se trabaja con datos privados o confidenciales. Con MLX, puedes ajustar de forma precisa estos modelos localmente en la Mac sin tener que usar los datos en la nube. Es eficiente, seguro y se integra perfecto en el flujo de trabajo de MLX.

MLX LM admite dos tipos de ajuste fino desde el inicio: de modelo completo y entrenamiento de bajo rango. En el completo, actualizamos todos los parámetros del modelo entrenado. Esto te da máxima flexibilidad, pero también requiere más recursos. El entrenamiento de adaptadores de bajo rango, agrega una pequeña cantidad de nuevos parámetros al modelo y entrena solo esos, manteniendo la red original bloqueada. Esto da un entrenamiento más rápido, ligero y con una memoria más eficiente en hardware local. Apliquemos esto al ajustar el modelo de Mistral en un conjunto de datos personalizado. Veamos lo fácil que es iniciar un trabajo de ajuste con MLX LM. Solo se necesita un comando y algunos argumentos clave. Especificamos el modelo, la ruta al conjunto de datos y cuánto tiempo queremos entrenarlo. Como la cuantificación está integrada en MLX, el comando mlx_lm.lora puede entrenar adaptadores sobre modelos cuantificados. Esto reduce drásticamente el uso de memoria sin sacrificar la capacidad de realizar ajustes.

Aquí, estamos entrenando en una versión cuantificada de 4 bits de Mistral, que reduce el uso de memoria unas 3.5 veces en comparación con la versión de precisión completa. Así, incluso con grandes modelos, el ajuste sigue siendo práctico y eficaz en la Mac. Ese comando de una sola línea es perfecto para un entrenamiento rápido cuando apenas comienzas. Si quieres ajustar el rendimiento, necesitarás más control sobre el proceso de entrenamiento. Ahí interviene el archivo de configuración de entrenamiento. MLX LM admite archivos de configuración que brindan control en el tamaño de la ruta, la velocidad de aprendizaje, los ajustes del optimizador, los intervalos de evaluación y más. Esto te permite adaptar la configuración a tus datos, hardware u objetivos de optimización y aprovechar al máximo tu adaptador. Veamos el ajuste en acción y cómo actualizar los conocimientos de un modelo. Preguntemos a Mistral 7b quién ganó el último Super Bowl.

La respuesta es correcta, pero está desactualizada. El límite del modelo está en que no tiene acceso a eventos recientes. Pero con el ajuste fino que podemos solucionarlo en unos minutos. Al entrenar en un conjunto de datos con preguntas y respuestas sobre el último Super Bowl, podemos actualizarlo y hacer que responda con precisión.

Después de unos minutos, el modelo puede dar respuestas actualizadas sobre equipos, jugadores, resultados y más.

Ahora podemos usar MLX LM para fusionar los adaptadores entrenados en el modelo base. Esto es muy útil para la implementación y el uso compartido porque produce un único modelo independiente fácil de distribuir y usar.

La fusión combina el adaptador con los pesos originales, lo que resulta en un modelo con la misma arquitectura y número de parámetros que la versión previa, pero con capacidades actualizadas. Así que se comporta como cualquier otro modelo, pero con ese conocimiento integrado.

Para esta fusión, usamos el comando mlx_lm.fuse. Calcula los pesos fusionados y guarda los resultados en la ruta en un solo paso. No es necesario cuantificar nada manualmente. MLX lo hace automáticamente y conserva la misma cuantificación del entrenamiento. Y si quieres compartir tu modelo recién perfeccionado, es igual de sencillo. Solo proporciona un repositorio de Huggin Face y el modelo fusionado se cargará y estará listo para usar. Ya usamos Python para generar texto y ajustar grandes modelos de lenguaje. Pero otra funcionalidad de MLX es que ofrece la misma sencillez y flexibilidad a Swift. Veamos lo fácil que es usar un gran modelo de lenguaje en Swift con MLX.

Este es un ejemplo de cómo cargar un modelo Mistral cuantificado y generar texto desde Swift. Y todo en tan solo 28 líneas de código. Importamos MLX y las bibliotecas de modelos de lenguaje. Creamos un contenedor de modelos que gestiona de forma segura el acceso al modelo y al identificador. Preparamos la entrada. Convertimos el texto de entrada al formato numérico que entiende el modelo. Por último, ejecutamos el bucle de generación e imprimimos el resultado, como en Python. Es el mismo flujo de trabajo y las mismas capacidades, pero en Swift. Veamos qué se necesita para conservar el historial de una conversación en varias interacciones con un modelo, como hicimos en Python. En Swift, esto requiere unas líneas más. La idea es la misma, necesitamos crear una caché de valores clave para reutilizarla en varias generaciones. Esto se hace con una línea adicional de código. Sin más complejidad. Para administrar la interacción de forma más precisa, usamos un iterador que establece el valor clave directamente y controla las generaciones paso a paso. Esto nos da la flexibilidad de administrar conversaciones de varios turnos y entradas avanzadas desde Swift. En esta sesión, vimos lo sencillo que es realizar inferencias, entrenamiento y cuantificaciones con MLX, ya sea mediante código o comandos de terminal. Las API del modelo de lenguaje de nivel superior que usamos y los kernels de Metal que las impulsan, son de código abierto. MLX proporciona operaciones básicas en C, C++, Python y Swift, con API de alto nivel en Python y Swift, lo que te da flexibilidad y control en toda la pila. Lo que lo hace especialmente potente para ejecutar modelos de lenguaje y flujos de aprendizaje automático en hardware de Apple. Ahora veamos qué sigue. Exploramos algunas funcionalidades clave de MLX LM, pero puedes hacer mucho más. Nuestra documentación profundiza en inferencia distribuida y entrenamiento, cuantificación aprendida y bucles de entrenamiento personalizados. Los repositorios MLX y MLX Swift ofrecen proyectos listos para ejecutar para tareas como generación de imágenes con modelos de difusión, reconocimiento de voz y entrenamiento de modelo completo. Ya sea que estés creando tu propia app con IA o explorando todas las posibilidades, lo que necesitas está a solo unos clics. Queremos ver las increíbles experiencias que crearás en el hardware de Apple con MLX y la potencia de los grandes modelos de lenguaje.