要約

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 層で回る構成を先に確保してから広げる。