Metal

Metal量子化は、Apple Silicon (MPS) デバイス上で、Hugging Face Hub にホストされている Metal カーネル (kernels-community/mlx-quantization-metal-kernels) を使用してアフィン量子化を実行します。これらのカーネルは MLX フレームワークに由来し、kernels ライブラリを介してコンパイルされます。

重みはグループごとのスケールとバイアス付きのuint32テンソルにパッキングされ、フォワードパスでは融合されたデクアンタイゼーションと行列積のMetalカーネル（affine_qmm_t）が使用されます。これにより、CPUとの往復なしに推論を完全にGPU上で実行しながら、メモリ使用量を低く抑えることができます。

サポートされているビット幅は2、4、8です。グループサイズは設定可能で、デフォルト値は64です。

要件

MPS 対応の Apple Silicon Mac (M1 / M2 / M3 / M4)
kernelsパッケージ:

pip install kernels

Metal カーネルは初回使用時に Hub から自動的にダウンロードされるため、手動でのコンパイルは不要です。

オンザフライで量子化

MetalConfig を渡すことで、モデルの読み込み時に任意のモデルを量子化できます。対象となるすべての nn.Linear レイヤーは、量子化されたバージョンに置き換えられます。

from transformers import AutoModelForCausalLM, AutoTokenizer, MetalConfig

quantization_config = MetalConfig(bits=4, group_size=64)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    device_map="mps",
    quantization_config=quantization_config,
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
inputs = tokenizer("Apple Silicon is", return_tensors="pt").to("mps")
output = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(output[0], skip_special_tokens=True))

事前量子化済みモデルの読み込み

チェックポイントに量子化済みウェイト（weightがパックされたuint32、scales、qbiases）が既に含まれている場合、それらは直接読み込まれます。再量子化は不要です。

from transformers import AutoModelForCausalLM, MetalConfig

model = AutoModelForCausalLM.from_pretrained(
    "your-org/model-metal-4bit",
    device_map="mps",
)

Dequantize

MPS がないマシンでは、事前量子化されたチェックポイントは自動的に float にデクォンタイズされ、モデルが CPU または CUDA でも使用可能な状態に保たれます。この動作を明示的に強制することもできます。

from transformers import AutoModelForCausalLM, MetalConfig

config = MetalConfig(dequantize=True)
model = AutoModelForCausalLM.from_pretrained(
    "your-org/model-metal-4bit",
    quantization_config=config,
    device_map="cpu",
)

レイヤーの除外

一部のレイヤー（例: lm_head）は、modules_to_not_convert を使用して量子化の対象から除外できます。

config = MetalConfig(bits=4, group_size=64, modules_to_not_convert=["lm_head"])

設定オプション

パラメータ	デフォルト	説明
`bits`	`4`	重みの量子化に使用するビット幅(2、4、または 8)
`group_size`	`64`	量子化グループごとの要素数
`modules_to_not_convert`	`None`	フル精度で保持するモジュール名のリスト
`dequantize`	`False`	非MPSデバイス向けに、デクォンタイズをfloatに強制します

Update on GitHub