What is the expected force precision for a float32 model?

[krarmd]

I'm working on using a float32 model in CP2K (cp2k/cp2k#4716) and I noticed that a float32 model outputs the atomic_energies in float64 and forces in float32. Is this the expected behavior? According to a recent Allegro publication (https://arxiv.org/abs/2504.16068, mixed precision section in appendix) I'm assuming that the force output would also be float64. (quote of relevant text)

The learned representations are then passed through a readout MLP, also
in single precision, and finally cast back to double precision
when producing the energy prediction. The backward pass
used to compute forces and other derivative quantities follows
this sequence in reverse, preserving the same precision transi-
tions.

Should I interpret this as float64 atomic_energies are converted back to float32 and then the gradient of the float32 atomic_energies is used to produce the forces in float32?

[cw-tan]

Hi @krarmd , I would expect all inputs and outputs (energies, forces) of the model to be in float64 even if the model internals are in float32. How are you inspecting this behavior? If there is a reproducer, I can try to look into it.

[krarmd]

I recompile CP2K with additional print statements to the get_data_ptr function (CP2K torch_c_api.cpp:45) to check the data type and found that the Allegro model outputs the atomic_energy as a Double and forces as float32 if the model_dtype is float32 (see the mentioned CP2K issue cp2k/cp2k#4716). This is a code section from CP2K manybody_allegro.F portion that runs the forward pass and returns an output dictionary.

      CALL torch_dict_create(outputs)
      CALL torch_model_forward(allegro_data%model, inputs, outputs)
      pot_allegro = 0.0_dp

      CALL torch_dict_get(outputs, "atomic_energy", atomic_energy_tensor)
      CALL torch_dict_get(outputs, "forces", forces_tensor)
      IF (allegro%do_allegro_sp) THEN
         CALL torch_tensor_data_ptr(atomic_energy_tensor, atomic_energy_sp)
         CALL torch_tensor_data_ptr(forces_tensor, forces_sp)
         allegro_data%force(:, :) = REAL(forces_sp(:, :), kind=dp)*allegro%unit_forces_val
         DO iat_use = 1, SIZE(unique_list_a)
            i = unique_list_a(iat_use)
            pot_allegro = pot_allegro + REAL(atomic_energy_sp(1, i), kind=dp)*allegro%unit_energy_val
         END DO
         DEALLOCATE (new_edge_cell_shifts_sp, pos_sp)

The line:
CALL torch_tensor_data_ptr(forces_tensor, forces_sp)
runs a C++ function that checks if the model_dtype matches the output tensor dtype. Here is a code snippet of the print statements I added to check the atomic_energy and forces dytpe.

static void *get_data_ptr(const torch_c_tensor_t *tensor,
                          const torch::Dtype dtype, const int ndims,
                          int64_t sizes[]) {
  std::cout << "Scalar type: " << tensor->scalar_type() << std::endl;
  std::cout << "dtype: " << dtype << std::endl;
  assert(tensor->scalar_type() == dtype);
  assert(tensor->ndimension() == ndims);

I found that when I ran a float32 model the atomic_energy was Double but forces was float32.
The assert(tensor->scalar_type() == dtype); command would fail I didn't load the forces as a single-precision array.

[krarmd]

I may have found the culprit. CP2K was setting the positions as float32 because the model dtype was float32. I think that was a remnant of the previous NequiP code structure. I am able to get Double precision for forces and atomic_energy by disabling CP2K's conversion of the positions from double to single precision which may be how the forces were being set to float32 in the ForceStressOutput function. Thanks for confirming that forces and atomic_energies should be in float64, it helped me trace the issue back to CP2K.