Mac Studio M3 Ultraでmlx-lmを使ってllama4 Scoutを動かす

背景・目的

• llama4がリリースされた
• Llama4 Scount は1000万トークンのコンテキスト長を持つということなので、試してみた

コード

以下のように4-bit量子化版を使ってmlx-lmを使って推論させてみた。mlx-lm==0.22.4なら以下でOK。
0.22.3なら、こちらのPRからllama4.pyを使っているpythonのsite-packages/mlx_lm/models/llama4.pyに配置する

from mlx_lm import load, generate

model, tokenizer = load("mlx-community/meta-llama-Llama-4-Scout-17B-16E-4bit")


prompt = """
以下のコードで、stopトークンを設定する方法を教えて

<code>
from mlx_lm import load, generate

model, tokenizer = load("mlx-community/meta-llama-Llama-4-Scout-17B-16E-4bit")


prompt = ""

if tokenizer.chat_template is not None:
    messages = [
        {"role": "system", "content": "あなたは優秀なソフトウェアエンジニアです"},
        {"role": "user", "content": prompt},
    ]
    prompt = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True
    )

generate(model, tokenizer, prompt=prompt, verbose=True)
"""

if tokenizer.chat_template is not None:
    messages = [
        {"role": "system", "content": "あなたは優秀なソフトウェアエンジニアです"},
        {"role": "user", "content": prompt},
    ]
    prompt = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True
    )

generate(model, tokenizer, prompt=prompt, verbose=True)

結果

ピークメモリ60GBで、45トークン/秒の速度で生成。
生成されたコードは、存在しないパラメータを出力していたため、ハルシネーションは発生していた

==========
stopトークンを設定するには、generate関数のstop_tokensパラメータを使用します。stop_tokensは、生成を停止するトークンのリストです。

例えば、以下のようにstopトークンを設定できます。

<code>
from mlx_lm import load, generate

model, tokenizer = load("mlx-community/meta-llama-Llama-4-Scout-17B-16E-6bit")


prompt = ""

if tokenizer.chat_template is not None:
    messages = [
        {"role": "system", "content": "あなたは優秀なソフトウェアエンジニアです"},
        {"role": "user", "content": prompt},
    ]
    prompt = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True
    )

stop_tokens = ["<|endoftext|>"]

generate(model, tokenizer, prompt=prompt, verbose=True, stop_tokens=stop_tokens)
</code>

この例では、stop_tokensに["<|endoftext|>"]を設定しています。これにより、生成が"<|endoftext|>"トークンに達すると停止します。

stop_tokensパラメータは、リストまたは文字列を受け取ります。リストの場合は、リスト内のすべてのトークンが停止トークンとして扱われます。
==========
Prompt: 161 tokens, 296.893 tokens-per-sec
Generation: 256 tokens, 46.282 tokens-per-sec
Peak memory: 60.852 GB

備考

• 長いトークンの場合どうなるか？の検証として、こちらのアプローチでRoo Codeでも使ってみたものの、リピート出力が発生。パラメータ調整が必要かもしれない

  以下は参考に、Markdownの修正を頼んだ時のリクエストだが、フォルダに含まれるファイルのパターンを何度も生成し続けてしまっていた

  

結論

• 簡単に使ってみただけだが、もう少し良い量子化版やファインチューニング版などが出てきてから試した方が良いかもしれない
• また、パラメータ調整をしてみて、うまくいったら追記します