1
$\begingroup$

Context:

I've implemented Muzero for the game Tic-tac-toe. Unfortunately, the self-play and training is very slow (like 10 hours until it plays quite well). I ran the python profiler to find the parts that take the most time. The result is that most time is spent doing Monte Carlo tree searches, specifically querying the neural networks for the next hidden state and the value, policy predictions.

While running self-play, my GPU is only on 30% load and my CPU is on max load (single process since GIL)

Note: In the muzero paper, they have done some fancy stuff like scaling some gradients, which I haven't implemented yet. This will probably also result in a small speedup

What I'm trying to do:

I want to speed up the self-play by running multiple MCTS's in different threads so that they pause whenever they want to query the neural network until enough other threads have queries for the network. Then I put all the queries into a batch and sent the batch to the network. Once I have the results, I return them to each thread, and they continue until they try to query the network again.

Reason:

Let's say I want to play 100 self-play games. I do 50 simulations per move, the average game length is 7 and the observation shape is (3,3,3). With my current approach, this would result in 100 * 7 * 50 = 35000 network queries of shape (1,3,3,3).

With the approach described above, I could run all 100 games at once and batch the network queries, resulting in 7 * 50 = 350 network queries of shape (100,3,3,3)

I hope that this will result in a significant speedup.

Questions:

  • What are your thoughts on this plan?
  • Any frameworks/PyTorch features that can help me with my plan?
  • How would you implement something like this?
  • How would you tackle problems like threads never waking up because the batch never gets full enough (if there's something better than timeouts), or possible race conditions?

Don't feel obligated to answer all of these questions.

$\endgroup$
3
  • $\begingroup$ you can just use minimax for tic-tac-toe... $\endgroup$
    – Alberto
    Feb 5 at 18:45
  • 1
    $\begingroup$ @Alberto Thats true but Tic-Tac-Toe is only a simple first step. Later on i will also let it learn other games with an action space to large for minimax. $\endgroup$
    – Lynix
    Feb 5 at 19:07
  • $\begingroup$ Yes, this could arguably be more focused, but I don't consider that sufficient justification for closing the question. $\endgroup$
    – mindcrime
    Mar 19 at 19:23

1 Answer 1

2
$\begingroup$

Batching: A Good Idea

You're right, batching is a great way to speed up AlphaZero or MuZero self-play! Your proposed solution of running multiple games in parallel is the easiest way to achieve some batching. There are other solutions, most notably Virtual loss, which allows you to get batches even from a single tree search. You can even combine both approaches.

Implementation Details

Unfortunately, I don't know of any PyTorch features that help with this, and in general Python is not the best language for concurrency and multithreading.

The way I implemented this in kZero (in Rust):

  • There is a single thread responsible for NN execution that listens to a channel/queue of work items. Once it has collected enough items, it evaluates the batch and sends the results back to each work item through a dedicated temporary response channel that was part of the original item.
  • Different games run on different threads or async tasks, send their evaluation requests to the main thread, and wait for the response.

The way to implement this in Python is with either threading and queues or with asyncio and its built-in queues.

See some of the diagrams in the kZero readme for some graphics to illustrate this.

An alternative is to share a single thread between multiple games and write a state machine to switch between them yourself. Effectively, this means you're implementing your own async-like system, this is tricky to get right! I did this at first but later switched to using async proper to simplify the code a lot.

Yet another alternative is to do what mctx does and let JAX "magically" handle all batching and concurrency. I'm not sure how well this actually works in practice, I haven't used JAX myself yet.

Deadlocks and Race Conditions

You can avoid the potential deadlock of batches never filling by just running enough games at once! For example, if you're only running a single NN execution thread with a batch size of 100, spawning at least 100 games will make deadlocks impossible. It's best to spawn at least twice as many games to get some concurrency, that way, by the time the first 100 requests have been evaluated, the next 100 are ready to go, and the NN thread never has to wait long.

The way to avoid deadlocks and race conditions caused by concurrency bugs is to write correct code! This is tricky to get right and basically an entire domain by itself. I've found that only communicating between threads using (concurrency-supporting) queues makes it relatively easy to write concurrency-bug-free code.

$\endgroup$
2
  • $\begingroup$ This is has helped me a lot. I really like the fact that you use something fast like rust for the mcts and just use python for the nn queries and training. For now im just using the library ray to run games in parallel. Its still pretty slow. For connect4 it takes 7min to run 100 games with 50 simulations per move. The training itself is pretty fast. I've implemented batchwise unrolling to generate predictions. I'm able to do 30 training steps per second at batch size 32. $\endgroup$
    – Lynix
    Mar 5 at 21:45
  • $\begingroup$ I'm happy to hear that! Yeah using a systems language for tree search and python for training is pretty common, LC0 and KataGo are doing the same thing. Indeed for AlphaZero selfplay is the part that's important to optimize, training is only a small fraction of the total cost. $\endgroup$ Mar 5 at 23:52

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .