2026-07-16 V14 256次元trainingのMPS OOM対策¶
現象¶
d_model=256, batch-size=64で既存約1.7M recordをbootstrap中、epoch 1 batch 14,200でMPS
OOMになった。直前のmemory logはmps_alloc=1086MBに対してmps_driver=43254MB、推奨上限は
25559MBだった。modelのlive tensorだけでなく、長時間実行でMPS driver
cacheと断片化が増えたことが主要因と判断した。
対応¶
- micro-batchを64から16へ縮小した。
- gradient accumulation 4でeffective batch 64とoptimizer更新回数を維持した。
- 50 micro-batchごとに
torch.mps.empty_cache()を呼ぶ。 - 5,000
micro-batchごとにmodel、optimizer、RNG、epoch内cursor、途中metricを
*.train-state.ptへ原子的に保存する。 pca train --resumeおよびpca selfplay-train --resumeから途中stateを復元する。- streaming再開時は
seed + epochで同じshuffleを再生成し、完了済みrecordをcollate前にskipする。
固定seedの小規模streaming学習をoptimizer step境界で意図的に中断し、resume後の全model tensorが中断なし実行とbitwise一致することをtestで確認した。
最初にbatch 14,200で失敗したrunには導入前のstateがなかったため、その地点自体からは再開できなかった。再実行ではbatch 5,000のstateが保存されており、後述する2回目のOOMからはそこを起点に再開できる。
再発と追加対策¶
micro-batch 16へ変更した再実行でもbatch 6,000でOOMが再発した。この時点でも
mps_alloc=758MBに対してmps_driver=43371MBであり、batch中のlive tensorが主因ではなかった。
PyTorchのdriver_allocated_memory()はallocator cacheだけでなくMPS/MPSGraph
frameworkの確保も含む。一方、
empty_cache()が解放するのはallocator内の未使用cacheに限られる。このため、batchごとにshapeが変わる現在のTransformer学習では、empty_cache()だけでdriver領域を長時間安定させられないと判断した。
- V14のMPS学習は3,000 micro-batchごとにstateを保存する。
- 保存後に
os.execv()で同じtrain commandを--resume付きで再実行し、MPS/MPSGraph process contextを解放する。 - CPU/CUDAではprocess再起動を行わない。
- stateにはJSONLのpath indexとbyte offset、shuffle RNG、10,000件bufferのfile referenceを保存する。2回目以降の再開では先頭からrecordを再decodeしない。
- 導入前に作られたbatch 5,000 stateも読み取れる。そのstateから最初に再開するときだけ、従来どおりprefixをreplayしてshuffle位置を復元する。
参考:
- PyTorch:
driver_allocated_memory - PyTorch:
empty_cache - PyTorch issue: MPS memory growth can persist despite
empty_cache
3,000 batch再起動前の再発¶
V14 cycle001のfinetuneをbatch 3,000から再開したところ、mps_driverは1.2 GBから増え続け、batch
5,800で43.3 GBに達して再びOOMになった。固定3,000 micro-batchの再起動予定より約200
batch早く失敗しており、固定間隔だけでは入力shapeの分布による増加速度へ追従できなかった。
train JSONLの先頭10,000 recordを調べると、state token長、盤面object数、history数、合法手数などを合わせたraw shapeは9,072種類あった。collatorが各batchの最大値へ動的paddingしていたため、MPSGraphが多数のshapeを保持するというログの挙動と整合する。
追加対応:
batch-padding-mode=mps_stableを追加し、V14で通常観測される範囲のpadding先を固定した。各値はminimumであり、範囲を超えるrecordは切り捨てない。mps-restart-driver-memory-mb=22000を追加した。optimizer更新境界でdriver memoryを確認し、固定3,000 batchより早く危険域へ達した場合はその場でstateを保存してprocessを置き換える。- padding profileとmemory thresholdはresume signatureの学習意味論から分離した。このため既存batch
3,000 stateを破棄せず、同じ
--resumecommandで対策を適用できる。