本文来自DataLearnerAI官方网站:重磅!苹果官方发布大模型框架:一个可以充分利用苹果统一内存的新的大模型框架MLX,你的MacBook可以一键运行LLaMA了 | 数据学习者官方网站(Datalearner)https://www.datalearner.com/blog/1051701871117729
苹果刚刚发布了一个全新的机器学习矿机MLX,这是一个类似NumPy数组的框架,目的是可以在苹果的芯片上更加高效地运行各种机器学习模型,当然最主要的目的是大模型。
- MLX框架简介
- 目前MLX支持的模型和使用方式
- 预训练模型的使用
- MLX性能实测
- MLX的相关资源和链接
MLX框架简介
根据官方的介绍,MLX的设计受到PyTorch、Jax和ArrayFile的启发,目的是设计一个对用户极其友好,但同时在训练和部署上也非常高效的框架。所以,它的接口你会非常熟悉,因为它的Python接口与NumPy很相似,而它的神经网络模型的接口和PyTorch非常类似。所以如果此前你使用Python编写相关模型代码,几乎是没有障碍的使用和切换。
MLX的关键特性如下:
熟悉的API:MLX 有一个 Python API,紧密遵循 NumPy。MLX 还有一个功能齐全的 C++ API,与 Python API 非常相似。MLX 拥有像mlx.nn
和 mlx.optimizers
这样的高级包,其 API 紧密遵循 PyTorch,以简化构建更复杂的模型。
可组合的函数转换:MLX 具有可组合的函数转换,用于自动微分、自动向量化和计算图优化。
延迟计算:MLX 中的计算是延迟的。数组仅在需要时才实体化。
动态图构建:MLX 中的计算图是动态构建的。更改函数参数的形状不会触发缓慢的编译,调试简单直观。
多设备:操作可以在任何支持的设备上运行(目前是 CPU 和 GPU)。
统一内存:MLX 与其他框架的一个显著区别是统一内存模型。MLX 中的数组存在于共享内存中。对 MLX 数组的操作可以在任何支持的设备类型上执行,无需移动数据。
从上面的描述中可以看到,MLX与其它框架的最大的不同点是可以充分使用苹果的统一内存来计算,而不需要搬运数据。关于统一内存架构相比显存和内存的方式运行大模型参考:https://www.datalearner.com/blog/1051698716733526M3系列芯片是苹果最新发布的芯片。也是当前苹果性能最好的芯片。由于苹果的统一内存架构以及它的超大内存,此前很多人发现可以使用苹果的电脑来运行大语言模型。尽管它的运行速度不如英伟达最先进的显卡,但是由于超大的内存(显存),它可以载入非常大规模的模型。而此次的M3芯片效果如何,本文做一个简单的分析。https://www.datalearner.com/blog/1051698716733526
目前MLX支持的模型和使用方式
从官方的示例看,苹果的MLX框架完成度很高,它支持transformer架构模型的训练、当前预训练模型如何LLaMA的推理、基于LoRA的高效参数微调、支持载入使用多模态大模型Stable Diffusion和Whisper等。
此外,MLX还有个配套框架教MLX Data,这是一个不依赖特定框架的数据加载库,可以与PyTorch、Jax 或 MLX 一起使用。MLX Data数据加载库的目标是高效但同时也具有灵活性,例如能够每秒加载和处理数千张图片,同时也能在生成的批次上运行任意 Python 转换。这个库的推出很有可能与苹果的统一内存数据的使用有关,有了它应该可以更加高效地在苹果设备上读写数据。
关于MLX的实测和开源情况参考DataLearnerAI的原文:重磅!苹果官方发布大模型框架:一个可以充分利用苹果统一内存的新的大模型框架MLX,你的MacBook可以一键运行LLaMA了 | 数据学习者官方网站(Datalearner)