GEMM results mismatch

[stepi9]

Description

Hi,
I am working on a deep learning project, where I need to use an AMD-GPU-based cluster to train the model. I am using jax, flax and optax as the main machine learning framework. When I run the python code on my personal computer, with an NVIDIA GPU, I do not receive any errors. However, when I try to run the same code on the AMD GPU on the cluster, I get the warning regarding the GEMM mismatch.

I have compiled, from source, the rocm-jax library with the latest tag (v0.6.0), and I have installed the package in a python venv. I have also tried with the official rocm/jax-community Docker image (via Apptainer), but I get the same messages.

The issue is not simply due to my code, as the same messages arise when I try to run one of Flax's example scripts (specifically https://github.com/google/flax/blob/main/examples/nnx_toy_examples/05_vae.py). When I try to run the code, I get:

E0611 14:26:23.597251  451832 buffer_comparator.cc:145] Difference at 9344: 205.843, expected 183.143
E0611 14:26:23.597299  451832 buffer_comparator.cc:145] Difference at 15892: 210.321, expected 187.659
E0611 14:26:23.597305  451832 buffer_comparator.cc:145] Difference at 15938: 199.099, expected 178.016
2025-06-11 14:26:23.597314: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.669495  451832 buffer_comparator.cc:145] Difference at 7792: 185.703, expected 208.131
E0611 14:26:23.669535  451832 buffer_comparator.cc:145] Difference at 12102: 206.548, expected 183.614
E0611 14:26:23.669549  451832 buffer_comparator.cc:145] Difference at 15948: 218.747, expected 194.395
2025-06-11 14:26:23.669557: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.682719  451832 buffer_comparator.cc:145] Difference at 1025: 10.3105, expected 8.15896
E0611 14:26:23.682732  451832 buffer_comparator.cc:145] Difference at 1029: 9.24706, expected 7.75471
E0611 14:26:23.682736  451832 buffer_comparator.cc:145] Difference at 1031: 10.8553, expected 9.51487
E0611 14:26:23.682739  451832 buffer_comparator.cc:145] Difference at 1032: 8.61683, expected 7.628
E0611 14:26:23.682743  451832 buffer_comparator.cc:145] Difference at 1035: 10.0826, expected 8.13794
E0611 14:26:23.682747  451832 buffer_comparator.cc:145] Difference at 1036: 9.36097, expected 8.08674
E0611 14:26:23.682750  451832 buffer_comparator.cc:145] Difference at 1038: 7.40636, expected 5.49972
E0611 14:26:23.682755  451832 buffer_comparator.cc:145] Difference at 1041: 9.11373, expected 8.06417
E0611 14:26:23.682758  451832 buffer_comparator.cc:145] Difference at 1044: 9.96564, expected 6.78432
E0611 14:26:23.682761  451832 buffer_comparator.cc:145] Difference at 1047: 10.3575, expected 8.61353
2025-06-11 14:26:23.682768: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.739910  451832 buffer_comparator.cc:145] Difference at 3137: 40.0469, expected 34.1624
E0611 14:26:23.739965  451832 buffer_comparator.cc:145] Difference at 3138: 34.1857, expected 29.3643
E0611 14:26:23.739970  451832 buffer_comparator.cc:145] Difference at 3139: 33.6209, expected 29.6085
E0611 14:26:23.739975  451832 buffer_comparator.cc:145] Difference at 3140: 33.4683, expected 29.7014
E0611 14:26:23.739982  451832 buffer_comparator.cc:145] Difference at 3143: 36.3527, expected 31.76
E0611 14:26:23.739986  451832 buffer_comparator.cc:145] Difference at 3144: 36.3201, expected 31.9039
E0611 14:26:23.740009  451832 buffer_comparator.cc:145] Difference at 3145: 37.0308, expected 32.787
E0611 14:26:23.740015  451832 buffer_comparator.cc:145] Difference at 3146: 33.6818, expected 29.4154
E0611 14:26:23.740021  451832 buffer_comparator.cc:145] Difference at 3150: 41.9207, expected 37.5318
E0611 14:26:23.740027  451832 buffer_comparator.cc:145] Difference at 3151: 33.3773, expected 28.8064
2025-06-11 14:26:23.740039: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.748676  451832 buffer_comparator.cc:145] Difference at 256: 33.6314, expected 29.6305
E0611 14:26:23.748690  451832 buffer_comparator.cc:145] Difference at 257: 35.9533, expected 31.5235
E0611 14:26:23.748695  451832 buffer_comparator.cc:145] Difference at 258: 34.0687, expected 29.0571
E0611 14:26:23.748702  451832 buffer_comparator.cc:145] Difference at 260: 34.1512, expected 30.1305
E0611 14:26:23.748709  451832 buffer_comparator.cc:145] Difference at 261: 35.5173, expected 31.2248
E0611 14:26:23.748714  451832 buffer_comparator.cc:145] Difference at 263: 37.6513, expected 33.6381
E0611 14:26:23.748720  451832 buffer_comparator.cc:145] Difference at 265: 34.4225, expected 30.5112
E0611 14:26:23.748724  451832 buffer_comparator.cc:145] Difference at 267: 40.5426, expected 36.3429
E0611 14:26:23.748729  451832 buffer_comparator.cc:145] Difference at 268: 35.4782, expected 30.2271
E0611 14:26:23.748736  451832 buffer_comparator.cc:145] Difference at 269: 34.3253, expected 30.6183
2025-06-11 14:26:23.748744: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.757706  451832 buffer_comparator.cc:145] Difference at 1025: 18.6348, expected 16.5961
E0611 14:26:23.757719  451832 buffer_comparator.cc:145] Difference at 1026: 15.6877, expected 13.5831
E0611 14:26:23.757724  451832 buffer_comparator.cc:145] Difference at 1038: 17.7971, expected 14.9409
E0611 14:26:23.757730  451832 buffer_comparator.cc:145] Difference at 1041: 18.4966, expected 16.4885
E0611 14:26:23.757736  451832 buffer_comparator.cc:145] Difference at 1043: 16.4748, expected 14.5351
E0611 14:26:23.757743  451832 buffer_comparator.cc:145] Difference at 1045: 15.9621, expected 14.029
E0611 14:26:23.757748  451832 buffer_comparator.cc:145] Difference at 1048: 16.9931, expected 14.4038
E0611 14:26:23.757755  451832 buffer_comparator.cc:145] Difference at 1050: 19.0769, expected 16.5237
E0611 14:26:23.757761  451832 buffer_comparator.cc:145] Difference at 1057: 14.1745, expected 12.4159
E0611 14:26:23.757766  451832 buffer_comparator.cc:145] Difference at 1064: 15.1762, expected 13.3455
2025-06-11 14:26:23.757774: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.765509  451832 buffer_comparator.cc:145] Difference at 384: 70.1718, expected 61.9522
E0611 14:26:23.765522  451832 buffer_comparator.cc:145] Difference at 394: 69.2255, expected 60.4526
E0611 14:26:23.765527  451832 buffer_comparator.cc:145] Difference at 397: 64.7409, expected 57.3438
E0611 14:26:23.765535  451832 buffer_comparator.cc:145] Difference at 404: 66.496, expected 59.6846
E0611 14:26:23.765540  451832 buffer_comparator.cc:145] Difference at 406: 67.2019, expected 59.6928
E0611 14:26:23.765547  451832 buffer_comparator.cc:145] Difference at 660: 71.3602, expected 63.6354
E0611 14:26:23.765554  451832 buffer_comparator.cc:145] Difference at 718: 64.1682, expected 71.8248
E0611 14:26:23.765563  451832 buffer_comparator.cc:145] Difference at 1691: 67.0263, expected 59.7899
E0611 14:26:23.765569  451832 buffer_comparator.cc:145] Difference at 1764: 59.8937, expected 67.8645
E0611 14:26:23.765576  451832 buffer_comparator.cc:145] Difference at 1765: 63.9905, expected 72.8709
2025-06-11 14:26:23.765583: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.769310  451832 buffer_comparator.cc:145] Difference at 384: 70.1718, expected 61.9522
E0611 14:26:23.769322  451832 buffer_comparator.cc:145] Difference at 394: 69.2255, expected 60.4526
E0611 14:26:23.769327  451832 buffer_comparator.cc:145] Difference at 397: 64.7409, expected 57.3438
E0611 14:26:23.769334  451832 buffer_comparator.cc:145] Difference at 404: 66.496, expected 59.6846
E0611 14:26:23.769339  451832 buffer_comparator.cc:145] Difference at 406: 67.2019, expected 59.6928
E0611 14:26:23.769347  451832 buffer_comparator.cc:145] Difference at 660: 71.3602, expected 63.6354
E0611 14:26:23.769353  451832 buffer_comparator.cc:145] Difference at 718: 64.1682, expected 71.8248
E0611 14:26:23.769363  451832 buffer_comparator.cc:145] Difference at 1691: 67.0263, expected 59.7899
E0611 14:26:23.769368  451832 buffer_comparator.cc:145] Difference at 1764: 59.8937, expected 67.8645
E0611 14:26:23.769373  451832 buffer_comparator.cc:145] Difference at 1765: 63.9906, expected 72.8709
2025-06-11 14:26:23.769380: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0611 14:26:23.805237  451832 buffer_comparator.cc:145] Difference at 1024: 9.26716, expected 7.62103
E0611 14:26:23.805249  451832 buffer_comparator.cc:145] Difference at 1025: 9.144, expected 8.11683
E0611 14:26:23.805254  451832 buffer_comparator.cc:145] Difference at 1029: 11.8019, expected 10.4442
E0611 14:26:23.805261  451832 buffer_comparator.cc:145] Difference at 1030: 9.69273, expected 7.69553
E0611 14:26:23.805268  451832 buffer_comparator.cc:145] Difference at 1031: 11.2271, expected 9.57831
E0611 14:26:23.805273  451832 buffer_comparator.cc:145] Difference at 1032: 10.0358, expected 8.47069
E0611 14:26:23.805279  451832 buffer_comparator.cc:145] Difference at 1034: 9.16138, expected 7.68747
E0611 14:26:23.805286  451832 buffer_comparator.cc:145] Difference at 1035: 10.7463, expected 9.28415
E0611 14:26:23.805291  451832 buffer_comparator.cc:145] Difference at 1037: 9.79935, expected 8.46282
E0611 14:26:23.805296  451832 buffer_comparator.cc:145] Difference at 1039: 9.90779, expected 8.7388
2025-06-11 14:26:23.805303: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.6.0.dev20250606+4952ad41a
jaxlib: 0.6.0.dev20250609
numpy:  2.3.0
python: 3.11.9 (main, Apr  3 2025, 09:59:09) [GCC 14.2.0]
device info: AMD Instinct MI210-3, 3 local devices"
process_count: 1
platform: uname_result(system='Linux', node='gpunode', release='5.14.0-503.33.1.el9_5.x86_64', version='#1 SMP PREEMPT_DYNAMIC Wed Mar 19 16:23:31 UTC 2025', machine='x86_64')

I compiled the library with ROCM 6.3.3

[ppanchad-amd]

Hi @stepi9. Internal ticket has been created to investigate this issue. Thanks!

[lucbruni-amd]

Hi @stepi9,

I was able to run the Flax example script you posted and did not encounter the GEMM results mismatch errors with the following setup (although with the 7900 XT) with ROCm 6.4.1.60401-83~24.04:

>>> jax.print_environment_info()
jax:    0.6.3.dev20250625+32c58ad7d
jaxlib: 0.6.3.dev20250702
numpy:  2.2.2
python: 3.12.3 (main, Jun 18 2025, 17:59:45) [GCC 13.3.0]
device info: Radeon RX 7900 XT-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='rocm', release='6.8.0-62-generic', version='#65-Ubuntu SMP PREEMPT_DYNAMIC Mon May 19 17:15:03 UTC 2025', machine='x86_64')

I also built from source but on the latest develop branch. Could you first try rebuilding on a later version to see if these errors are still being hit?

Also, please share your exact steps to reproduce the error starting from building Jax from source until you run the Flax example script on both baremetal and the rocm/jax-community Docker image. Thanks!

[lucbruni-amd]

Closing this issue due to inactivity and as I cannot reproduce the issue on a newer Jax source build.

Please try on the latest source build to confirm the fix. If there are still problems, feel free to reopen this ticket and we will continue investigation. Thanks!

[stepi9]

Hi @lucbruni-amd,
Unfortunately I have been a bit busy lately, sorry for the delayed response. To build from source, I am following the official AMD guide.

(P.S.: I can't reopen the issue)

I previously installed the following components:

• numactl 2.0.19
• clang 21.0.0git (https://github.com/llvm/llvm-project.git 18e51314c41ea3093f28659cd15095778dfe88f7)
• cmake 3.26.5

I cloned jax rocm in my home folder and I checked out the latest commit on the rocm-jaxlib-v0.6.0 branch.

Then, in the cloned-repository folder, I ran:

python3 ./build/build.py build --wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt --rocm_version=60 --rocm_path=/opt/rocm-6.3.3 --clang_path=$HOME/llvm-project/build/bin/clang --bazel_options="--linkopt=-L$NUMACTL_HOME/lib"

I created a new python environment with the following commands:

mkdir ~/python_venvs
cd ~/python_venvs
python3 -m venv ROCM_ENV

source ~/python_venvs/ROCM_ENV/bin/activate

pip install -U pip
pip install -U numpy scipy matplotlib seaborn pandas flax optax jaxopt datasets

cd ~/jax
python3 setup.py develop 
pip install -U dist/*.whl # to install the wheels
pip install -U . # to install jax

My environment:

jax:    0.6.0.dev20250801+9890a51d7
jaxlib: 0.6.0.dev20250805
numpy:  2.3.2
python: 3.11.9 (main, Apr  3 2025, 09:59:09) [GCC 14.2.0]
device info: AMD Instinct MI210-3, 3 local devices"
process_count: 1
platform: uname_result(system='Linux', node='gpunode', release='5.14.0-503.33.1.el9_5.x86_64', version='#1 SMP PREEMPT_DYNAMIC Wed Mar 19 16:23:31 UTC 2025', machine='x86_64')

With pip list:

jax                 0.6.0.dev20250801+9890a51d7
jax-rocm60-pjrt     0.6.0.dev20250805
jax-rocm60-plugin   0.6.0.dev20250805
jaxlib              0.6.0.dev20250805
jaxopt              0.8.5
flax                0.11.0

I then cloned the flax repository in my home folder and I launched the 05_vae.py example. This is the full output:

X_train: (60000, 28, 28) float32
X_test: (10000, 28, 28) float32
E0805 15:10:06.393002 2171903 buffer_comparator.cc:145] Difference at 1025: 10.3105, expected 8.15896
E0805 15:10:06.393031 2171903 buffer_comparator.cc:145] Difference at 1029: 9.24706, expected 7.75471
E0805 15:10:06.393036 2171903 buffer_comparator.cc:145] Difference at 1031: 10.8553, expected 9.51487
E0805 15:10:06.393040 2171903 buffer_comparator.cc:145] Difference at 1032: 8.61683, expected 7.628
E0805 15:10:06.393044 2171903 buffer_comparator.cc:145] Difference at 1035: 10.0826, expected 8.13794
E0805 15:10:06.393047 2171903 buffer_comparator.cc:145] Difference at 1036: 9.36097, expected 8.08674
E0805 15:10:06.393050 2171903 buffer_comparator.cc:145] Difference at 1038: 7.40636, expected 5.49972
E0805 15:10:06.393054 2171903 buffer_comparator.cc:145] Difference at 1041: 9.11373, expected 8.06417
E0805 15:10:06.393057 2171903 buffer_comparator.cc:145] Difference at 1044: 9.96564, expected 6.78432
E0805 15:10:06.393061 2171903 buffer_comparator.cc:145] Difference at 1047: 10.3575, expected 8.61353
2025-08-05 15:10:06.393072: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.449790 2171903 buffer_comparator.cc:145] Difference at 384: 70.1718, expected 61.9522
E0805 15:10:06.449825 2171903 buffer_comparator.cc:145] Difference at 394: 69.2255, expected 60.4526
E0805 15:10:06.449829 2171903 buffer_comparator.cc:145] Difference at 397: 64.7409, expected 57.3438
E0805 15:10:06.449832 2171903 buffer_comparator.cc:145] Difference at 404: 66.496, expected 59.6846
E0805 15:10:06.449837 2171903 buffer_comparator.cc:145] Difference at 406: 67.2019, expected 59.6928
E0805 15:10:06.449846 2171903 buffer_comparator.cc:145] Difference at 660: 71.3602, expected 63.6354
E0805 15:10:06.449850 2171903 buffer_comparator.cc:145] Difference at 718: 64.1682, expected 71.8248
E0805 15:10:06.449855 2171903 buffer_comparator.cc:145] Difference at 1691: 67.0263, expected 59.7899
E0805 15:10:06.449859 2171903 buffer_comparator.cc:145] Difference at 1764: 59.8937, expected 67.8645
E0805 15:10:06.449864 2171903 buffer_comparator.cc:145] Difference at 1765: 63.9905, expected 72.8709
2025-08-05 15:10:06.449873: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.451395 2171903 buffer_comparator.cc:145] Difference at 384: 70.1718, expected 61.9522
E0805 15:10:06.451407 2171903 buffer_comparator.cc:145] Difference at 394: 69.2255, expected 60.4526
E0805 15:10:06.451410 2171903 buffer_comparator.cc:145] Difference at 397: 64.7409, expected 57.3438
E0805 15:10:06.451413 2171903 buffer_comparator.cc:145] Difference at 404: 66.496, expected 59.6846
E0805 15:10:06.451416 2171903 buffer_comparator.cc:145] Difference at 406: 67.2019, expected 59.6928
E0805 15:10:06.451420 2171903 buffer_comparator.cc:145] Difference at 660: 71.3602, expected 63.6354
E0805 15:10:06.451425 2171903 buffer_comparator.cc:145] Difference at 718: 64.1682, expected 71.8248
E0805 15:10:06.451430 2171903 buffer_comparator.cc:145] Difference at 1691: 67.0263, expected 59.7899
E0805 15:10:06.451434 2171903 buffer_comparator.cc:145] Difference at 1764: 59.8937, expected 67.8645
E0805 15:10:06.451438 2171903 buffer_comparator.cc:145] Difference at 1765: 63.9906, expected 72.8709
2025-08-05 15:10:06.451443: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.499677 2171903 buffer_comparator.cc:145] Difference at 3137: 40.0469, expected 34.1624
E0805 15:10:06.499716 2171903 buffer_comparator.cc:145] Difference at 3138: 34.1857, expected 29.3643
E0805 15:10:06.499719 2171903 buffer_comparator.cc:145] Difference at 3139: 33.6209, expected 29.6085
E0805 15:10:06.499722 2171903 buffer_comparator.cc:145] Difference at 3140: 33.4683, expected 29.7014
E0805 15:10:06.499727 2171903 buffer_comparator.cc:145] Difference at 3143: 36.3527, expected 31.76
E0805 15:10:06.499730 2171903 buffer_comparator.cc:145] Difference at 3144: 36.3201, expected 31.9039
E0805 15:10:06.499735 2171903 buffer_comparator.cc:145] Difference at 3145: 37.0308, expected 32.787
E0805 15:10:06.499739 2171903 buffer_comparator.cc:145] Difference at 3146: 33.6818, expected 29.4154
E0805 15:10:06.499742 2171903 buffer_comparator.cc:145] Difference at 3150: 41.9207, expected 37.5318
E0805 15:10:06.499745 2171903 buffer_comparator.cc:145] Difference at 3151: 33.3773, expected 28.8064
2025-08-05 15:10:06.499754: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.511758 2171903 buffer_comparator.cc:145] Difference at 256: 33.6314, expected 29.6305
E0805 15:10:06.511772 2171903 buffer_comparator.cc:145] Difference at 257: 35.9533, expected 31.5235
E0805 15:10:06.511776 2171903 buffer_comparator.cc:145] Difference at 258: 34.0687, expected 29.0571
E0805 15:10:06.511779 2171903 buffer_comparator.cc:145] Difference at 260: 34.1512, expected 30.1305
E0805 15:10:06.511782 2171903 buffer_comparator.cc:145] Difference at 261: 35.5173, expected 31.2248
E0805 15:10:06.511785 2171903 buffer_comparator.cc:145] Difference at 263: 37.6513, expected 33.6381
E0805 15:10:06.511788 2171903 buffer_comparator.cc:145] Difference at 265: 34.4225, expected 30.5112
E0805 15:10:06.511792 2171903 buffer_comparator.cc:145] Difference at 267: 40.5426, expected 36.3429
E0805 15:10:06.511795 2171903 buffer_comparator.cc:145] Difference at 268: 35.4782, expected 30.2271
E0805 15:10:06.511799 2171903 buffer_comparator.cc:145] Difference at 269: 34.3253, expected 30.6183
2025-08-05 15:10:06.511804: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.517640 2171903 buffer_comparator.cc:145] Difference at 9344: 205.843, expected 183.143
E0805 15:10:06.517669 2171903 buffer_comparator.cc:145] Difference at 15892: 210.321, expected 187.659
E0805 15:10:06.517673 2171903 buffer_comparator.cc:145] Difference at 15938: 199.099, expected 178.016
2025-08-05 15:10:06.517680: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.525489 2171903 buffer_comparator.cc:145] Difference at 7792: 185.703, expected 208.131
E0805 15:10:06.525512 2171903 buffer_comparator.cc:145] Difference at 12102: 206.548, expected 183.614
E0805 15:10:06.525524 2171903 buffer_comparator.cc:145] Difference at 15948: 218.747, expected 194.395
2025-08-05 15:10:06.525530: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.531965 2171903 buffer_comparator.cc:145] Difference at 1024: 9.26716, expected 7.62103
E0805 15:10:06.531978 2171903 buffer_comparator.cc:145] Difference at 1025: 9.144, expected 8.11683
E0805 15:10:06.531982 2171903 buffer_comparator.cc:145] Difference at 1029: 11.8019, expected 10.4442
E0805 15:10:06.531985 2171903 buffer_comparator.cc:145] Difference at 1030: 9.69273, expected 7.69553
E0805 15:10:06.531988 2171903 buffer_comparator.cc:145] Difference at 1031: 11.2271, expected 9.57831
E0805 15:10:06.531991 2171903 buffer_comparator.cc:145] Difference at 1032: 10.0358, expected 8.47069
E0805 15:10:06.531993 2171903 buffer_comparator.cc:145] Difference at 1034: 9.16138, expected 7.68747
E0805 15:10:06.531998 2171903 buffer_comparator.cc:145] Difference at 1035: 10.7463, expected 9.28415
E0805 15:10:06.532001 2171903 buffer_comparator.cc:145] Difference at 1037: 9.79935, expected 8.46282
E0805 15:10:06.532005 2171903 buffer_comparator.cc:145] Difference at 1039: 9.90779, expected 8.7388
2025-08-05 15:10:06.532010: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
E0805 15:10:06.539354 2171903 buffer_comparator.cc:145] Difference at 1025: 18.6348, expected 16.5961
E0805 15:10:06.539367 2171903 buffer_comparator.cc:145] Difference at 1026: 15.6877, expected 13.5831
E0805 15:10:06.539371 2171903 buffer_comparator.cc:145] Difference at 1038: 17.7971, expected 14.9409
E0805 15:10:06.539374 2171903 buffer_comparator.cc:145] Difference at 1041: 18.4966, expected 16.4885
E0805 15:10:06.539379 2171903 buffer_comparator.cc:145] Difference at 1043: 16.4748, expected 14.5351
E0805 15:10:06.539382 2171903 buffer_comparator.cc:145] Difference at 1045: 15.9621, expected 14.029
E0805 15:10:06.539386 2171903 buffer_comparator.cc:145] Difference at 1048: 16.9931, expected 14.4038
E0805 15:10:06.539389 2171903 buffer_comparator.cc:145] Difference at 1050: 19.0769, expected 16.5237
E0805 15:10:06.539392 2171903 buffer_comparator.cc:145] Difference at 1057: 14.1745, expected 12.4159
E0805 15:10:06.539396 2171903 buffer_comparator.cc:145] Difference at 1064: 15.1762, expected 13.3455
2025-08-05 15:10:06.539403: E external/xla/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc:1165] Results do not match the reference. This is likely a bug/unexpected loss of precision.
Epoch 0 loss: 0.30997368693351746
Epoch 1 loss: 0.22965705394744873
Epoch 2 loss: 0.21164396405220032
Epoch 3 loss: 0.20139960944652557
Epoch 4 loss: 0.19378237426280975
Epoch 5 loss: 0.18765537440776825
Epoch 6 loss: 0.18277066946029663
Epoch 7 loss: 0.17905406653881073
Epoch 8 loss: 0.176287442445755
Epoch 9 loss: 0.17389754951000214
Epoch 10 loss: 0.17164503037929535
Epoch 11 loss: 0.16956108808517456
Epoch 12 loss: 0.16839267313480377
Epoch 13 loss: 0.1673182100057602
Epoch 14 loss: 0.16584056615829468
Epoch 15 loss: 0.16545917093753815
Epoch 16 loss: 0.16392460465431213
Epoch 17 loss: 0.16328173875808716
Epoch 18 loss: 0.1624574214220047
Epoch 19 loss: 0.16176360845565796

[lucbruni-amd]

Hi @stepi9,

According to the logs, this seems to be related to the autotuner.

Could you try disabling xla_gpu_autotune_level, which may be causing the messages?

Here is some context regarding this setting:

xla_gpu_autotune_level: sets matmul, triton GEMM and convolution auto-tuning level (default is 4).

Autotuner generally works as follows: it takes a list of candidate solutions (algorithms) from the underlying math library (rocBLAS, hipBlasLt, MIOpen, etc.) for a given problem, runs each of them several times, and picks up the fastest one. Supported autotune levels are:
0 - autotuning is disabled
1 - autotuning is enabled
2 - same as level 1 but with input buffers initialization (so that the input operands for the operation being tuned are initialized with random but valid values)
3 - same as level 2 but with input buffers reinitialization after benchmarking each candidate solution. This might be necessary if some input buffer is also used as an output buffer: e.g., consider a GEMM operation: D = alphaAB + beta*C, where alpha/beta are non-zero, and D and C share the same buffer.
4 - same as level 3 but also enables correctness check. The correctness check works as follows: the autotuner chooses the first solution to be the reference one, and checks the accuracy of all candidate solutions versus the reference solution. Note that, the solutions with failed accuracy check are not discarded on level 4 (XLA only prints a corresponding warning message). Besides, it also performs redzone-checking (out-of-bound access) for the output buffer. If this is the case, the faulty solution is discarded.
5 - same as level 4 but also discards solutions based on correctness check (instead of just printing the warning message).

An environment variable example of setting this flag to disable the autotuner would be like:

export XLA_FLAGS="--xla_gpu_autotune_level=0"

Then, you can run the script again, which imports JAX after the flag is set (it should pick up on the flag(s) during initialization).

Let me know if this helps you run your workload. Thanks!

[stepi9]

Hi @lucbruni-amd,

Thank you for your reply.

I have tried to apply the suggested fix to the example flax/examples/nnx_toy_examples/05_vae.py.

With autotuning disabled, I do not get the warning anymore. However, it completely changes the output of the programme. Moreover, it seems that the model is not learning correctly.

Here is the loss function with XLA_FLAGS unset:

Epoch 0 loss: 0.30997368693351746
Epoch 1 loss: 0.22965705394744873
Epoch 2 loss: 0.21164394915103912
Epoch 3 loss: 0.20139962434768677
Epoch 4 loss: 0.19378238916397095
Epoch 5 loss: 0.18765632808208466
Epoch 6 loss: 0.18275947868824005
Epoch 7 loss: 0.17905645072460175
Epoch 8 loss: 0.17630773782730103
Epoch 9 loss: 0.1739160120487213
Epoch 10 loss: 0.17164623737335205
Epoch 11 loss: 0.16959188878536224
Epoch 12 loss: 0.16839230060577393
Epoch 13 loss: 0.1672750860452652
Epoch 14 loss: 0.16583149135112762
Epoch 15 loss: 0.16540634632110596
Epoch 16 loss: 0.16393810510635376
Epoch 17 loss: 0.16326674818992615
Epoch 18 loss: 0.1624736338853836
Epoch 19 loss: 0.16171644628047943

Instead, here is the loss function with the autotuning disabled:

Epoch 0 loss: 0.34972521662712097
Epoch 1 loss: 0.3227783143520355
Epoch 2 loss: 0.32643717527389526
Epoch 3 loss: 0.3294863998889923
Epoch 4 loss: 0.3330540359020233
Epoch 5 loss: 0.33656197786331177
Epoch 6 loss: 0.34022581577301025
Epoch 7 loss: 0.3431639075279236
Epoch 8 loss: 0.3488299548625946
Epoch 9 loss: 0.3514985740184784
Epoch 10 loss: 0.35641488432884216
Epoch 11 loss: 0.35863131284713745
Epoch 12 loss: 0.3658508360385895
Epoch 13 loss: 0.3671998679637909
Epoch 14 loss: 0.37013572454452515
Epoch 15 loss: 0.37649545073509216
Epoch 16 loss: 0.3780072033405304
Epoch 17 loss: 0.3807724118232727
Epoch 18 loss: 0.3812345266342163
Epoch 19 loss: 0.3828718662261963

For reference, this is the output on (default) Jax CUDA:

Epoch 0 loss: 0.3101576566696167
Epoch 1 loss: 0.2296341359615326
Epoch 2 loss: 0.21127766370773315
Epoch 3 loss: 0.20121970772743225
Epoch 4 loss: 0.19375291466712952
Epoch 5 loss: 0.18759869039058685
Epoch 6 loss: 0.18283958733081818
Epoch 7 loss: 0.17909470200538635
Epoch 8 loss: 0.17629297077655792
Epoch 9 loss: 0.1739780753850937
Epoch 10 loss: 0.17155811190605164
Epoch 11 loss: 0.1695699691772461
Epoch 12 loss: 0.16842646896839142
Epoch 13 loss: 0.1673978865146637
Epoch 14 loss: 0.16583457589149475
Epoch 15 loss: 0.1656794399023056
Epoch 16 loss: 0.1639333963394165
Epoch 17 loss: 0.1634102314710617
Epoch 18 loss: 0.16268375515937805
Epoch 19 loss: 0.16166573762893677

[i-chaochen]

Hi @stepi9 could you try it with export XLA_FLAGS="--xla_gpu_enable_cublaslt=true ?

[stepi9]

Hi @stepi9 could you try it with export XLA_FLAGS="--xla_gpu_enable_cublaslt=true ?

Hi @i-chaochen,
That options does not seem to have any immediate effect on the bug.

With export XLA_FLAGS="--xla_gpu_enable_cublaslt=true I still get the GEMM mismatch warnings.

With export XLA_FLAGS="--xla_gpu_enable_cublaslt=true --xla_gpu_autotune_level=0" I do not get the warning, but the model does not learn.

[i-chaochen]

@stepi9 Thanks for the feedback, could you try the newer rocm? e.g., ROCm6.4.2?

[stepi9]

@i-chaochen unfortunately I am not able to update the GPU drivers on this machine, since I am not the administrator