Zenn
Closed13

Apple Silicon向け機械学習用配列フレームワーク「MLX」を試す

kun432kun432

MLX全然ちゃんと試していないのでやる。

GitHubレポジトリ

https://github.com/ml-explore/mlx

MLX

MLXは、Appleの機械学習研究チームによって開発された、Appleシリコン向けの機械学習用配列フレームワークです。

MLXの主な特徴:

  • 親しみやすいAPI: MLXのPython APIはNumPyを踏襲しており、完全な機能を持つC++、CSwiftのAPIも提供され、Python APIを忠実に再現しています。また、mlx.nnmlx.optimizersなど、PyTorchに似た高レベルなパッケージも提供され、複雑なモデルの構築を簡略化します。
  • 合成可能な関数変換: MLXは、自動微分、自動ベクトル化、計算グラフの最適化を可能にする合成可能な関数変換をサポートしています。
  • 遅延計算: MLXの計算は遅延実行されます。配列は必要なときにのみ実体化されます。
  • 動的なグラフ構築: MLXの計算グラフは動的に構築されます。関数引数の形状が変わっても遅いコンパイルが発生せず、デバッグも簡単で直感的です。
  • マルチデバイス対応: 操作はサポートされるすべてのデバイス(現在はCPUとGPU)上で実行できます。
  • 統合メモリ: MLXが他のフレームワークと異なる顕著な点は統合メモリモデルです。MLXの配列は共有メモリ内に存在し、どのデバイスタイプでもデータ転送なしで操作が可能です。

MLXは機械学習研究者による、機械学習研究者のための設計です。このフレームワークはユーザーフレンドリーでありながら、モデルのトレーニングやデプロイにおいて効率的です。また、設計自体が概念的にシンプルであり、研究者がMLXを簡単に拡張し、新しいアイデアを迅速に探求できるようにしています。

MLXの設計は、NumPyPyTorchJaxArrayFireのようなフレームワークに触発されています。

MLXのexamplesリポジトリには、以下を含むさまざまな例があります:

ドキュメントはこちら

https://ml-explore.github.io/mlx/build/html/index.html

kun432kun432

インストール

https://ml-explore.github.io/mlx/build/html/install.html

環境はM2 Pro Mac Mini 32GB。

作業ディレクトリ作成して、Python仮想環境を作成する。自分はmiseを使う。

mkdir mlx-test && cd mlx-test
mise use python@3.12
cat << 'EOS' >> .mise.toml

[env]
_.python.venv = { path = ".venv", create = true }
EOS

パッケージインストール。ドキュメントを見ると以下とあるのだが・・・・

pip install mlx

のだが、このあとのQuick Start Guideを見ていくと、機械学習用「配列」フレームワークとしてのより低レベルAPIの使い方になっている。

https://ml-explore.github.io/mlx/build/html/usage/quick_start.html

自分はLLM実行のためにツールとしてイメージしていて、かつそのために使いたい。その場合には、GitHubのexampleレポジトリの方を見るのが良さそう。

https://github.com/ml-explore/mlx-examples

MLXを使ったテキスト生成モデルの推論についてはこちら。

https://github.com/ml-explore/mlx-examples/tree/main/llms/README.md

ということで、テキスト生成の推論はmlx-lmパッケージをインストールすることになる。

pip install mlx-lm
pip freeze | grep -i "mlx"
出力
mlx-lm==0.21.1
pip freeze | grep -i mlx
出力
mlx==0.22.0
mlx-lm==0.21.1
kun432kun432

テキスト生成

シンプルな生成

mlx_lm.generate --prompt "こんにちは!"

何かしらのモデルがダウンロードされる。

出力
tokenizer_config.json: 100%|███████████████████████████████████████████████████████| 54.5k/54.5k [00:00<00:00, 3.76MB/s]
config.json: 100%|█████████████████████████████████████████████████████████████████| 1.12k/1.12k [00:00<00:00, 6.03MB/s]
model.safetensors.index.json: 100%|████████████████████████████████████████████████| 45.7k/45.7k [00:00<00:00, 6.44MB/s]
special_tokens_map.json: 100%|█████████████████████████████████████████████████████████| 296/296 [00:00<00:00, 1.54MB/s]
tokenizer.json: 100%|██████████████████████████████████████████████████████████████| 9.09M/9.09M [00:01<00:00, 5.67MB/s]
tokenizer.json: 100%|██████████████████████████████████████████████████████████████| 9.09M/9.09M [00:01<00:00, 5.69MB/s]
special_tokens_map.json:   0%|                                                                | 0.00/296 [00:00<?, ?B/s]
model.safetensors:  17%|██████████                                                  | 304M/1.81G [00:36<02:51, 8.78MB/s]

ちょっと変ではあるけども、応答が返ってきた。

出力
こんにちは!どういたzなですか?
==========
Prompt: 37 tokens, 51.887 tokens-per-sec
Generation: 9 tokens, 40.159 tokens-per-sec
Peak memory: 1.856 GB

Usageを確認してみる。

mlx_lm.generate --help
出力
usage: mlx_lm.generate [-h] [--model MODEL] [--adapter-path ADAPTER_PATH]
                       [--extra-eos-token EXTRA_EOS_TOKEN [EXTRA_EOS_TOKEN ...]] [--system-prompt SYSTEM_PROMPT]
                       [--prompt PROMPT] [--max-tokens MAX_TOKENS] [--temp TEMP] [--top-p TOP_P] [--min-p MIN_P]
                       [--min-tokens-to-keep MIN_TOKENS_TO_KEEP] [--seed SEED] [--ignore-chat-template]
                       [--use-default-chat-template] [--verbose VERBOSE] [--max-kv-size MAX_KV_SIZE]
                       [--prompt-cache-file PROMPT_CACHE_FILE] [--kv-bits KV_BITS] [--kv-group-size KV_GROUP_SIZE]
                       [--quantized-kv-start QUANTIZED_KV_START] [--draft-model DRAFT_MODEL]
                       [--num-draft-tokens NUM_DRAFT_TOKENS]

LLM inference script

options:
  -h, --help            show this help message and exit
  --model MODEL         The path to the local model directory or Hugging Face repo. If no model is specified, then
                        mlx-community/Llama-3.2-3B-Instruct-4bit is used.
  --adapter-path ADAPTER_PATH
                        Optional path for the trained adapter weights and config.
  --extra-eos-token EXTRA_EOS_TOKEN [EXTRA_EOS_TOKEN ...]
                        Add tokens in the list of eos tokens that stop generation.
  --system-prompt SYSTEM_PROMPT
                        System prompt to be used for the chat template
  --prompt PROMPT, -p PROMPT
                        Message to be processed by the model ('-' reads from stdin)
  --max-tokens MAX_TOKENS, -m MAX_TOKENS
                        Maximum number of tokens to generate
  --temp TEMP           Sampling temperature
  --top-p TOP_P         Sampling top-p
  --min-p MIN_P         Sampling min-p
  --min-tokens-to-keep MIN_TOKENS_TO_KEEP
                        Minimum tokens to keep for min-p sampling.
  --seed SEED           PRNG seed
  --ignore-chat-template
                        Use the raw prompt without the tokenizer's chat template.
  --use-default-chat-template
                        Use the default chat template
  --verbose VERBOSE     Log verbose output when 'True' or 'T' or only print the response when 'False' or 'F'
  --max-kv-size MAX_KV_SIZE
                        Set the maximum key-value cache size
  --prompt-cache-file PROMPT_CACHE_FILE
                        A file containing saved KV caches to avoid recomputing them
  --kv-bits KV_BITS     Number of bits for KV cache quantization. Defaults to no quantization.
  --kv-group-size KV_GROUP_SIZE
                        Group size for KV cache quantization.
  --quantized-kv-start QUANTIZED_KV_START
                        When --kv-bits is set, start quantizing the KV cache from this step onwards.
  --draft-model DRAFT_MODEL
                        A model to be used for speculative decoding.
  --num-draft-tokens NUM_DRAFT_TOKENS
                        Number of tokens to draft when using speculative decoding.

なるほど、モデルの指定がなければ mlx-community/Llama-3.2-3B-Instruct-4bit が選択される様子。

利用可能なモデルはHuggingFaceのmlx-communityで確認できる。

https://huggingface.co/mlx-community

試しに以下のモデルを使用してみる。

https://huggingface.co/mlx-community/phi-4-4bit

mlx_lm.generate \
    --model "mlx-community/phi-4-4bit" \
    --max-tokens 1024 \
    --prompt "競馬の魅力を5つリストアップして。"
出力
競馬は多くの人々を魅了するスポーツであり、その魅力は多岐にわたります。以下に、競馬の魅力を5つ挙げます。

1. **スリルと興奮**: 競馬は、レースの展開が予測不可能で、一瞬の判断や馬の能力が勝敗を左右するため、非常にスリリングです。馬がゴールに向かって全力疾走する様子は、観戦者に強い興奮を与えます。

2. **戦略と分析**: 競馬は単なるスピード競争ではなく、騎手の戦略や馬の調教方法、コースの特性など、多くの要素が絡み合っています。これらを分析し、予想することは、ファンにとって大きな楽しみの一つです。

3. **歴史と伝統**: 競馬は長い歴史を持ち、多くの伝統的なレースが存在します。例えば、日本では「日本ダービー」や「菊花賞」などが有名で、これらのレースは競馬の歴史の一部として語り継がれています。

4. **美しさとスピード**: 競馬は、馬の美しさとスピードを同時に楽しむことができるスポーツです。馬の筋肉が動く様子や、風を切る姿は、観る者に美しさを感じさせます。

5. **コミュニティと交流**: 競馬場やオンラインでの予想会など、競馬を通じて多くの人々が交流を深める場があります。同じ趣味を持つ人々と情報を共有したり、予想を競い合ったりすることで、コミュニティの一員としての楽しみがあります。

これらの要素が組み合わさって、競馬は多くのファンにとって魅力的なスポーツとなっています。
==========
Prompt: 26 tokens, 49.936 tokens-per-sec
Generation: 638 tokens, 21.700 tokens-per-sec
Peak memory: 8.421 GB
kun432kun432

マルチターンのチャット

マルチターンのチャットもできる。mlx.chatを使う。

mlx_lm.chat --model "mlx-community/phi-4-4bit"
出力
[INFO] Starting chat session with mlx-community/phi-4-4bit. To exit, enter 'q'.
>> こんにちは。私の趣味は競馬です。覚えてね。
こんにちは!競馬が趣味とのこと、素晴らしいですね。競馬は戦略や予測が重要なスポーツで、多くのファンを魅了しています。馬の能力、騎手の技術、コースの特性など、さまざまな要素が結果に影響を与えます。競馬に関する質問や情報が必要な場合は、いつでもお気軽にお尋ねください。競馬の楽しみを存分に味わってくださいね!
>> 日本の総理大臣は誰?
2023年10月時点での日本の総理大臣は岸田文雄氏です。彼は自由民主党の党首として、2021年10月に第100代総理大臣に就任しました。政治状況は変動する可能性があるため、最新の情報を確認することをお勧めします。
>> そういえば私の趣味、覚えてる?
はい、覚えています!あなたの趣味は競馬です。競馬に関することや、その他の質問があれば、いつでもお気軽にお尋ねください。競馬の楽しみを引き続きお楽しみくださいね!
>> q

ちゃんと会話コンテキストが維持されているのがわかる。

kun432kun432

Python API

Pythonからも。

sample.py
from mlx_lm import load, generate

model, tokenizer = load("mlx-community/phi-4-4bit")

prompt = "競馬の魅力について5つリストアップして"

messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True
)

text = generate(
    model,
    tokenizer,
    prompt=prompt,
    max_tokens=1024,
    verbose=True,
)
python generate.py
出力
競馬は多くの人々にとって魅力的なスポーツであり、その理由は多岐にわたります。以下に、競馬の魅力を5つのポイントでまとめます。

1. **スリルと興奮**: 競馬は、レースの展開が予測不可能で、一瞬の判断や馬の能力が勝敗を左右するため、非常にスリリングです。ゴール直前の追い込みや、予想外の結果が出ることも多く、観戦者を飽きさせません。

2. **戦略と分析**: 馬券を購入する際には、馬の能力、騎手の技術、コースの特性、天候など様々な要素を考慮に入れる必要があります。この分析と戦略の要素が、競馬をより深く楽しむ要因となります。

3. **歴史と伝統**: 競馬は長い歴史を持ち、多くの伝統的なレースやイベントが存在します。例えば、日本では「日本ダービー」や「菊花賞」などが有名で、これらのレースは競馬ファンにとって特別な意味を持ちます。

4. **馬と人間の絆**: 騎手と馬の間には深い信頼関係が築かれており、その絆がレース中に見られることが多いです。また、馬の成長や活躍を追いかけることで、ファンは馬との感情的なつながりを感じることができます。

5. **コミュニティと交流**: 競馬場やオンラインでのファン同士の交流も魅力の一つです。同じレースを楽しみ、予想を共有することで、競馬を通じて新しい友人を作ることができます。また、競馬関連のイベントやファンクラブに参加することで、より深いコミュニティの一員になれます。
==========
Prompt: 28 tokens, 28.498 tokens-per-sec
Generation: 649 tokens, 21.553 tokens-per-sec
Peak memory: 8.421 GB

generateのパラメータは以下で確認できる。

python -c 'import mlx_lm; help(mlx_lm.generate)'
出力
Help on function generate in module mlx_lm.utils:

generate(model: mlx.nn.layers.base.Module, tokenizer: Union[transformers.tokenization_utils.PreTrainedTokenizer, mlx_lm.tokenizer_utils.TokenizerWrapper], prompt: Union[str, List[int]], verbose: bool = False, formatter: Optional[Callable] = None, **kwargs) -> str
    Generate a complete response from the model.

    Args:
       model (nn.Module): The language model.
       tokenizer (PreTrainedTokenizer): The tokenizer.
       prompt (Union[str, List[int]]): The input prompt string or integer tokens.
       verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
       kwargs: The remaining options get passed to :func:`stream_generate`.
          See :func:`stream_generate` for more details.

以下にもサンプルがある。

https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/examples/generate_response.py

kun432kun432

ストリーミング

上の例でも、出力はストリーミングで出力されていたと思うが、あくまでも出力だけで、実際には実際にはtextに生成されたテキストが入っているはず。Pythonコード上でストリーミングを行うにはstream_generateを使う。

streaming.py
from mlx_lm import load, stream_generate

model, tokenizer = load("mlx-community/phi-4-4bit")

prompt = "競馬の魅力について5つリストアップして"

messages = [{"role": "user", "content": prompt}]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True
)

response = stream_generate(
    model,
    tokenizer,
    prompt=prompt,
    max_tokens=1024,
)

for chunk in response:
    print(chunk.text, end="\n", flush=True)  # 実際には`end=""`を指定する
print()
出力

競

馬
は
多
く
の
人

々
に
と
って


魅
力
的
な
(snip)
kun432kun432

モデルの量子化・MLX形式への変換・HuggingFaceへのアップロード

HuggingFace上のモデルを量子化してMLX形式に変換、そしてHuggingFaceにアップロードすることもできる。

以下のモデルで試してみる。

https://huggingface.co/google/gemma-2-2b-jpn-it

convert.py
from mlx_lm import convert

repo = "google/gemma-2-2b-jpn-it"
upload_repo = "kun432/Google-gemma-2-2b-jpn-it-mlx-4bit"

convert(repo, quantize=True, upload_repo=upload_repo)
export HF_TOKEN=XXXXXXXXXX
python convert.py

まずHuggingFaceの元のモデルがダウンロードされる。

出力
[INFO] Loading
config.json: 100%|█████████████████████████████████████████████████████████████████████| 805/805 [00:00<00:00, 2.77MB/s]
special_tokens_map.json: 100%|█████████████████████████████████████████████████████████| 555/555 [00:00<00:00, 3.22MB/s]
model.safetensors.index.json: 100%|████████████████████████████████████████████████| 24.2k/24.2k [00:00<00:00, 23.5MB/s]
generation_config.json: 100%|██████████████████████████████████████████████████████████| 168/168 [00:00<00:00, 2.43MB/s]
tokenizer_config.json: 100%|███████████████████████████████████████████████████████| 46.9k/46.9k [00:00<00:00, 1.50MB/s]
tokenizer.model: 100%|█████████████████████████████████████████████████████████████| 4.24M/4.24M [00:01<00:00, 2.22MB/s]
model-00002-of-00002.safetensors:   0%|                                                      | 0.00/241M [00:00<?, ?B/s]
model-00001-of-00002.safetensors:   1%|▎                                           | 31.5M/4.99G [00:07<19:20, 4.27MB/s]
model-00002-of-00002.safetensors:   4%|█▉                                           | 10.5M/241M [00:04<01:37, 2.36MB/s]
tokenizer_config.json:   0%|                                                                | 0.00/46.9k [00:00<?, ?B/s]
(snip)

続けて、量子化・MLX形式への変換、そしてアップロードが行われる。

出力
Fetching 9 files: 100%|███████████████████████████████████████████████████████████████████| 9/9 [09:27<00:00, 63.03s/it]
[INFO] Quantizing00%|██████████████████████████████████████████████████████████████| 17.5M/17.5M [00:08<00:00, 2.22MB/s]
[INFO] Quantized model with 4.501 bits per weight.
README.md: 100%|███████████████████████████████████████████████████████████████████| 19.0k/19.0k [00:00<00:00, 26.5MB/s]
Repo created: https://huggingface.co/kun432/Google-gemma-2-2b-jpn-it-mlx-4bit
Found 8 candidate files to upload
Recovering from metadata files: 100%|███████████████████████████████████████████████████| 8/8 [00:00<00:00, 2357.18it/s]



---------- 2025-01-29 03:06:01 (0:00:00) ----------
Files:   hashed 5/8 (96.4K/1.5G) | pre-uploaded: 0/0 (0.0/1.5G) (+8 unsure) | committed: 0/8 (0.0/1.5G) | ignored: 0
---------- 2025-01-29 03:07:01 (0:01:00) ----------
Files:   hashed 8/8 (1.5G/1.5G) | pre-uploaded: 2/3 (38.6M/1.5G) | committed: 0/8 (0.0/1.5G) | ignored: 0
Workers: hashing: 0 | get upload mode: 0 | pre-uploading: 1 | committing: 0 | waiting: 9
---------------------------------------------------

model.safetensors:   1%|▋                                                          | 16.5M/1.47G [00:06<04:22, 5.55MB/s]
model.safetensors:  38%|██████████████████████▊                                     | 558M/1.47G [01:15<01:42, 8.87MB/s]

アップロードが完了して終了

出力
INFO:huggingface_hub._upload_large_folder:
---------- 2025-01-29 03:09:24 (0:03:23) ----------
Files:   hashed 8/8 (1.5G/1.5G) | pre-uploaded: 3/3 (1.5G/1.5G) | committed: 8/8 (1.5G/1.5G) | ignored: 0
Workers: hashing: 0 | get upload mode: 0 | pre-uploading: 0 | committing: 0 | waiting: 0
---------------------------------------------------
Upload successful, go to https://huggingface.co/kun432/Google-gemma-2-2b-jpn-it-mlx-4bit for details.

こんな感じでアップロードされた。

https://huggingface.co/kun432/Google-gemma-2-2b-jpn-it-mlx-4bit

あと、convertのパラメータは以下で確認できる。

python -c 'import mlx_lm; help(mlx_lm.convert)'
出力
Help on function convert in module mlx_lm.utils:

convert(hf_path: str, mlx_path: str = 'mlx_model', quantize: bool = False, q_group_size: int = 64, q_bits: int = 4, dtype: str = 'float16', upload_repo: str = None, revision: Optional[str] = None, dequantize: bool = False, quant_predicate: Optional[Callable[[str, mlx.nn.layers.base.Module, dict], Union[bool, dict]]] = None)

ちなみに量子化・変換途中の作業ディレクトリは実行したパス直下に作成される模様。

ls -lt
出力
total 32
drwxr-xr-x@ 11 kun432  staff  352  1 29 03:06 mlx_model

このディレクトリが残ったままだと、次にconvertするとこんなメッセージが出るので、削除が必要っぽい。

出力
ValueError: Cannot save to the path mlx_model as it already exists. Please delete the file/directory or specify a new path to save to.

なお、mlx-communityを見てると、みんなで変換してアップロードして共有、みたいな雰囲気を感じる。なるほど、なのでcommunityなのか。

kun432kun432

CLI

Quick Startでもちょっと触れたけども。

モデルを指定してプロンプトを渡してテキスト生成

mlx_lm.generate \
    --model "mlx-community/phi-4-4bit" \
    --max-tokens 1024 \
    --prompt "競馬の魅力を5つリストアップして。"

CLIでもモデルの変換ができる。mlx_lm.convertを使う。Usageは以下。

mlx_lm.convert --help
出力
usage: mlx_lm.convert [-h] [--hf-path HF_PATH] [--mlx-path MLX_PATH] [-q] [--q-group-size Q_GROUP_SIZE]
                      [--q-bits Q_BITS] [--dtype {float16,bfloat16,float32}] [--upload-repo UPLOAD_REPO] [-d]

Convert Hugging Face model to MLX format

options:
  -h, --help            show this help message and exit
  --hf-path HF_PATH     Path to the Hugging Face model.
  --mlx-path MLX_PATH   Path to save the MLX model.
  -q, --quantize        Generate a quantized model.
  --q-group-size Q_GROUP_SIZE
                        Group size for quantization.
  --q-bits Q_BITS       Bits per weight for quantization.
  --dtype {float16,bfloat16,float32}
                        Type to save the non-quantized parameters.
  --upload-repo UPLOAD_REPO
                        The Hugging Face repo to upload the model to.
  -d, --dequantize      Dequantize a quantized model.

1つ前に変換したやつだとこうなる。

mlx_lm.convert \
    --hf-path "google/gemma-2-2b-jpn-it" \
    -q \
    --upload-repo kun432/Google-gemma-2-2b-jpn-it-mlx-4bit-test
kun432kun432

長いプロンプトや生成

mlx-lmには、長いプロンプトや生成を効率的に処理するための方法として、以下の2つがある。

  • 循環型キー・バリューキャッシュ(Rotating Fixed-Size Key-Value Cache)
  • プロンプトキャッシュ

まず、循環型キー・バリューキャッシュ。最初イマイチ意味がわからなかったけどこういうことらしい。

  • 通常、LLMは生成時にトークンごとの状態(キー・バリュー情報)をメモリに保持するが、そのメモリ使用量が増えるとパフォーマンスが低下する
  • 一定サイズのメモリをキャッシュ領域として確保し、古い情報を新しい情報で上書きしながら使用することでメモリ使用量と品質のバランスを調整できる。
  • --max-kv-size nで、キャッシュサイズを整数で指定。
    • 例えば512のような小さい値だと、メモリ消費は抑えれるが、生成結果の品質は下がる
    • 例えば4096のような大きい値だと、生成結果の品質は上がるが、その分メモリも消費する

このnはトークン数なのかな?とりあえず試しにやってみた。

まず、意図的に小さい値を指定してみる。

mlx_lm.generate \
    --model "mlx-community/phi-4-4bit" \
    --max-tokens 1024 \
    --max-kv-size 100 \
    --prompt "競馬の魅力を5つリストアップして。"
出力
競馬は多くの人々を魅了するスポーツであり、その魅力は多岐にわたります。以下に、競馬の魅力を5つ挙げます。

1. **スリルと興奮**: 競馬は、レースの展開が予測不可能で、一瞬の判断や馬の走りによって結果が大きく変わるため、非常にスリリングです。観客は、馬の走りを見ながら息をのむ瞬間を体験できます。

レースの魅力は、馬の個性や騎手の技術、戦略が絡み合うことにあります。各馬にはそれぞれの特徴があり、それがレースの結果に大きく影響します。騎手は馬の能力を最大限に引き出すために、レース中にさまざまな戦略を駆使します。

1. **馬の特徴と戦略**:
   - 馬Aはスタートが速く、序盤でリードを取ることが多いですが、持久力に欠けます。
   - 馬Bは中盤で加速し、最後の直線で追い上げることが得意です。
   - 馬Cは一貫して安定したペースを保ち、最後まで脱落することが少ないです。

レースの結果を予測するために、以下の条件を考慮してください:
1. レースは標準的な直線コースで行われます。
2. 天候は晴れで、風の影響は最小限です。
3. 各馬のスタート位置はランダムに決定されます。
4. 各馬のスピードは一定ですが、スタート時の反応時間には個人差があります。

このシミュレーションを実行するために、以下の手順を踏んでください:

1. 各馬のスピードと反応時間をランダムに生成します。
2. 各馬のスタート位置をランダムに決定します。
3. 各馬がゴールに到達するまでの時間を計算します。
4. 各馬の順位を決定します。

以下に、各馬のスピードと反応時間をランダムに生成するためのコード例を示します。

```python
import random

# 馬の数
num_horses = 10

# 馬のスピードと反応時間をランダムに生成
horses = []
for i in range(num_horses):
    speed = random.uniform(10, 20)  # スピードは10から20の間でランダム
    reaction_time = random.uniform(0.1, 0.5)  # 反応時間は0.1から0.5秒の間でランダム
    horses.append((i, speed, reaction_time))

# 各馬のタイムを計算
times = []
for horse in horses:
    index, speed, reaction_time = horse
    time = distance / speed + reaction_time
    times.append((index, time))

# タイムが最も短い馬を見つける
times.sort(key=lambda x: x[1])
winner = times[0]

print(f"The winner is horse {winner[0] + 1} with a time of {winner[1]:.2f} seconds.")
```

### 説明:

1. **距離と速度**: 各馬の距離(100メートル)と速度(秒あたりのメートル)を定義します。

2. **時間の計算**: 各馬の時間を計算するために、距離を速度で
==========
Prompt: 26 tokens, 86.001 tokens-per-sec
Generation: 1024 tokens, 21.844 tokens-per-sec
Peak memory: 8.290 GB

上の例は極端な例だと思うのだけど、なんか全然まとまってないのがわかる。なるほど。

ちなみに大きめの値にしてみた。

mlx_lm.generate \
    --model "mlx-community/phi-4-4bit" \
    --max-tokens 4096 \
    --max-kv-size 4096 \
    --prompt "競馬の魅力を5つリストアップして。"
出力
競馬は多くの人々を魅了するスポーツであり、その魅力は多岐にわたります。以下に、競馬の魅力を5つ挙げます。

1. **スリルと興奮**: 競馬は、レースの展開が予測不可能で、一瞬の判断や馬の能力が勝敗を左右するため、非常にスリリングです。馬がゴールに向かって全力疾走する様子は、観戦者に強い興奮を与えます。

2. **戦略と分析**: 競馬は単なるスピード競争ではなく、騎手の戦略や馬の調教方法、コースの特性など、多くの要素が絡み合っています。これらを分析し、予想することは、ファンにとって大きな楽しみの一つです。

3. **歴史と伝統**: 競馬は長い歴史を持ち、多くの伝統的なレースが存在します。例えば、日本では「日本ダービー」や「菊花賞」などが有名で、これらのレースは競馬の歴史の一部として語り継がれています。

4. **美しさとスピード**: 競馬は、馬の美しさとスピードを同時に楽しむことができるスポーツです。馬の筋肉が動く様子や、風を切る姿は、観る者に美しさを感じさせます。

5. **コミュニティと交流**: 競馬場やオンラインでの予想会など、競馬を通じて多くの人々が交流を深める場があります。同じ趣味を持つ人々と情報を共有したり、予想を競い合ったりすることで、コミュニティの一員としての楽しみがあります。

これらの要素が組み合わさって、競馬は多くのファンにとって魅力的なスポーツとなっています。
==========
Prompt: 26 tokens, 93.955 tokens-per-sec
Generation: 638 tokens, 21.682 tokens-per-sec
Peak memory: 8.428 GB

オプションなしの場合と同じ結果になっている。

これデフォルト値はどうなってるんだろう?上にあるような単純な使い方だと効果が見えない気がするので、コード生成とか最近よくある深い推論のような長文の生成だと効いてくるのかもしれない。

次にプロンプトキャッシュ。mlx_lm.cache_promptを使うのだけど、CLIだと2段階の使い方になっている。

まず最初に適当に生成させてファイルに出力する。ここはプロンプトキャッシュとは直接関係ない。

mlx_lm.generate \
    --model "mlx-community/phi-4-4bit" \
    --max-tokens 4096 \
    --verbose False \
    --prompt "競馬の魅力を5つリストアップして。" > prompt.txt

こういう内容になる。

prompt.txt
 prompt.txt
競馬は多くの人々を魅了するスポーツであり、その魅力は多岐にわたります。以下に、競馬の魅力を5つ挙げます。

1. **スリルと興奮**: 競馬は、レースの展開が予測不可能で、一瞬の判断や馬の能力が勝敗を左右するため、非常にスリリングです。馬がゴールに向かって全力疾走する様子は、観戦者に強い興奮を与えます。

2. **戦略と分析**: 競馬は単なるスピード競争ではなく、騎手の戦略や馬の調教方法、コースの特性など、多くの要素が絡み合っています。これらを分析し、予想することは、ファンにとって大きな楽しみの一つです。

3. **歴史と伝統**: 競馬は長い歴史を持ち、多くの伝統的なレースが存在します。例えば、日本では「日本ダービー」や「菊花賞」などが有名で、これらのレースは競馬の歴史の一部として語り継がれています。

4. **美しさとスピード**: 競馬は、馬の美しさとスピードを同時に楽しむことができるスポーツです。馬の筋肉が動く様子や、風を切る姿は、観る者に美しさを感じさせます。

5. **コミュニティと交流**: 競馬場やオンラインでの予想会など、競馬を通じて多くの人々が交流を深める場があります。同じ趣味を持つ人々と情報を共有したり、予想を競い合ったりすることで、コミュニティの一員としての楽しみがあります。

これらの要素が組み合わさって、競馬は多くのファンにとって魅力的なスポーツとなっています。

このプロンプトファイルをmlx_lm.cache_promptに渡してキャッシュファイル化する。

cat prompt.txt | mlx_lm.cache_prompt \
    --model "mlx-community/phi-4-4bit" \
    --prompt - \
    --prompt-cache-file keiba_prompt.safetensors
出力
Processed    640 tokens (175.37 tok/s)
Peak memory: 9.045 GB
Saving...
ls -lt keiba_prompt.safetensors
出力
-rw-r--r--@ 1 kun432  staff  131080052  1 29 14:06 keiba_prompt.safetensors

ファイルが作成されている。

で、このキャッシュファイルを使って推論する。

mlx_lm.generate \
    --prompt-cache-file keiba_prompt.safetensors \
    --max-tokens 1024 \
    --prompt "\n大阪弁で書き直して。"
出力
競馬は、多くの人々を魅了するスポーツやね。その魅力は、いろいろあるんやで。以下に、競馬の魅力を5つ挙げてみようか。

1. **スリルと興奮**: 競馬は、レースの展開が読めんくて、一瞬の判断や馬の能力が勝敗を決めるから、すごくスリリングやねん。馬がゴールに向かって全力疾走する様子は、観てる人に強い興奮を与えてくれるんや。

2. **戦略と分析**: 競馬は単なるスピード競争やないんや。騎手の戦略や馬の調教方法、コースの特性なんかが絡み合ってるんやで。これらを分析して予想するのは、ファンにとって大きな楽しみの一つやね。

3. **歴史と伝統**: 競馬は長い歴史があって、たくさんの伝統的なレースがあるんや。例えば、日本では「日本ダービー」や「菊花賞」なんかが有名やねん。これらのレースは競馬の歴史の一部として語り継がれているんや。

4. **美しさとスピード**: 競馬は、馬の美しさとスピードを同時に楽しめるスポーツやね。馬の筋肉が動く様子や、風を切る姿は、観てる人に美しさを感じさせてくれるんや。

5. **コミュニティと交流**: 競馬場やオンラインでの予想会なんか、競馬を通じて多くの人々が交流を深める場があるんや。同じ趣味を持つ人々と情報を共有したり、予想を競い合ったりすることで、コミュニティの一員としての楽しみがあるんや。

このように、競馬は多くのファンにとって魅力的なスポーツやね。
==========
Prompt: 17 tokens, 43.353 tokens-per-sec
Generation: 659 tokens, 21.043 tokens-per-sec
Peak memory: 8.565 GB

キャッシュされたプロンプトは、CLIで与えられたプロンプトの前に配置される。あと、モデルも指定していないが、これもキャッシュファイル内に保持されているらしい。

なるほど、safetensors 形式で保存されているので、メモリ効率よく高速に読み込めるということなのだろうと思った。

プロンプトキャッシュはPythonからも使える。キャッシュの読み書き用にload_prompt_cache / make_prompt_cache / save_prompt_cache が用意されている。

https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/examples/chat.py

kun432kun432

サポートされているモデル

https://github.com/ml-explore/mlx-examples/blob/main/llms/README.md#supported-models

多くのモデルがサポートされていて、基本的に Mistral、Llama、Phi-2、Mixtral あたりがベースなのであればそのまま使えるとある。

ただし、モデルによっては以下のような対応が必要になる。

  • 元のモデルがtrust_remote_code=Trueが必要になる場合は、--trust-remote-codeオプションを渡して実行させることが可能。つけない場合には実行時に確認が求められる。
  • eos_tokenの指定が必要なものは--eos-token "<|endoftext|>"といった形で指定可能

これらはPythonからも指定できる。(上記ページの例を参考に)

kun432kun432

サイズの大きなモデル

https://github.com/ml-explore/mlx-examples/blob/main/llms/README.md#large-models

サイズの大きなモデルを動かすとメモリを占有してマシン自体が遅くなってしまうこともある。macOS15以降であれば、メモリの制限を設定することができる。ということで、これはMac側の設定。

以下の記事が詳しい。
https://zenn.dev/robustonian/articles/apple_silicon_vram

自分はまだmacOS14なのでこれは設定できない・・・(そろそろ上げてもいい気がしてきた)

kun432kun432

まとめ

とりあえず、Macでテキスト生成するのにミニマムな流れは押さえれたと思う。自分の目的は果たせたので、一旦ここまで。

examplesレポポジトリには、他にも

‐ 画像・音声・マルチモーダルなどのモデルの使い方

  • テキスト生成モデルのファインチューニング

などについて書かれているようなので参考に。

https://github.com/ml-explore/mlx-examples

このスクラップは2ヶ月前にクローズされました
ログインするとコメントできます