En diciembre de 2023, Apple lanzó su nuevo Marco de aprendizaje profundo MLX, un marco de matriz para el aprendizaje automático en el silicio de Apple, desarrollado por su equipo de investigación de aprendizaje automático. Este tutorial explorará el marco y demostrará la implementación del modelo Mistral-7B localmente en una MacBook Pro (MBP). Configuraremos una interfaz de chat local para interactuar con el modelo implementado y probaremos su rendimiento de inferencia en términos de tokens generados por segundo. Además, profundizaremos en la API de MLX para comprender las palancas disponibles para alterar el comportamiento del modelo e influir en el texto generado.
Como es habitual, el código está disponible en un repositorio público de GitHub: https://github.com/marshmellow77/mlx-deep-dive
El nuevo marco de aprendizaje automático de Apple, MLX, ofrece ventajas notables sobre otros marcos de aprendizaje profundo con su arquitectura de memoria unificada para el aprendizaje automático en el silicio de Apple. A diferencia de los marcos tradicionales como PyTorch y Jax, que requieren costosas copias de datos entre CPU y GPU, MLX mantiene los datos en la memoria compartida accesible para ambos. Este diseño elimina la sobrecarga de las transferencias de datos, lo que facilita una ejecución más rápida, particularmente con los grandes conjuntos de datos comunes en el aprendizaje automático. Para tareas complejas de aprendizaje automático en dispositivos Apple, la arquitectura de memoria compartida de MLX podría generar aceleraciones significativas. Esta característica hace que MLX sea muy relevante para los desarrolladores que buscan ejecutar modelos en dispositivos, como iPhones.
Con la experiencia de Apple en diseño de silicio, MLX insinúa las interesantes capacidades que podrían integrarse en sus chips para futuras aplicaciones de IA en dispositivos. El potencial de MLX para acelerar y optimizar las tareas de aprendizaje automático en las plataformas Apple lo convierte en un marco que los desarrolladores deben tener en cuenta.
Antes de implementar el modelo, se requiere cierta configuración. En primer lugar, es imprescindible instalar determinadas bibliotecas. Recuerde crear un entorno virtual antes de proceder con las instalaciones:
pip install mlx-lm