要約
Apple Silicon (M5 Max 128GB) で mlx-lm を使い、MoE モデル (Gemma 4 26B-A4B、128 expert) を LoRA 学習したところ、全 30 層を対象にすると iter 335〜353 で必ず RuntimeError: [metal::malloc] Resource limit (499000) exceeded で停止しました。ピークメモリは 87GB で set_memory_limit(95GB) の内側。これはメモリ量ではなく Metal の buffer 個数の上限で、MLX 側の未修正の問題 (mlx-lm#1185) です。8 層に絞ると同じ設定で 3,410 iters 完走します。
環境
- mlx 0.32.2 / mlx-lm 0.31.3、macOS、M5 Max 128GB
-
mlx-community/gemma-4-26b-a4b-it-4bit(30 層、MoE 128 expert) - LoRA rank 32、batch 1、
--grad-accumulation-steps 4、--grad-checkpoint、--mask-prompt、--max-seq-length 12288 - mlx-lm の既定では attention・dense MLP・expert (
LoRASwitchLinear)・router のすべての線形層に LoRA が入る
何が 499000 なのか
MLX の Metal allocator は live な MTLBuffer の個数を数えていて、iogpu.rsrc_limit という sysctl が無い Mac では 499,000 を上限に使います (mlx#1718 で導入)。mx.set_memory_limit / set_wired_limit / set_cache_limit はバイト数の制御で、この個数には効きません。
観測
| 構成 | 結果 |
|---|---|
| 30 層 / accum 4 | iter 335 で停止 (peak 87GB) |
| 30 層 / accum 4 (再走) | iter 353 で停止 |
| 30 層 / accum 4、最長系列 3 例 × 8 iters の探り | 通過 (peak 82GB) |
| 8 層 / accum 4 | 3,410 iters 完走 (peak 66GB) |
短い探りが通って本走が ~350 で死ぬので、iter とともに個数が累積する型です。mlx-lm は毎 step mx.clear_cache() を呼ぶのでキャッシュではなく、活性の buffer が増え続けています。mlx-lm#1185 の報告 (Qwen3.5/3.6 MoE、Gemma3) と同じ症状で、mx.compile された学習 step と MoE の gather_qmm (routing で毎回形が変わる) の相互作用が疑われています。2026-09 時点で未修正です。
もう 1 つの壁: 系列長
同じ環境で、系列長 ~14k トークンまでは学習でき、15k を超えると set_memory_limit を上げても Metal OOM になります。attention の backward が系列長の 2 乗で効くためで、cap の引き上げでは解けません。長いプロンプトは学習前に圧縮する (候補一覧を間引いて番号を振り直すなど) しかありませんでした。
回避策 (序列)
-
MLX_MAX_OPS_PER_BUFFERを小さくして command buffer を早く commit する (速度影響小。効果は経験報告) -
MLX_DISABLE_COMPILE=1(issue で有効との報告あり。約 2 倍遅い) -
lora_parameters.keysで対象を attention に限定する (個数は減るが、LoRA の文献では MLP / expert 側が学習の鍵とされているので最後の手) - 層数を減らす (8 層は完走。16〜20 は未検証)
検証結果
| 回避策 | 結果 |
|---|---|
MLX_MAX_OPS_PER_BUFFER=8 |
効果なし。初回とまったく同じ iter 335 で同じエラー (個数の増え方は決定的) |
MLX_DISABLE_COMPILE=1 |
完走 (30 層、682 iters、peak 83GB)。速度は compile 有りの約 15% 減で、心配した「2 倍遅」ではなかった |
つまり 2026-09 時点で、mlx-lm で MoE を全層 LoRA 学習するなら MLX_DISABLE_COMPILE=1 が事実上の必須設定です。この設定で学習したモデルは、8 層で学習したものより明確に良い出力 (JSON 末尾のリスト欄が教師の 85〜96% まで伸び、繰り返しも消えた) になったので、回避策を入れてでも全層に入れる価値がありました。
教訓
- 「メモリは余っているのに落ちる」ときは、エラー文の単位を読む。個数の壁はバイトの調整では解けない。
- 累積型の壁は短い探りでは検出できない。探りで安心せず、本走の途中 checkpoint を必ず残す。
- MoE + 全層 LoRA + mlx-lm は 2026-09 時点でまだ地雷がある。8 層で回る構成を先に確保してから広げる。