disabling DeePMD TF32 calculation in LAMMPS

[yjw0510]

Dear DeePMD users,

I'm trying to use my DeePMD model to do molecular dynamics simulation in LAMMPS.

I have two different types of GPUs (RTX 3090 and A100).

I found that both GPUs support TF32 calculations.
However, the lack of precision does make an undesired effect in the simulation so I want to disable them.

Generally, NVIDIA_TF32_OVERRIDE=0 should turn off TF32 but it seems Tensorflow does not do that.
(https://servicedesk.surf.nl/wiki/display/WIKI/Deep+Learning+on+A100+GPUs)

I'm using a compressed model using DP_INTERFACE_PREC=low

Slurm message

with CUBLASLT_LOG_LEVEL=5, I found the log below.

[njzjz]

For C++ applications, one needs to call tensorflow::enable_tensor_float_32_execution(false) to disable tensor float 32.
See https://github.com/tensorflow/tensorflow/blob/5dcfc51118817f27fad5246812d83e5dccdc5f72/tensorflow/core/platform/tensor_float_32_utils.h#L21.

[yjw0510]

Thank you for the answer. Then I think the only way is to recompile the code with the tensorflow::enable_tensor_float_32_execution(false) Is there any way to disable it without recompiling and without touching the original DeePMD code as much as possible?

[njzjz]

LAMMPS provide a way to load a external DSO, where you can call TF. See https://docs.lammps.org/plugin.html