Performance is lacking

[vrnvorona]

When checking your repo I see benchmarks, but running Leduc CFR example on both CPU and GPU doesn't give anything even close to this performance. On 9950x I get like 30 it/s.

Please provide benchmarking code for cfrx side if example is not full picture (for example doesn't use vmap extensively idk).

MCCFR gives 190 it/s on GPU, 9.4k it/s on CPU

The fact that it works better on CPU is a huge sign of problem :(

GPU is 5090 btw.

[Egiob]

Please provide benchmarking code for cfrx side if example is not full picture (for example doesn't use vmap extensively idk).

I confirm the code used for benchmarking is not different than the one in the MCCFR trainer, only that I don't count exploitability measures (which is normal), in particular, no additional vmap were used.

The fact that it works better on CPU is a huge sign of problem :(

GPU is 5090 btw.

No it's typically not, the reason it works better on CPU is because pgx envs are slower on GPU when they are not heavily parallelized. For benchmarking I used the vanilla version of MCCFR that uses a single copy of the environment. Were you to write a version leveraging env parallelization, you would most likely see speedups.

When checking your repo I see benchmarks, but running Leduc CFR example on both CPU and GPU doesn't give anything even close to this performance. On 9950x I get like 30 it/s.

MCCFR gives 190 it/s on GPU, 9.4k it/s on CPU

That's a fair point, I had a hard time reproducing the result today since I didn't touch the code for a while. I figured out it's a JAX problem, see jax-ml/jax#24501. Rolling back to jax==0.4.31 or using the flag XLA_FLAGS=--xla_cpu_use_thunk_runtime=false gives the reported speed for me. Please tell me if that works for you.

[svtuck]

FWIW here is what I got back to 57k mc iterations /s with

pip list
Package                   Version     Editable project location
------------------------- ----------- --------------------------------
absl-py                   2.4.0
anyio                     4.12.1
appnope                   0.1.4
argon2-cffi               25.1.0
argon2-cffi-bindings      25.1.0
arrow                     1.4.0
asttokens                 3.0.1
async-lru                 2.2.0
attrs                     25.4.0
babel                     2.18.0
beautifulsoup4            4.14.3
bleach                    6.3.0
certifi                   2026.1.4
cffi                      2.0.0
cfrx                      0.0.2     
charset-normalizer        3.4.4
chex                      0.1.90
comm                      0.2.3
contourpy                 1.3.3
cycler                    0.12.1
debugpy                   1.8.20
decorator                 5.2.1
defusedxml                0.7.1
etils                     1.13.0
executing                 2.2.1
fastjsonschema            2.21.2
flax                      0.10.4
fonttools                 4.61.1
fqdn                      1.5.1
fsspec                    2026.2.0
h11                       0.16.0
httpcore                  1.0.9
httpx                     0.28.1
humanize                  4.15.0
idna                      3.11
importlib_resources       6.5.2
ipykernel                 7.2.0
ipython                   9.10.0
ipython_pygments_lexers   1.1.1
ipywidgets                8.1.8
isoduration               20.11.0
jax                       0.4.31
jaxlib                    0.4.31
jaxtyping                 0.3.9
jedi                      0.19.2
Jinja2                    3.1.6
json5                     0.13.0
jsonpointer               3.0.0
jsonschema                4.26.0
jsonschema-specifications 2025.9.1
jupyter                   1.1.1
jupyter_client            8.8.0
jupyter-console           6.6.3
jupyter_core              5.9.1
jupyter-events            0.12.0
jupyter-lsp               2.3.0
jupyter_server            2.17.0
jupyter_server_terminals  0.5.4
jupyterlab                4.5.4
jupyterlab_pygments       0.3.0
jupyterlab_server         2.28.0
jupyterlab_widgets        3.0.16
kiwisolver                1.4.9
lark                      1.3.1
markdown-it-py            4.0.0
MarkupSafe                3.0.3
matplotlib                3.10.8
matplotlib-inline         0.2.1
mdurl                     0.1.2
mistune                   3.2.0
ml_dtypes                 0.5.4
msgpack                   1.1.2
nbclient                  0.10.4
nbconvert                 7.17.0
nbformat                  5.10.4
nest-asyncio              1.6.0
notebook                  7.5.3
notebook_shim             0.2.4
numpy                     2.4.2
opt_einsum                3.4.0
optax                     0.2.5
orbax-checkpoint          0.6.4
packaging                 26.0
pandocfilters             1.5.1
parso                     0.8.6
pexpect                   4.9.0
pgx                       2.0.1
pgx-minatar               0.5.1
pillow                    12.1.1
pip                       26.0
platformdirs              4.9.2
prometheus_client         0.24.1
prompt_toolkit            3.0.52
protobuf                  6.33.5
psutil                    7.2.2
ptyprocess                0.7.0
pure_eval                 0.2.3
pycparser                 3.0
Pygments                  2.19.2
pyparsing                 3.3.2
python-dateutil           2.9.0.post0
python-json-logger        4.0.0
PyYAML                    6.0.3
pyzmq                     27.1.0
referencing               0.37.0
requests                  2.32.5
rfc3339-validator         0.1.4
rfc3986-validator         0.1.1
rfc3987-syntax            1.1.0
rich                      14.3.3
rpds-py                   0.30.0
scipy                     1.17.0
Send2Trash                2.1.0
setuptools                82.0.0
six                       1.17.0
soupsieve                 2.8.3
stack-data                0.6.3
svgwrite                  1.4.3
tensorstore               0.1.81
terminado                 0.18.1
tinycss2                  1.4.0
toolz                     1.1.0
tornado                   6.5.4
tqdm                      4.66.6
traitlets                 5.14.3
treescope                 0.1.10
typing_extensions         4.15.0
tzdata                    2025.3
uri-template              1.3.0
urllib3                   2.6.3
wadler_lindig             0.1.7
wcwidth                   0.6.0
webcolors                 25.10.0
webencodings              0.5.1
websocket-client          1.9.0
widgetsnbextension        4.0.15
zipp                      3.23.0

[vrnvorona]

Which is, no offense, pretty pointless as with modern hardware it's not possible to run 4.x.x JAX due to drivers and CUDA requirements for those GPUs.

[svtuck]

What speedup are you expecting to get from a gpu here? You either need to substantially rewrite the cfr code to reduce the number of dispatches to the gpu, or you just dont use cuda. Even PIOSolver doesn't benefit from gpu.