Energy only evaluations with pretrained_mlip

[InfluenceFunctional]

What would you like to report?

I'm using UMA to compute energies on large, heterogeneous batches of molecular crystals, as fast as possible.

predictor = pretrained_mlip.load_predict_unit(model_path)
uma_batch = atomicdata_list_to_batch([AtomicData.from_ase(atoms, task_name='omc') for atoms in data_list])
out = predictor.predict(uma_batch)

The force calculations seem to inevitably call .backward() within the model, and my profiler is telling me that this is taking up a very large fraction of the total compute time (local RTX 5080, cluster A100). This is causing a significant slowdown in my overall workflows, as I do not need the forces or stresses.

I have played with the overrides in load_predict_unit, as well as a number of internal model flags, to try to turn off force calculation. These either have had no effect, or broken the model.

Is it possible to deactivate force calculations / internal backprop, to get the fastest possible energy-only evaluations?

[misko]

Hi @InfluenceFunctional ,
I had this same issue as well when trying to do energy calculations. I wasn't sure if it would be useful to others. I don't think it's super easy to do in the current code without changes.

We don't have an official solution yet, but if you are blocked try #1714 .

Pass energy_only=True for inference_settings to use it, if it does help you please let me know so we can prioritize this

[InfluenceFunctional]

Thanks for the quick reply! It would indeed be very useful for me as I'm using it on-line in a training workflow where UMA is the bottleneck. I tried hacking something like this together:

predictor.tasks.pop("omc_forces")

predictor.tasks.pop("omc_stress")

predictor.dataset_to_tasks['omc'].pop(-1)

predictor.dataset_to_tasks['omc'].pop(-1)

predictor.model.module.backbone.regress_forces = False

predictor.model.module.backbone.regress_stress = False

which yielded the below (with our without changing the tasks, just from changing the regress options).

I think the models may be hard coded to expect some forces information? I didn't dig past this point.

Traceback (most recent call last):
File "C:\Users\mikem\AppData\Local\Programs\PyCharm\plugins\python\helpers-pro\pydevd_asyncio\pydevd_asyncio_utils.py", line 117, in _exec_async_code
result = func()
^^^^^^
File "", line 1, in
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\units\mlip_unit\predict.py", line 53, in collated_predict
preds = predict_fn(predict_unit, data, undo_element_references)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\units\mlip_unit\predict.py", line 272, in predict
output = self.model(data_device)
^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\optim\swa_utils.py", line 244, in forward
return self.module(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\models\base.py", line 186, in forward
out.update(self.output_heads[k](data, emb))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\common\utils.py", line 80, in cls_method
return f(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\utils_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\models\uma\escn_moe.py", line 260, in forward
head_output = self.head(data, emb)
^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\common\utils.py", line 80, in cls_method
return f(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\torch\utils_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\models\uma\escn_md.py", line 661, in forward
[data["pos_original"], emb["displacement"]],
~~~~^^^^^^^^^^^^^^^^
File "C:\Users\mikem\venvs\csd_mxt_gfn\Lib\site-packages\fairchem\core\datasets\atomic_data.py", line 567, in getitem
return getattr(self, idx)
^^^^^^^^^^^^^^^^^^
AttributeError: 'AtomicData' object has no attribute 'pos_original'