Low performance on TPUs

[wtedw]

PGX on TPUs seems to be slower than CPUs.
With a TPU v3-8, PGX is only achieving 1638 steps / sec on the game of chess.

Minimal Reproducible Example
PGX CPU vs TPU Test (512 env) (with sharding)

• https://gist.github.com/wtedw/e7332e8d99acd0132be5f82c389d8f60
• 512 envs * 512 game steps = 262,144 steps.
• Runtime = 2 min 40s = 160 sec
• 262,144 steps / 160 sec = 1,638 steps/sec

PGX CPU vs TPU Test (64 env) (single device)

• https://gist.github.com/wtedw/6e070e7ebed33dc52f798960a2789c75

Running around 8192 envs seems to be the limit. With split sharding across 8 devices, it takes about 1 hour and 27 minutes. If more than 8192 envs are used, there will be memory issues during JIT AOT compilation.

[sotetsuk]

Hi,

Thank you for opening the issue!

Given your comment, I did some preliminary experiments and I have confirmed similar results.
I noticed that the throughput on TPUs (v2) was indeed slow on colab, even though the same program scales appropriately on colab A100 GPU.

It may require further investigation. Let's use this issue.

TODOs may include:

• Try the latest TPU
• Test other games
• Find out bottleneck (I guess it's legal action mask computation)

[wtedw]

Hi @sotetsuk ,

I figured out what the bottleneck was! We just needed to do more matrix-multiplications / dot_general ops.
For the game of chess, we can improve the throughput from 1638 steps/s to 60k - 110k steps/s
For the game of hex (11x11, no swapping), we can even do 3 million - 42 million steps/s.

Unfortunately, I did the hex optimization work a while ago, so I'll just explain the chess changes.
Here's a draft of all the changes I've made for the chess env: wtedw@f9ff64c

The problem

The issue is that gather and scatter operations can be extremely slow on TPUs.
e.g:

• arr[0] (gather operation)
• arr.at[1].set(...) (scatter operation)

To get a perspective of how slow these operations are, here's a profile of a MCTS function call which calls chess_env.step() and does model inference (256/32 alphazero-like conv model).
~90% of the runtime is spent doing gather/scatter/scatter-add operations.

The solution

The solution is to transform these operations into matrix multiplication / dot_general operations which will make use of the underutilized MXU units. These tricks can be applied to any Pgx env to run faster on TPUs, and any JAX+TPU code for that matter (https://github.com/google-research/t5x/blob/63d9addf628c6d8c547a407a32095fcb527bb20b/t5x/examples/scalable_t5/layers.py#L280-L284).

Gather replacement

Suppose we want to retrieve a value from one of the 64 chess positions in state.board.
We can construct a one-hot mask with a 1 at idx and 0 elsewhere, then compute a dot product with the board.

For example, suppose we have a 2x2 board that is represented as a (4,) boolean array, and we want to do board[0]

mask = [1, 0, 0, 0] # Build one hot mask, shape is (4,)
result = mask @ board # (4,) @ (4,) -> ()
return result

In my commit, we do this with _pieces_at

def _pieces_at(board: Array, idx: Array) -> Array: 
    ..
    pieces = jax.nn.one_hot(idx, 64, dtype=board.dtype) @ board
    ...

Scatter replacement

To set a value into one of the 64 spots in state.board, we can use the same one hot mask trick along with an addition operation.

For example, suppose we had a 2x2 board that is represented as a (4,) boolean array, and we want to do board.at[3].set(val)

mask = [0, 0, 0, 1] # Build one hot mask
mask_inv = [1, 1, 1, 0] # Get the inverse
b = mask_inv * board # Select contents of the board
# b is now [?, ?, ?, 0] , we can set any T/F value if we add an array that looks like [0, 0, 0, val]
val_arr = mask * val # [0, 0, 0, val]
return b + val_arr

In my commit, we do this w/

def _set_pieces(board: Array, idx: Array, val: Array) -> Array:
    ....
    mask = _one_hot64(idx, dtype=board.dtype)           # (k, 64)
    return board * (1 - mask.max(axis=0)) + (mask.T @ val)

Other tricks

There are a lot of other optimizations that I haven't covered, but in general, we want to replace these operations w/ MXU friendly ops. Some notable changes not mentioned

• (scatter-add) Replace jnp.nonzero with lax.top_k
• (gather) CAN_MOVE, the one hot trick is too slow (I assume because the array is large 7x64x64), so we compute it on the fly
• (scatter) Replace legal_action_mask.at[...].set(True) w/ OR operations
• _is_attacked, completely written w/ a suggestion from Gemini. To be honest, I have no idea how it works.
• Instead of indexing into state.board_history, vmap a function over all the 8 board arrays

The result

We convert most of the scatter/gather ops into dot_general, and improve the same run_mcts runtime significantly.
Unfortunately, the before and after runs used different batch sizes, but the speed‑up is still substantial.

env.step speedup

When we rerun a similar Minimal Reproducible Example notebook on the same TPU v3-8, we can get a better estimation of these improvements.

Running chess_env.step for some batch size, for 128 iterations yields:

Batch	Wall_s	Steps_per_s
512	1.03635	63237.6
1024	1.26141	103909
2048	2.41507	108545
4096	4.7402	110605
8192	9.82247	106753
16384	19.3019	108650
32768	43.306	96852.6

https://gist.github.com/wtedw/22fb87cf4b30f494505e86c9a57e9792

Similarly for the game of hex (no swapping, 11x11)

Batch	Wall_s	Steps_per_s
512	0.0173938	3.76777e+06
1024	0.0180474	7.26266e+06
2048	0.0208172	1.25926e+07
4096	0.0259398	2.02117e+07
8192	0.0360053	2.91228e+07
16384	0.0562982	3.72508e+07
32768	0.0978984	4.28434e+07

https://gist.github.com/wtedw/e1061f95fca4843c17075956e6ae5305

Correctness

Many of these optimizations were co-developed w/ LLMs, so I'm still suspicious that the optimized chess_env.step function isn't semantically correct. However, the tests in test_chess.py all pass

# python -m pytest ~/repos/pgx/tests/test_chess.py

========================================= test session starts ==========================================
platform linux -- Python 3.12.7, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/tedpsw/repos/pgx
configfile: pyproject.toml
plugins: jaxtyping-0.2.36
collected 7 items                                                                                      

tests/test_chess.py .......                                                                      [100%]

==================================== 7 passed in 112.89s (0:01:52) =====================================

I've also done a selfplay test across 400k game states and checked if this chessopt_env.step produces the same thing as the original chess_env.step. The code is messy, unfortunately, but it goes something like

# chess0 is the original Pgx chess class
# env_reg is the original Pgx chess env

# chess is the scatter/gather optimized Pgx chess class
# env is the optimized Pgx chess env


while True:

        prev_env_state = ...
        env_state_opt = env.autostep(prev_env_state, action, env_rngs)

        # Convert the original state to the type expected by the regular env
        state_reg_input = chess0_lib.State(**vars(prev_env_state))
        # Step the regular environment with the converted state and same action
        env_state_reg = env_reg.autostep(state_reg_input, action, env_rngs)
        
        # Compare the two resulting states
        def _compare_fn(opt_state, reg_state, step_num):
            are_equal_tree = jax.tree_util.tree_map(np.array_equal, opt_state, reg_state)
            if not jax.tree.all(are_equal_tree):
                print(f"🚨 STATE MISMATCH (via callback) DETECTED! Step: {step_num}")
        jax.debug.callback(_compare_fn, env_state_opt, env_state_reg, step)

So far, I haven't encountered any mismatches.

Final thoughts

There's a lot to unpack, but I thought it'd be best to relay these optimization tricks to you, other contributors, and anyone who's working with TPUs. Pgx has been immensely helpful in my projects, and I really appreciate all the work that’s gone into it. If you'd like me to submit a PR, please let me know!

[sotetsuk]

@wtedw

Awesome! 🚀 Thank you for sharing your findings!! 🙏
It's really interesting and useful! I believe this report is really helpful for other users too!

If you'd like me to submit a PR, please let me know!

Unfortunately, I have no resource to review large PRs for now but I'm really interested in merging your efforts in the future!
As your changes are relatively large, before merging your efforts, we need

• Test enhancement: we need something like regression tests as you did in the notebook
• Continuous Performance measurement: we need CI for performance measurement for both GPU/TPU to ensure the performance improvement for both devices

So I'll let you know once we are ready to merge you efforts (Aug or Sep in 2025?). A that time, I may ask you to divide your changes into small chunks.
Again, thank you so much for your efforts!! 🙏