蘋果為自家芯片打造開源框架MLX,實現Llama 7B并在M2 Ultra上運行
2020 年 11 月,蘋果推出 M1 芯片,其速度之快、功能之強大一時令人驚艷。2022 年蘋果又推出了 M2,今年 10 月,M3 芯片正式登場。
蘋果在發布芯片的同時,也非常注重在其芯片上訓練和部署 AI 模型的能力。
蘋果推出的 ML Compute 可用于在 Mac 上進行 TensorFlow 模型的訓練。PyTorch 則支持在 M1 版本的 Mac 上進行 GPU 加速的 PyTorch 機器學習模型訓練,使用蘋果 Metal Performance Shaders (MPS) 作為后端來實現。這些使得 Mac 用戶能夠在本地訓練神經網絡。
現在,蘋果宣布推出專門在 Apple 芯片上用于機器學習的開源陣列框架 ——MLX。
MLX 是專門為機器學習研究人員設計的,旨在有效地訓練和部署 AI 模型。框架本身的設計在概念上也很簡單。研究人員能夠輕松地擴展和改進 MLX,以快速探索、測試新的想法。MLX 的設計靈感來自 NumPy、PyTorch、Jax 和 ArrayFire 等框架。
項目地址:https://github.com/ml-explore/mlx
MLX 項目貢獻者之一、Apple 機器學習研究團隊(MLR)研究科學家 Awni Hannun 展示了一段使用 MLX 框架實現 Llama 7B 并在 M2 Ultra 上運行的視頻。
MLX 迅速引起機器學習研究人員的關注。TVM、MXNET、XGBoost 作者,CMU 助理教授,OctoML CTO 陳天奇轉推表示:「蘋果芯片又有新的深度學習框架了。」
有網友評價 MLX 稱,蘋果再次「重造了輪子」。
圖源:https://twitter.com/ofervic/status/1732305883814596953
MLX 特性、示例
在該項目中,我們可以看到,MLX 有以下一些主要特性。
熟悉的 API。MLX 擁有非常像 NumPy 的 Python API,以及功能齊備的 C++ API(與 Python API 非常相似)。MLX 還有更高級的包(比如 mlx.nn 和 mlx.optimizers),它們的 API 很像 PyTorch,可以簡化構建更復雜的模型。
可組合函數變換。MLX 擁有自動微分、自動矢量化和計算圖優化的可組合函數變換。
惰性計算。MLX 中的計算是惰性的,陣列只有在需要時才被實例化。
動態圖構建。MLX 中的計算圖構建是動態的,改變函數參數的形狀不會導致編譯變慢,并且 debug 很簡單、容易上手。
多設備。任何支持的設備上(如 CPU 和 GPU)都可以運行操作。
統一內存。MLX 與其他框架的顯著差異在于統一內存,陣列共享內存。MLX 上的操作可以在任何支持的設備類型上運行,無需移動數據。
此外,項目中提供了多種使用 MLX 框架的示例,比如 MNIST 示例可以很好地讓你學習如何使用 MLX。
圖源:https://github.com/ml-explore/mlx-examples/tree/main/mnist
MLX 還有其他更多有用的示例,包括如下:
- Transformer 語言模型訓練;
- LLaMA 大規模文本生成和 LoRA 微調;
- Stable Diffusion 生成圖片;
- OpenAI 的 Whisper 語音識別。
更詳細的文檔可參閱:https://ml-explore.github.io/mlx/build/html/install.html#