fix: Remove import of anemoi training in compile (#705)

## Description
removed a leftover import of anemoi.training in anemoi models. this
created a circular dependancy which broke some downstream CI tests.

I had to refactor the code, because now we import torc.nn.module and not
pl.LightningModule, so the syntax is different

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

Author: cathalobrien

models/src/anemoi/models/utils/compile.py

@@ -8,16 +8,14 @@
 # nor does it submit to any jurisdiction.
 
 import logging
-from functools import reduce
 from importlib.util import find_spec
 
 import torch
 import torch_geometric
 from hydra.utils import get_class
 from numpy import unique
 from omegaconf import DictConfig
-
-from anemoi.training.train.tasks.base import BaseGraphModule
+from torch.nn import Module
 
 LOGGER = logging.getLogger(__name__)
 
@@ -58,7 +56,7 @@ def _meets_library_versions_for_compile() -> bool:
     return version_req and has_triton
 
 
-def mark_for_compilation(model: BaseGraphModule, compile_config: DictConfig | None) -> BaseGraphModule:
+def mark_for_compilation(model: Module, compile_config: DictConfig | None) -> Module:
     """Marks modules within 'model' for compilation, according to 'compile_config'.
 
     Modules are not compiled here. The compilation will occur
@@ -83,20 +81,11 @@ def mark_for_compilation(model: BaseGraphModule, compile_config: DictConfig | No
             options = entry.get("options", default_compile_options)
 
             LOGGER.debug("%s will be compiled with the following options: %s", str(module), str(options))
-            compiled_module = torch.compile(module, **options)  # Note: the module is not compiled yet
+            module.forward = torch.compile(module.forward, **options)  # Note: the function is not compiled yet
             # It is just marked for JIT-compilation later
             # It will be compiled before its first forward pass
             compiled_modules.append(entry.module)
 
-            # Update the model with the new 'compiled' module
-            # go from "anemoi.models.layers.conv.GraphTransformerConv"
-            # to obj(anemoi.models.layers.conv)
-            parts = name.split(".")
-            parent = reduce(getattr, parts[:-1], model)
-            # then set obj(anemoi.models.layers.conv).GrapTransformerConv = compiled_module
-            LOGGER.debug("Replacing %s with a compiled version", str(parts[-1]))
-            setattr(parent, parts[-1], compiled_module)
-
     LOGGER.info("The following modules will be compiled: %s", str(unique(compiled_modules)))
 
     return model

models/tests/utils/test_compile.py

@@ -109,7 +109,7 @@ def test_compile() -> None:
     result_compiled = ln_compiled.forward(x_in, cond)
 
     # check the function was compiled
-    assert hasattr(ln_compiled, "_compile_kwargs")
+    assert hasattr(ln_compiled.forward, "_torchdynamo_orig_callable")
 
     # check the result of the compiled function matches the uncompiled result
     assert torch.allclose(result, result_compiled)
@@ -145,7 +145,7 @@ def test_compile_layer_kernel() -> None:
     result_compiled = mhsa_compiled.forward(x, shapes, batch_size)
 
     # check the function was compiled
-    assert hasattr(mhsa_compiled.projection, "_compile_kwargs")
+    assert hasattr(mhsa_compiled.projection.forward, "_torchdynamo_orig_callable")
 
     # check the result of the compiled function matches the uncompiled result
     assert torch.allclose(result, result_compiled)