Comienza a usar MLX para el chip de Apple

Hola, soy Awni. Hoy tengo el placer de presentarte MLX. MLX es una estructura de arreglo de código abierto diseñada específicamente para el chip de Apple. Se usa para realizar desde cálculos básicos hasta aprendizaje automático en dispositivos Apple. Ideal para generar texto con modelos de lenguaje grandes, generar imágenes, audio o videos con los últimos modelos. Puedes usarlo para entrenar, ajustar o personalizar los modelos de aprendizaje automático desde la Mac. Te contaré un poco sobre qué es MLX y dónde es útil. También cubriré los conceptos básicos del uso de MLX en Python, la instalación y las operaciones de arreglo. Después, te contaré algunas de las funciones clave que diferencian a MLX de otras estructuras. Veremos algunas de las herramientas de MLX para hacer que el aprendizaje automático sea más rápido en el chip de Apple.

En la última sección, te presentaré brevemente MLX Swift y te mostraré cómo empezar a utilizarlo. Vamos a sumergirnos en una introducción a MLX. MLX se creó desde cero para ser rápido en el chip de Apple, donde puede ejecutarse en el CPU o acelerarse en el GPU. Puedes utilizar MLX para muchísimas aplicaciones, desde operaciones de pequeña escala hasta de gran escala. Está diseñado para ser fácil de usar y flexible, sin comprometer la velocidad ni la eficiencia.

MLX tiene una API central que sigue de cerca a NumPy. A menudo, puedes usarlo como reemplazo directo para acelerar la mayoría de los cálculos numéricos. MLX también tiene todas las herramientas necesarias para el aprendizaje automático. Las API de alto nivel son similares a PyTorch y JAX. Si conoces alguna de esas estructuras, MLX te resultará familiar y será aún más fácil de usar. Puedes usar MLX para un aprendizaje automático más avanzado directamente en el dispositivo. Por ejemplo, LM Studio, una popular app para generar texto con grandes modelos de lenguaje directo en Mac, usa MLX.

Puedes utilizar el paquete MLX LM, creado sobre MLX, para generar texto y ajustar modelos de miles de parámetros. Para obtener más información, consulta la sesión “Explore large language models on Apple silicon with MLX”. MLX tiene una API de Python con todas las funciones, que resulta útil para la creación rápida de prototipos. También tiene una API en Swift, que incluye todos los paquetes para crear y optimizar redes neuronales. MLX también tiene API en C++ y C. Puedes usar MLX en uno de estos lenguajes para ejecutar los modelos de aprendizaje automático en el chip de Apple. El software MLX es de código abierto y cuenta con una licencia MIT permisiva. El software principal está disponible en GitHub, junto con ejemplos y paquetes creados con API de Python y Swift. MLX también tiene una comunidad activa de creadores de modelos en Hugging Face. Muchos modelos recientes ya están en el Hugging Face de la comunidad MLX, y cada día se cargan nuevos modelos. La forma más fácil de comenzar a usar MLX en Python es instalarlo desde PyPi. Solo debes ejecutar una línea en la terminal, pip3 install mlx. MLX es fácil de usar. Para empezar a realizar cálculos en arreglos, solo abre un archivo Python e importa MLX. Luego, puedes crear algunos arreglos y realizar operaciones básicas en ellos. Por ejemplo, sumemos dos arreglos de enteros.

También puedes ver fácilmente la información sobre un arreglo, como su forma y el tipo de datos. La API de Python MLX es similar a NumPy. Las operaciones normalmente tienen los mismos nombres y firmas, y se comportan de la misma manera. Si vienes de NumPy o una estructura similar, MLX te resultará familiar y será fácil comenzar a usarlo. Ahora que ya sabes qué es MLX, dónde es útil y algunos conceptos básicos, veamos las funciones clave. Memoria unificada, evaluación perezosa, transformación de funciones y paquetes para optimizar redes neuronales. MLX está diseñado para aprovechar lo mejor del chip de Apple. Esto incluye un nuevo modelo de programación específico para la memoria unificada.

La mayoría de los sistemas para el aprendizaje automático tienen un GPU discreto con memoria independiente. El chip de Apple, en cambio, tiene una arquitectura de memoria unificada. Esto significa que el CPU y el GPU comparten la misma memoria física.

Para trabajar con memoria unificada, MLX es diferente de lo que quizá conozcas de las estructuras tradicionales. El cálculo sigue a los datos en las estructuras tradicionales. El cálculo se realiza en el CPU si el arreglo está en la memoria del CPU. El cálculo se realiza en el GPU si el arreglo está en la memoria del GPU.

En MLX, los arreglos se asignan en la memoria unificada. Nunca necesitarás copiarlos en ningún lugar para usarlos en todos los dispositivos compatibles.

En cambio, para ejecutar una operación en un dispositivo, debes especificar el dispositivo para la operación en sí. Por ejemplo, aquí estamos sumando a y b en el GPU y multiplicando a y b en el CPU. Estas operaciones incluso pueden ejecutarse en paralelo y MLX administrará automáticamente las dependencias. Otra función clave de MLX es el cálculo perezoso. Para que MLX sea lo más eficiente posible, cuenta con un motor de ejecución perezosa.

Cuando se llama a una operación como una suma en dos arreglos, no se produce ningún cálculo real. Sino que se crea un gráfico de cálculo, como este.

El arreglo C aún no se calculó. Solo se calcula si realmente necesitas usarlo. Por ejemplo, imprimir C o convertir C desde MLX a una lista de Python forzará su cálculo. Puedes forzar explícitamente la evaluación del gráfico usando mx.eval. El cálculo perezoso tiene varias ventajas. Al disociar la construcción y la ejecución del gráfico, MLX puede realizar acciones antes de calcular resultados. Con el cálculo perezoso solo pagas por lo que usas. La transformación de funciones es otra de las funciones clave de MLX. Lo elevan de una estructura de arreglo rápido a una herramienta mucho más potente.

Las transformaciones de funciones son transformaciones que toman funciones de entrada y devuelven otras de salida.

MLX tiene varias transformaciones de funciones. Por lo general, entran en una de dos categorías. Las transformaciones para la diferenciación automática o transformaciones para optimizar gráficos de cálculo. Por ejemplo, puedes usar una transformación de función para el cálculo automático del gradiente de una función.

Supongamos que tienes una función básica que calcula el seno de su entrada. Para tomar el gradiente de esta función, puedes usar la transformación de función mx.grad. La salida de mx.grad es una nueva función que, cuando la llamas en un arreglo, te da la derivada. Las transformaciones de funciones son componibles arbitrariamente. Puedes tomar la segunda derivada solo usando mx.grad en la salida de mx.grad.

El resultado es otra función que, cuando se llama en un arreglo, proporciona la segunda derivada del signo.

MLX tiene dos paquetes de nivel superior para facilitar las acciones con redes neuronales.

El primero es mlx.nn, una biblioteca modular que se utiliza para construir redes neuronales. El segundo es mlx.optimizers, una biblioteca de algoritmos de optimización comunes. Los dos paquetes pueden usarse independientemente o también pueden integrarse perfectamente entre sí.

El paquete mlx.nn tiene toda la funcionalidad que necesitas para construir redes neuronales. nn.module es la clase base principal de la que heredan todas las capas y contenedores. Expone métodos útiles para acceder a parámetros, cargar parámetros, guardar parámetros y mucho más.

La biblioteca nn también viene con un conjunto de capas estándar listas para usar, pero también puedes crearlas.

Las funciones de pérdida y los métodos de inicialización comunes se incluyen en los subpaquetes nn.losses y nn.init.

Veamos cómo se puede construir una red neuronal multicapa simple con mlx.nn.

El primer paso es crear una clase personalizada que herede de nn.Module. En este caso, utilizaremos una red neuronal simple de una capa oculta. Creamos las capas lineales dentro del método de inicialización del módulo. Luego implementamos la función de llamada, que calcula la salida del módulo en la entrada dada. La función de llamada llama a la primera capa lineal, aplica la función de activación relu y llama a la segunda capa. El paquete de red neuronal de nivel superior está diseñado para parecerse a estructuras de aprendizaje automático.

Comparemos la implementación del modelo MLX con la misma implementación en PyTorch. Son casi idénticos, con solo dos pequeñas diferencias en la función que calcula la salida. Si ya creaste modelos en PyTorch, cambiar a MLX debería ser muy sencillo. Ahora veamos el uso de MLX para que las cargas de trabajo de aprendizaje automático sean lo más eficientes posible. Empezaré mostrándote cómo compilar funciones para acelerarlas. Luego, hablaré del subpaquete mx.fast y una API para agregar tus Metal Kernels personalizados. Después de eso, mostraré cómo usar la cuantificación para reducir el uso de memoria y acelerar la ejecución. Por último, mostraré cómo usar MLX para distribuir un cálculo entre muchas máquinas. Casi todos los cálculos realistas de MLX constarán de funciones que realizan varias operaciones en arreglos. Una forma sencilla de hacer que funciones como esa sean más rápidas es con la transformación de mx.compile. Supongamos que tienes una función que realiza algunas operaciones elemento por elemento, como la función GELU. El gráfico de cálculo de esta función tiene varios nodos. Cada uno de estos nodos corresponde a un lanzamiento de kernel de GPU separado.

La compilación del gráfico utiliza todos estos núcleos separados en un único núcleo fusionado. Esto puede ahorrar ancho de banda de memoria y sobrecarga de ejecución de gráficos y lograr un cálculo más eficiente.

Usar mx.compile es tan fácil como decorar la función con la transformación de función mx.compile. Para operaciones más complejas, es probable que puedan ejecutarse más rápido usando mx.fast. Por ejemplo, muchos de los componentes básicos de un modelo de transformador usan operaciones en mx.fast. Estos incluyen codificaciones posicionales, capas de normalización y atención de producto puntual de escala. Las operaciones en mx.fast son más especializadas y están altamente optimizadas para ser lo más rápidas posible. También son altamente configurables, por lo que pueden acelerar muchas variaciones de un cálculo determinado. La operación de atención escalada por producto puede tomar una máscara opcional como entrada. La máscara puede ser aditiva, booleana o una cadena que indica el tipo de máscara. Veamos más de cerca a la norma RMS, que es una de las operaciones en mx.fast. La norma RMS se utiliza en casi todos los modelos de lenguaje grandes basados en transformadores modernos. Una implementación básica que usa operaciones MLX da como resultado un gráfico de cálculo grande.

Pero podemos reemplazar toda la implementación con una sola llamada a mx.fast.rms_norm. El código es más simple, el gráfico de cálculo solo tiene un único nodo y el cálculo en sí será mucho más rápido.

MLX tiene una API para agregar Metal Kernel personalizados para cuando convenga una implementación más personalizada. Escribes el kernel personalizado y MLX hace el resto, incluida la compilación y la ejecución justo a tiempo. Estos kernels están escritos en Metal, que es un lenguaje y API de Apple para ejecutar funciones en los GPU de Apple. Al construir el kernel, se pasa una cadena fuente de código Metal, así como información sobre las entradas y salidas.

Se llama al kernel especificando el tamaño de la cuadrícula y las formas y tipos de la salida. MLX trata la llamada al kernel de la misma manera que cualquier otra operación. Crea un nodo en el gráfico de cálculo y se realiza su evaluación de forma perezosa.

Otra herramienta de nuestro kit de herramientas es la cuantificación. Los modelos grandes necesitan mucha memoria y mucho ancho de banda de memoria para ser rápidos. La precisión necesaria para el entrenamiento es mucho mayor que para la inferencia para casi la misma calidad. Reducir la precisión te permite colocar modelos más grandes en la memoria y ejecutarlos más rápido.

Si el modelo está en precisión de punto flotante de 32 bits, primero puedes reducir los requisitos de memoria a la mitad. Si 16 bits son demasiados, MLX tiene rutinas integradas para cuantificar arreglos para que sean aún más pequeños. Por ejemplo, puedes cuantificar usando 4 bits por elemento para reducir más los requisitos de memoria.

Para cuantificar una matriz, usa mx.quantize. Especifica el número de bits a usar por elemento y el tamaño del grupo. Este determina la cantidad de elementos en la matriz cuantificada que comparten un valor de escala y punto cero. Cuanto más pequeños sean los bits y mayor el tamaño del grupo, más pequeño y rápido será el resultado. MLX tiene varias opciones de bits y tamaños de grupos para brindarte la mayor flexibilidad posible.

Puedes multiplicar vectores o matrices de entrada no cuantificada utilizando mx.quantized_matmul.

Usa mx.dequantize para recuperar una aproximación a la entrada original.

MLX.mn también tiene una utilidad práctica para cuantificar un módulo en un solo comando. Digamos que tienes un modelo que es una pila de una capa de incrustación más varias capas lineales. Puedes cuantificar el modelo con nn.quantize. El comando de cuantificación acepta una devolución de llamada opcional para que tengas un control más preciso.

En modelos de lenguaje grandes, la cuantificación reduce mucho el uso de memoria y aumenta los tokens por segundo. A veces, una sola máquina no es suficiente. Por ejemplo, es posible que desees generar texto con un modelo grande que no cabe en la memoria. O puede que estés realizando tareas fáciles de paralelizar y pueden ser mucho más rápidas si se usan varias máquinas.

MLX viene de fábrica con la capacidad de ejecutar cálculos arbitrarios en varias máquinas. Las máquinas se pueden conectar a través de Ethernet o Thunderbolt.

Usa el subpaquete mx.distributed para distribuir cálculos entre varias máquinas. El mx.distributed es principalmente un conjunto de operaciones de comunicación.

Por ejemplo, all_Sum agrega los arreglos de entrada de todas las máquinas. La salida de all_Sum es la entrada sumada, que es la misma para cada máquina. Veamos cómo sumar un arreglo en varias máquinas. Inicia el backend distribuido con mx.distributed.init. Este paso es opcional, pero debes hacerlo si necesitas acceder al grupo de comunicación. El grupo de comunicación tiene información útil adjunta, como número total de procesos e índice de proceso actual.

Crea un arreglo con un solo valor en cada proceso y llama a mx.distributed.all_sum para sumar los arreglos.

MLX tiene un iniciador práctico para ejecutar un programa MLX en varias máquinas. Para ejecutar un programa en 4 máquinas, utiliza mlx.launch con las 4 direcciones IP del host. Hasta ahora, solo estuvimos usando MLX en Python. En muchos casos, es posible que prefieras la facilidad y flexibilidad de Python. En otros casos, es posible que prefieras Swift. Por esa razón, MLX tiene una API completa en Swift. Está creada con base en Metal y funciona muy bien en macOS, iOS, iPadOS, visionOS y más. Comenzar a utilizar MLX en Swift es tan fácil como agregarlo como un paquete al proyecto Xcode. Haz clic en el proyecto y, a continuación, en el signo más en la pestaña de dependencias del paquete. Luego, ingresa la URL del repositorio GitHub de MLX Swift y haz clic en Add Package. Eso es todo lo que se necesita para comenzar a crear con MLX Swift. Para facilitar la transición entre Python y Swift, las API de ambos son intencionalmente similares. A continuación, se muestra una comparación lado a lado de un fragmento de código Python y su contraparte MLX Swift. Crear arreglos, realizar operaciones e inspeccionar metadatos es casi lo mismo en Swift que en Python. Todas las funciones principales de MLX en Python se aplican igualmente a MLX Swift. Vimos una descripción general amplia de muchas de las funciones clave de MLX. Consulta el sitio web de MLX para obtener más información sobre la estructura. Las API de Python y Swift tienen un repositorio de ejemplos, que contiene muchos ejemplos de aprendizaje automático. Estos ejemplos son una gran forma de aprender sobre MLX y excelentes puntos de partida para desarrollar con MLX. Gracias por acompañarnos. Nos entusiasma ver lo que crearás con MLX.