MCTSの超並列化について

将棋ソフトで使われている技術は将棋以外の分野で役に立つことは少ないのだが、Deep Learning系の将棋ソフトで使われている技術、例えばMCTS(モンテカルロ木探索)は、わりと広範な応用事例がある。

Practical Massively Parallel Monte-Carlo Tree Search Applied to Molecular Design : https://arxiv.org/abs/2006.10504

この論文のタイトルは、「分子設計に応用される実用的な超並列モンテカルロ木探索」(あってる?)で、MCTSを分子設計(希望する性質を持つ分子式の発見)に応用している。MCTSの大規模分散メモリ環境で(要するにPCが1000台ぐらいあって、メモリは共有していないという条件で)超並列を実現するというものである。

「え?この手法で(MCTSを用いている)将棋ソフトとか囲碁ソフトとかクラスター化できるの?」

と将棋ソフトとか囲碁ソフトの開発者の間でこの論文が話題になっていた。

私も少し興味があって、チラ見してみた。以下に感想というか、この論文手法を将棋ソフトとか囲碁ソフトへ適用できるのかについてざっと書いてみたい。

// チラ見なので間違って可能性はわりとあります。

MCTSとは?

MCTSではUCB1アルゴリズムに従い、最適なchild node(子ノード。nodeとは局面のこと)を選んでいく。MCTSの概略はWikipediaの説明でだいたいわかる。

// モンテカルロ木探索(Wikipedia) : https://ja.wikipedia.org/wiki/%E3%83%A2%E3%83%B3%E3%83%86%E3%82%AB%E3%83%AB%E3%83%AD%E6%9C%A8%E6%8E%A2%E7%B4%A2

MCTSでは訪問回数が一番多かったchild nodeが期待値最大となる。この収束証明は、以下にある。
// node = 局面、child node(子ノード = ある指し手で進めた局面)ぐらいの意味。root node = 探索開始局面。

// Finite-time Analysis of the Multiarmed Bandit
Problem : https://link.springer.com/content/pdf/10.1023/A:1013689704352.pdf

MCTSは、root nodeからよさげなchild nodeを辿っていく。最終的に、root nodeで訪問回数が最大のchild nodeを知りたいので、探索木を下手に分割することはできない。例えば、αβ探索のクラスター化でよくあるような、root nodeでの分割(異なるPCに割り振る)は、訪問回数がおかしくなってしまうので駄目なのである。

MCTSではleaf node(探索の末端局面)でplayout(rollout)を行う。これは一般的には対局シミュレーションを終局までやることを意味するのだが、その時間が勿体ないので将棋ソフトや囲碁ソフトではここで対局のシミュレーションを行わずにその局面をNN(ニューラルネットワーク)に評価させる。

この「評価」の部分(機械学習で言うところの推論)だけを異なるPCにお願いするという方法は考えられる。これもMCTSのクラスター化手法としてはアリだと思うが、1台のメインPCが探索をしていて、残りのPCがすべてNNの推論専用というのは負荷のバランスが悪く、台数が増えてくるとスケールアウトしない。

// 逆に言うと5~10台程度なら、こういう愚直なクラスター化手法もアリだという話もある。

分散MCTS(MCTSのクラスター並列化)

そこで、大規模でスケールアウトさせようと思うなら探索nodeごとにPCを分割する方法が考えられる。nodeによって決定論的に担当PCが決まる方法があれば良い。

冒頭の論文では、この「担当PC」のことを「home worker」と呼んでいるようだ。

将棋では、局面のhash keyとしてZobrist Hashという手法を用いている。要するに、局面に対応する64bit(やねうら王では128bit,256bitも選択可)の値が一意に決まる。

// 連載やねうら王miniで遊ぼう!7日目 : https://yaneuraou.yaneu.com/2015/12/16/%E9%80%A3%E8%BC%89%E3%82%84%E3%81%AD%E3%81%86%E3%82%89%E7%8E%8Bmini%E3%81%A7%E9%81%8A%E3%81%BC%E3%81%86%EF%BC%817%E6%97%A5%E7%9B%AE/
// 将棋ソフトに高速化は本当に必要なのか? : https://yaneuraou.yaneu.com/2015/12/17/%E5%B0%86%E6%A3%8B%E3%82%BD%E3%83%95%E3%83%88%E3%81%AB%E9%AB%98%E9%80%9F%E5%8C%96%E3%81%AF%E6%9C%AC%E5%BD%93%E3%81%AB%E5%BF%85%E8%A6%81%E3%81%AA%E3%81%AE%E3%81%8B%EF%BC%9F/

同じ局面であれば同じhash keyの値になるから、ある指し手で進めて、その局面のhash keyを求める。例えば、hash keyの下位2bitの値によって4台のPCのどのPCがhome workerに担当させるか決定する。その局面を担当するPCがどのPCであるかは他のPCに問い合わせることなくわかるので、そいつにお願いすれば良いのである。これでMCTSはクラスター並列化できているというわけである。(これが基本的なアイデアで、この手法自体は、冒頭の論文の共著者である美添 一樹さんが昔に考案されている。)

指し手で局面を進めるごとにhome worker(担当PC)が変わるので、そこで通信が都度発生する。この通信コスト(CPU時間)も通信による遅延もわりとあるが、冒頭の論文ではPCは1024台とかあるような超並列環境を想定しているので、通信によって、1台当たりの性能が1台単体の時の1/10の性能になろうが、全体がスケールアウトするなら、1台の時の (1024台×1/10=)102.4倍の速度で探索できているからそれでいいだろ、ということではある。

超分散MCTS

ところが、そうは言っても、MCTSにはback propagation(AlphaZeroの用語ではbackup)という工程がある。leaf nodeでの評価の結果(playoutの結果)をroot nodeへ伝達しなければならないのだ。この伝達において、上のアルゴリズムだとroot node付近を担当するhome workerで通信競合が起きる。1024台というような超並列を達成する時には、ここがボトルネックとなる。

そこで、root付近まで毎回back propagationするのは控えるというのが、冒頭の論文の提案である。root nodeからleaf nodeまでの経路上の、あるnodeでのUCB1最大の子ノードが変わらないなら、そこより上流nodeまでback propagationしなくていいだろ、というアイデアである。これにより、上のようなボトルネックが解消できるので超並列でスケールアウトするよ、ということのようだ。

ちなみにこれに似たアイデア自体は、やねうら王のdf-pnルーチンでも使っている。

https://github.com/yaneurao/YaneuraOu/blob/fb415337f90249148230516b98a9f6ab240b75f5/source/mate/mate_dfpn.hpp#L689

df-pnで、rootから現在探索中のnodeの経路上のあるnodeにおいてbestな子ノードが変わる時に、現在のnodeから上流nodeに戻るというアルゴリズムになっている。

より一般化して言うと、「rootから何らかのchild nodeの選択アルゴリズムによってleaf nodeまで到達して、leaf nodeで何かしら行う」というのを1ラウンド(1周)として、それを反復するようなアルゴリズムがあったとして、次のラウンドでは、root nodeからまたやりなおす必要はなく、leaf nodeからその経路上のnodeで、best childが変化するところまで上流(root node方向)に移動して、そのnodeからスタートすれば良い」ということである。

// MCTSでは、back propagationでrewardの情報がrootまでなかなか伝わらないことにはなるけども、その経路のnodeでのbest childが変わっていない時は、root nodeまで伝えなくとも探索結果に大きな影響はないとみなせる。

本論文手法の将棋ソフト、囲碁ソフトへの応用について

DL系の将棋ソフトでは局面の合流を扱っていない。

// 少なくともdlshogi、ふかうら王では扱っていない。他の将棋ソフトもおそらく同様。囲碁でもたぶんそう。

将棋では、そこまでの経路(手順)によって千日手が成立するかどうかが変わるため、局面の合流を真面目に扱うのは、とても難しいのだ。おそらく囲碁でも同様だと思う。(同形反復は禁止なので、手順によって次の手が合法手になるかどうかが変わるから。)

そこで、冒頭の論文の手法を適用するならば経路が違うなら、同一局面でも異なるnodeだとみなす必要がある。これは、経路までを含めた(Zobrist Hashによる)hash keyが必要である。

次に、1手進めるごとに次のhome workerに探索をバトンタッチすることになるわけだが、このオーバーヘッドが馬鹿にならない。将棋では1手進めるのは移動した駒の移動元の升を空きにして、移動先の升に移動元にあった駒を置くだけである。(成りとか駒取りはもうちょっと複雑だが) 要するに局面の差分更新はわずかなコストでできる。囲碁も打った石が1つ追加されるだけなので、盤面は広いものの、差分更新のコストはわずかである。このような大きな盤面(+そこまでの指し手すべて)をネットワーク上でやりとりするのは、わりと無駄である。なるべくなら、このやりとりのコストをなるべく削減したい。

// 少なくとも、まずhash keyを相手のhome workerにまず渡して、それがcacheされているか問い合わせてから、cacheされていないのであれば、局面とそこまでの手順(これがないと千日手の判定などができない)を送るというような通信プロトコルになっていたほうが良いと思う。

そこで、例えば、home workerの割当てをもう少し賢くする方法が考えられる。あるnodeのhash keyからhome workerを決めるのは、さきほどのルールで良いのだが、そうやって決まったhome workerは、(root nodeから数えて)5の倍数の手数の時しか変更しないだとか。まあ、何かしら通信量を減らすようなアイデアが必要だと思う。

それとは別の問題として、例えばdlshogiではA100×8(8GPU)のとき、DNN_Batch_Size = 128 , UCT_Threads = 4とかで探索している。つまりは、4スレッド×8GPU = 32スレッドが、それぞれ128局面leaf nodeが溜まったらNN.forward()を呼び出すというような実装になっている。root nodeに乗ってくるVirtual Lossは平均的には32×128×3(?) = 12,288ぐらいあって、これによってかなり歪(いびつ)な探索になってしまっているし、これのせいで並列化効率を大きく下げている。

このような状態ではメモリ分散環境のような厳しい条件下ではなく、メモリは共有していて(1つのCPU)、A100×64のような環境を想定した場合であっても、Virtual Lossのせいでそもそも並列化効率がよろしくないのであるが、この部分を先に解決しないといけないように思う。

長々と書いたが、結局のところ、冒頭の論文が将棋や囲碁で適用可能なのかどうかは、私にはよくわからない。

↓↓↓詳しい人、コメント欄で教えてもらえると助かります。↓↓↓

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です