MLX — это фреймворк для исследований в области машинного обучения для Apple Silicon, разработанный компанией Apple.
Некоторые ключевые особенности MLX включают:
- Знакомые API: MLX имеет API на языке Python, близкий к NumPy. MLX также имеет полнофункциональные API для C++, C и Swift, которые в точности повторяют API Python. В MLX есть пакеты более высокого уровня, такие как mlx.nn и mlx.optimizers, с API, близкими к PyTorch, что упрощает построение более сложных моделей.
- Композитные преобразования функций: MLX поддерживает композитные преобразования функций для автоматического дифференцирования, автоматической векторизации и оптимизации вычислительных графов.
- Ленивые вычисления: Вычисления в MLX являются ленивыми. Массивы материализуются только тогда, когда это необходимо.
- Динамическое построение графов: Графы вычислений в MLX строятся динамически. Изменение формы аргументов функций не вызывает медленной компиляции, а отладка проста и интуитивно понятна.
- Мультиустройства: Операции могут выполняться на любом из поддерживаемых устройств (в настоящее время это CPU и GPU).
- Унифицированная память: Заметным отличием MLX от других фреймворков является унифицированная модель памяти. Массивы в MLX находятся в общей памяти. Операции над массивами MLX могут выполняться на любом из поддерживаемых типов устройств без передачи данных.
MLX разработан исследователями машинного обучения для исследователей машинного обучения. Фреймворк задуман как удобный для пользователя, но при этом эффективный для обучения и развертывания моделей. Дизайн самого фреймворка также концептуально прост.
В репозитории примеров MLX есть множество примеров, включая обучение языковой модели трансформера, генерация крупномасштабных текстов с помощью LLaMA и тонкая настройка с помощью LoRA, генерация изображений с помощью Stable Diffusion, распознавание речи с помощью Whisper от OpenAI.