Gumbel dlshogiを作る その５（マルチプロセス対応）

前回、自己対局で10万局面生成するのに、1時間14分かかったので、マルチプロセス化して高速化することを検討した。

マルチプロセス化

シングルプロセスで実行してた自己対局の処理をそのままマルチプロセスで実行する。
GPUの推論は競合すると速度低下するため、排他制御する。

torch.multiprocessingとmodel.share_memory()を使用すると、モデルのパラメータを複数プロセスで共有することができるが、試したところTorchScriptのモデルではエラーになるため動かせなかった。

そこで、Python標準のmultiprocessingを使用して実装する。
モデルはそれぞれのプロセスでロードするため、プロセス数分GPUメモリを消費する。
GPUのメモリは、訓練時に大きなバッチサイズで学習する場合に大量に消費するが、推論のみであればGPUメモリは大量に消費しないため、プロセスごとにパラメータをロードしても問題ない。

ChatGPTによると、forkでマルチプロセスにすると、CUDAで問題が起きやすいため、spawnが推奨のようなので、明示的にspawnを使うようにする。

データの書き出し

訓練データの書き出し部分は、プロセスごとに出力すると合計件数の計算や生成速度の測定が難しくなるため、データの書き出しは専用プロセスにして、multiprocessing.Manager().Queue()で、各ワーカプロセスからデータを送信して、まとめて書き出すようにする。

実装

シングルプロセスの処理とは分けて、マルチプロセス用の自己対局処理を追加した。
探索部分の処理は共通なので、コードの理解のしやすさは損なっていない。

def selfplay_worker_mp(
    lock,
    model_path,
    batch_size,
    max_num_considered_actions,
    num_simulations,
    queue,
    amp,
):
    """マルチプロセス用の自己対局ワーカー"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.jit.load(model_path)
    model.to(device)
    model.eval()

    init_table(max_num_considered_actions, num_simulations)

    actors = [
        Actor(max_num_considered_actions, num_simulations, queue)
        for _ in range(batch_size)
    ]

    torch_features = torch.empty(
        (batch_size, FEATURES_NUM, 9, 9),
        dtype=torch.float32,
        pin_memory=True,
        requires_grad=False,
    )
    input_features = torch_features.numpy()

    while True:
        for i, actor in enumerate(actors):
            actor.next()
            make_input_features(actor.board, input_features[i])

        # Evaluate with lock
        with lock:
            with autocast(enabled=amp):
                with torch.no_grad():
                    logits, value = model(torch_features.to(device))

        for i, actor in enumerate(actors):
            actor.step.prior_logits = logits[i].cpu().numpy()
            actor.step.value = value[i].cpu().numpy()[0]

            invalid_actions = _get_invalid_actions(actor.board)
            actor.step.prior_logits = _mask_invalid_actions(
                actor.step.prior_logits, invalid_actions
            )


def selfplay_multiprocess(
    model_path,
    batch_size,
    max_num_considered_actions,
    num_simulations,
    output_dir,
    num_positions,
    amp,
    num_processes,
):
    """マルチプロセスで自己対局を実行する"""
    import multiprocessing as mp

    mp.set_start_method("spawn")

    queue = mp.Manager().Queue()
    lock = mp.Lock()

    writer_process = mp.Process(
        target=write_training_data, args=(queue, output_dir, num_positions)
    )
    writer_process.start()

    processes = []
    for _ in range(num_processes):
        p = mp.Process(
            target=selfplay_worker_mp,
            args=(
                lock,
                model_path,
                batch_size,
                max_num_considered_actions,
                num_simulations,
                queue,
                amp,
            ),
            daemon=True,
        )
        p.start()
        processes.append(p)

    try:
        writer_process.join()
        # Writerが終了したら、全局面生成完了
    except KeyboardInterrupt:
        print("\nTerminating self-play processes.")
    finally:
        for p in processes:
            if p.is_alive():
                p.terminate()
            p.join()
        if writer_process.is_alive():
            writer_process.terminate()
        writer_process.join()
        print("All processes terminated.")

動作確認

プロセス数を変えて、生成速度を確認した。
プロセス数を2,4,8,16にした場合の、生成速度(局面/秒)は以下の通り。

線形に近い形で、生成速度が上がっている。

GPUの使用率と、GPUメモリは、16プロセスでも余裕がある。

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 566.36                 Driver Version: 566.36         CUDA Version: 12.7     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4090      WDDM  |   00000000:01:00.0  On |                  Off |
| 30%   52C    P2            158W /  450W |   13017MiB /  24564MiB |     44%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

この実験では、10ブロック192フィルタの小さいモデルを使用しているので、大きなモデルではGPU使用率が上がると予想する。

まとめ

自己対局の高速化のため、Python標準のmultiprocessingを用いてマルチプロセス化を実装した。
プロセス数を増やすことで自己対局の生成速度はほぼ線形に向上し、GPUのリソースにも余裕があることが確認できた。
これで、Pythonのみでもかなりの速度で自己対局による訓練データ生成が可能になる。