Appleが自社チップ用のオープンソースフレームワークMLXを開発、Llama 7Bを実装しM2 Ultraで動作

2020年11月、Appleは速度と強力な機能の点で驚異的なM1チップを発売しました。 2022年にAppleはM2を発売し、今年10月にはM3チップが正式にデビューしました。

Apple はチップをリリースする際に、チップ上で AI モデルをトレーニングおよび展開する機能にも大きな重点を置いています。

Apple の ML Compute を使用して、Mac 上で TensorFlow モデルをトレーニングできます。 PyTorch は、Apple Metal Performance Shaders (MPS) をバックエンドとして使用し、Mac の M1 バージョンで GPU アクセラレーションされた PyTorch 機械学習モデル トレーニングをサポートします。これにより、Mac ユーザーはニューラル ネットワークをローカルでトレーニングできるようになります。

現在、Apple は、Apple チップ上での機械学習専用のオープンソース配列フレームワークである MLX のリリースを発表しました。

MLX は、機械学習研究者が AI モデルを効率的にトレーニングおよび展開できるように特別に設計されています。フレームワーク自体の設計も概念的にはシンプルです。研究者は MLX を簡単に拡張および改善して、新しいアイデアを迅速に探索およびテストできます。 MLX のデザインは、NumPy、PyTorch、Jax、ArrayFire などのフレームワークからインスピレーションを得ています。

プロジェクトアドレス: https://github.com/ml-explore/mlx

Apple の機械学習研究 (MLR) チームの研究者であり、MLX プロジェクトの貢献者の 1 人である Awni Hannun 氏は、MLX フレームワークを使用して Llama 7B を実装し、それを M2 Ultra で実行するビデオを公開しました。

MLX はすぐに機械学習研究者の注目を集めました。 TVM、MXNET、XGBoost の作者であり、カーネギーメロン大学の助教授、OctoML の CTO である Chen Tianqi 氏は、「Apple チップには新しいディープラーニング フレームワークが搭載されています」とリツイートしました。

一部のネットユーザーはMLXについてコメントし、Appleが再び「車輪の再発明」をしたと述べた。

画像出典: https://twitter.com/ofervic/status/1732305883814596953

MLX の機能、例

このプロジェクトでは、MLX には次の主な機能があることがわかります。

使い慣れた API 。 MLX には、NumPy に非常によく似た Python API と、フル機能の C++ API (Python API に非常によく似ています) があります。 MLX には、API が PyTorch と非常によく似ている、より複雑なモデルの構築を簡素化できる高レベル パッケージ (mlx.nn や mlx.optimizers など) もあります。

構成可能な関数変換。 MLX は、自動微分、自動ベクトル化、計算グラフ最適化を備えた構成可能な関数変換を特徴としています。

遅延評価。 MLX での計算は遅延され、配列は必要な場合にのみインスタンス化されます。

動的グラフ構築。 MLX の計算グラフ構築は動的であり、関数パラメータの形状を変更してもコンパイル速度が低下することはなく、デバッグもシンプルで使いやすいです。

マルチデバイス。操作は、CPU や GPU など、サポートされている任意のデバイスで実行できます。

統合メモリ。 MLX と他のフレームワークの大きな違いは、配列がメモリを共有する統合メモリです。 MLX での操作は、モバイル データを必要とせずに、サポートされている任意のデバイス タイプで実行できます。

さらに、このプロジェクトでは、MNIST の例など、MLX フレームワークを使用するさまざまな例が提供されており、MLX の使用方法を学ぶのに適しています。

画像ソース: https://github.com/ml-explore/mlx-examples/tree/main/mnist

MLX には、次のような便利な例が他にもたくさんあります。

• Transformer言語モデルのトレーニング。
• LLaMA 大規模テキスト生成と LoRA 微調整。
• 安定した拡散により画像が生成されます。
• OpenAI の Whisper 音声認識。

詳細なドキュメントについては、https://ml-explore.github.io/mlx/build/html/install.html# を参照してください。