Undocumented PolarMACE dtype inconsistency (float32 vs float64) and unclear intended precision

[godot11]

The released (excellent) MACE-POLAR models appear to use float32, while other foundation models seem to use float64:

import torch

for mname in [
    "MACE-POLAR-1-L.model", "MACE-POLAR-1-M.model", "MACE-POLAR-1-S.model", 
    "MACE-matpes-pbe-omat-ft.model", "MACE-matpes-r2scan-omat-ft.model", "MACE-omol-0-extra-large-1024.model", 
    "mace-omat-0-medium.model", "mace_agnesi_medium.model", "mace_agnesi_small.model",
]:
    model = torch.load(f"models/{mname}")
    dtype = next(model.parameters()).dtype
    print(f"{mname}  ---  {dtype}")

Output:

MACE-POLAR-1-L.model  ---  torch.float32              
MACE-POLAR-1-M.model  ---  torch.float32              
MACE-POLAR-1-S.model  ---  torch.float32              
MACE-matpes-pbe-omat-ft.model  ---  torch.float64     
MACE-matpes-r2scan-omat-ft.model  ---  torch.float64  
MACE-omol-0-extra-large-1024.model  ---  torch.float64
mace-omat-0-medium.model  ---  torch.float64          
mace_agnesi_medium.model  ---  torch.float64          
mace_agnesi_small.model  ---  torch.float64           

I couldn't find this being mentioned of this in the documentation, and it may lead to confusion or issues with MACE workflows. For example, eval_configs.py fails with PolarMACE and --default-dtype=float64, which is the default option (see ACEsuit/mace#1429); while mace_run_train accepts --default-dtype=float64 and produces a float64 model.

The documentation also suggests training with --dtype=float64, which makes the intended precision of PolarMACE unclear.

This also made me confused whether PolarMACE models are inherently float32 (i.e., trained and intended to operate at this precision), or if the float32 dtype arise from a technicality, while the underlying training or expected usage is float64, as seem to be the case with other models? If the polar models are fundamentally float32, then is it still fine to conduct finetuning in float64, considering that it would not only adapt the model but also upscale it's numerical precision?

[ilyes319]

Hey @godot11, you are right the models were trained indeed with float32, which is different from the other models. Though you can use them with float64 for simulations, which should not be a problem, but be careful and validate your results by cross checking with float32 results.

Fine-tuning in float64 will indeed upscale the precision, which might lead to different behaviour, which also you should be careful of. But in general, things are well behaved from my experience.