Qhapaq Lab

バッチ推論

バッチ推論は、複数の局面をまとめて方策・価値ネットワークへ入力し、policy と value を一度に計算する実装技術である。

AlphaZero 型のモンテカルロ木探索(MCTS)では、葉ノードを展開するときにニューラルネットワークで局面を評価する。1 局面ずつ GPU に投げると、GPU の並列計算能力を十分に使えず、呼び出しや転送のオーバーヘッドも目立つ。そこで、複数の探索が到達した葉局面を集め、まとめて推論する。

何をまとめるか

典型的には、次のような流れになる。

1. MCTS を複数回進める
2. 未評価の葉局面をキューまたは配列へ積む
3. 入力特徴を batch 次元つきのテンソルへ詰める
4. GPU / 推論バックエンドでまとめて forward する
5. 各局面の policy/value を対応するノードへ戻す
6. 探索経路へ value をバックアップする

このとき、バッチサイズは「一度の推論に何局面をまとめるか」を決める重要なパラメータである。大きいほど GPU 効率は上がりやすいが、探索側が十分な未評価局面を集められない場合や、探索結果の反映が遅れる場合には、単純に大きければよいとは限らない。

dlshogiでの例

dlshogi では、USI エンジン側と自己対局生成側のどちらにもバッチ推論の仕組みがある。

USI 探索の usi/UctSearch.cpp では、policy_value_batch_maxsize 個ぶんの入力特徴、出力バッファ、policy_value_batch を確保する。探索中に QueuingNode が呼ばれると、make_input_features で局面特徴を作り、policy_value_batch[current_policy_value_batch_index] にノードと手番などを保存して、current_policy_value_batch_index を増やす。

その後 EvalNode で、現在たまっている局面数を policy_value_batch_size として nn_forward(policy_value_batch_size, ...) に渡す。返ってきた policy logits と value は、各ノードの合法手 prior と勝率値へ展開される。

自己対局生成の selfplay/self_play.cpp でも同様に、policy_value_batch_maxsize を使って複数の playout を集め、EvalNode でまとめて推論する。コマンドラインでは gpu_id batchsize [gpu_id batchsize]* の形で、GPU ごとにバッチサイズを指定できる。

USI エンジンでは DNN_Batch_Size, DNN_Batch_Size2, ... の USI オプションが SetThread へ渡され、GPU ごとの policy_value_batch_maxsize になる。

推論バックエンドとの関係

TensorRT 版の usi/nn_tensorrt.cpp では、NNTensorRT::forward が batch_size を受け取り、入力テンソルの batch 次元を設定してから CUDA へ転送し、context->enqueue(...) でまとめて推論する。出力も batch_size 個ぶんの policy と value としてホスト側へ戻される。

ONNX Runtime 版の usi_onnxruntime/nn_onnxruntime.cpp でも、入力 shape は batch_size を先頭次元に持つ。convert_model_to_onnx.py では、固定バッチでない場合に input1, input2, output_policy, output_value の 0 次元を batch_size として dynamic axes にしている。

つまり、探索部の「未評価局面を集める設計」と、推論バックエンドの「batch 次元を処理する設計」は対になっている。

Virtual lossとの関係

バッチ推論では、複数の playout を先に進めて未評価局面を集める。その間、複数スレッドや複数 playout が同じ有望経路へ集中すると、同じ局面ばかりキューに入ってしまう。そこでVirtual lossを使い、評価中の経路を一時的に選ばれにくくして探索の重複を減らす。

ただし Virtual loss は「バッチを大きくする魔法」ではない。強い prior や極端に有望な手がある局面では、衝突を完全には避けられない。Lc0 の開発資料でも、大きなバッチを集めること自体が難しい最適化課題として扱われている。

NNキャッシュとの関係

NN キャッシュは、過去に推論した局面の policy/value を保存する仕組みである。

  • キャッシュヒットした局面は、GPU バッチへ送らずにすぐ使える。
  • キャッシュミスした局面だけをバッチへ積める。
  • 終端局面やキャッシュヒットを先に処理する out-of-order evaluation 的な設計と相性がある。

バッチ推論は「まとめて計算する」技術であり、NN キャッシュは「同じ計算を避ける」技術である。どちらもニューラルネットワーク評価のコストを下げるが、効き方は異なる。

調整上の注意

バッチサイズを大きくすると、GPU のスループットは上がりやすい。一方で、次のような副作用がある。

  • 1 回の推論が返るまでの待ち時間が増える。
  • 探索結果のバックアップが遅れ、古い統計で次の playout を進めやすくなる。
  • 大きな入力・出力バッファが必要になる。
  • 小さい探索回数や短い持ち時間では、十分なバッチが集まらないことがある。
  • バッチサイズ、探索スレッド数、GPU 数、NN キャッシュサイズを一緒に調整する必要がある。

実戦エンジンでは、最善のバッチサイズはネットワークの大きさ、GPU、持ち時間、探索設定によって変わる。自己対局データ生成ではスループット重視、対局エンジンではレイテンシと探索品質のバランス重視、というように目的によっても適値が異なる。

関連項目

参考にしたホームページ