TadaoYamaokaの日記

山岡忠夫Homeで公開しているプログラムの開発ネタを中心に書いていきます。

将棋AIの進捗 その50(メモリ使用量削減)

dlshogiは、1ノードにつき平均2KBのメモリを消費する(局面の平均合法手を80とした場合)。
通常GPU 1枚で探索した場合、NVIDIA RTX 3090で最大4.5万NPS程度なので、5分探索したとすると、探索ノード数は1350万ノードで、約27GBのメモリを消費する。

複数枚GPUを使用して、定跡を作成する場合には、さらに多くのノードを探索することになる。
5000万ノードを探索すると、100GBのメモリが必要になる。
なお、使用メモリ量が搭載物理メモリ以上になると、スワップが発生するため、NPSが極端に落ちる。

dlshogiのノードのメモリ管理

各ノードには、候補手ごとに、

  • 候補手を表すMove型
  • ポリシーネットワークの確率
  • 訪問回数
  • バックアップされた価値の累計
  • 子ノードへのポインタ

などを記録している。

問題箇所

Discordでやねうらお氏から、子ノードへのポインタを候補手すべてについて保持せず、実際に訪問して必要になった時点で確保すれば削減できるとアドバイスをもらった。
Leela Chess Zeroでも、Edge_Iteratorという仕組みでリンクリストのような構造で、必要になった時点で追加する方式を採用している。
しかし、リンクリストの管理を行うために、ノードの排他制御が必要になっているため、dlshogiでは採用したくなかった。
そのため、候補手の分、管理領域を予め確保して、配列にインデックスでアクセスするという方式をとっている。

メモリ使用量削減

速度を重視して、この方式は変えたくないのだが、終端ノードで一度も子ノードに訪問していない場合は、子ノードの管理領域を確保しておく必要がないことに気付いた。
初回にいずれかの子ノードの訪問した時点で、管理領域を確保するようにすれば、メモリを節約できる。
つまり、ゲーム木の終端ノードの数だけ、子ノードの管理領域の分のメモリが削減される。

そのように管理領域を確保するタイミング変更して、どれくらいメモリを節約できるか測定した。

5000万ノード探索時のメモリ使用量

変更前後で、5000万ノードの探索を行い、ps auxでRSS(物理メモリの使用サイズ)を3秒おきに確認した。
f:id:TadaoYamaoka:20201208234751p:plain

変更前と比べて、86%に減らせることが分かった。

探索速度

探索速度が低下していないか、floodgateの棋譜からサンプリングした100局面を使用したベンチマークプログラムで確認した。

V100×8で測定した結果
平均NPS
変更前 227648
変更後 230801

NPSは低下していないことが確認できた。

まとめ

dlshogiのノードのメモリ確保タイミングを変更するとで、メモリ使用量を削減した。
これでも、まだメモリ使用量が多いので、Leela Chess Zeroのようなメモリ管理方式にするかは悩みどころである。
定跡作成時でも、5000万ノードも探索できれば十分な気もするので、ひとまずこれ以上の変更は行わないことにする。

探索ノード数が増えると、floatの桁落ちで精度が落ちるという問題も見つかっているので、そちらも別途改善予定。

2020/12/11 追記

上記の修正を行い、変更前後で勝率を比較すると、弱くなることが判明した。
メモリ使用量より強さを優先したいため、変更は取り込まないことにする。

   # PLAYER              :  RATING  ERROR  POINTS  PLAYED   (%)  CFS(%)    W    D    L  D(%)
   1 dlshogi             :    10.3   20.4   144.0     272    53      84  134   20  118     7
   2 memory_efficient    :   -10.3   20.4   128.0     272    47     ---  118   20  134     7