cagrikymk / JAX-ReaxFF

JAX-ReaxFF: A Gradient Based Framework for Extremely Fast Optimization of Reactive Force Fields
GNU General Public License v3.0
53 stars 22 forks source link

CUDA_ERROR_MISALIGNED_ADDRESS while running JAX-ReaxFF #23

Closed houdingyu closed 1 week ago

houdingyu commented 3 weeks ago

Hi, recently I am using JAX-ReaxFF to optimize some Si-N parameters for my system. However, I ran into the error "CUDA_ERROR_MISALIGNED_ADDRESS" and I did not understand why. If delete some items in the trainset.in file, then the program can run normally. However, those items are necessary for my system, and I cannot detect any problem in those items. I have attached my input files. Could you please have a look at why JAX-ReaxFF can not run with the trainset_err.in file? Thank you very much! Si_N_opt.zip

houdingyu commented 2 weeks ago

In addition, I found that the CPU version can run normally.

cagrikymk commented 1 week ago

Hello, sorry for the delayed fix. This issue stems from XLA (the compiler used in JAX). Upgrading JAX to >=0.4.26 fixes the issue. I will push an update soon.