Let AOTAutogradCache cache graphs with custom autograd.Functions

UMA's forward contains custom autograd.Functions. AOTAutogradCache refuses to
cache any graph containing one unless
torch._functorch.config.autograd_cache_allow_custom_autograd_functions is set,
so today the largest graph misses the cache and is retraced and recompiled from
scratch on every process start. Warm start ends up paying most of a cold
compile.

Turning it on halves warm compile time. On uma-s-1p2 with 1000 atoms, a second
process run against caches populated by the first:

  autograd_cache_allow_custom_autograd_functions=False   22.6s
  autograd_cache_allow_custom_autograd_functions=True    10.2s

which is a >2x reduction (3 runs each, 22.79/22.35/22.69 vs 10.86/9.92/9.76,
measured as the union of top-level compile events from tlparse). Cold start is
unaffected, 53.5s vs 54.0s, within run-to-run noise: the flag decides only
whether the second run can reuse the first run's work.

PyTorch intends to flip this flag to True by default, but that is waiting on
thorough internal testing across models. Nothing about UMA needs to wait for it:
the model's custom autograd.Functions are ordinary differentiable ops with no
state that would make a cached graph unsafe to reuse, so we can opt in now and
drop this line once the default changes.

This change was authored with the assistance of an AI coding agent.

Author: anijain2305

src/fairchem/core/units/mlip_unit/predict.py

@@ -539,6 +539,12 @@ def _lazy_init(self, data: AtomicData) -> None:
             # The model's scalars are fixed at inference, so this skips dynamo's
             # TensorifyScalarRestartAnalysis retrace during compile.
             torch._dynamo.config.specialize_float = True
+            # The model uses custom autograd.Functions, which AOTAutogradCache
+            # refuses to cache unless this is set, so without it the largest
+            # graph is re-compiled from scratch on every process start.
+            import torch._functorch.config as _functorch_config
+
+            _functorch_config.autograd_cache_allow_custom_autograd_functions = True
             self.model = torch.compile(self.model, dynamic=True)
 
         self.lazy_model_intialized = True