feat: compile transformer gnn (#181)

Compile GraphTransformerConv



<!-- readthedocs-preview anemoi-training start -->
----
📚 Documentation preview 📚:
https://anemoi-training--181.org.readthedocs.build/en/181/

<!-- readthedocs-preview anemoi-training end -->

<!-- readthedocs-preview anemoi-graphs start -->
----
📚 Documentation preview 📚:
https://anemoi-graphs--181.org.readthedocs.build/en/181/

<!-- readthedocs-preview anemoi-graphs end -->

<!-- readthedocs-preview anemoi-models start -->
----
📚 Documentation preview 📚:
https://anemoi-models--181.org.readthedocs.build/en/181/

<!-- readthedocs-preview anemoi-models end -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Harrison Cook <[email protected]>
Co-authored-by: Icedoom <[email protected]>
Co-authored-by: Cathal OBrien <[email protected]>
Co-authored-by: Matthew Chantry <[email protected]>
Co-authored-by: Ana Prieto Nemesio <[email protected]>

Author: 7 people

models/src/anemoi/models/schemas/models.py

@@ -12,7 +12,9 @@
 import logging
 from enum import Enum
 from typing import Annotated
+from typing import Any
 from typing import Literal
+from typing import Optional
 from typing import Union
 
 from pydantic import BaseModel as PydanticBaseModel
@@ -226,6 +228,8 @@ class BaseModelSchema(PydanticBaseModel):
         discriminator="target_",
     )
     "GNN decoder schema."
+    compile: Optional[list[dict[str, Any]]] = Field(None)
+    "Modules to be compiled"
 
 
 class NoiseInjectorSchema(BaseModel):

models/src/anemoi/models/utils/__init__.py

@@ -0,0 +1,8 @@
+# (C) Copyright 2025 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.

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

@@ -0,0 +1,102 @@
+# (C) Copyright 2025 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# 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
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _get_compile_entry(module: str, compile_config: DictConfig) -> DictConfig | None:
+    """Search the compile config for an entry c module name.
+
+    module: str -> full module name e.g. 'anemoi.models.layers.conv.GraphTransformerConv'
+    compile_config : DictConfig -> The 'compile' entry within the models config
+
+    returns: None, if 'module' is not listed within 'compile_config'. Otherwise returns the modules entry.
+
+    """
+    for entry in compile_config:
+        if get_class(entry["module"]) is type(module):
+            return entry
+
+    return None
+
+
+def _meets_library_versions_for_compile() -> bool:
+    """Returns True if minimum library versions for compilation in Anemoi is met."""
+    has_triton = True
+    if find_spec("triton") is None:
+        msg = "Triton not installed! Consider installing Triton to "
+        msg += "enable compilation and improve speed and memory usage."
+        LOGGER.warning(msg)
+        has_triton = False
+
+    version_req = torch.__version__ >= "2.6" and torch_geometric.__version__ >= "2.6"
+
+    if not version_req:
+        msg = "Minimum library versions for compilation not met. "
+        msg += f"torch: v{torch.__version__}<2.6 or torch_geometric: v{torch_geometric.__version__}<2.6. "
+        msg += "Please upgrade these libraries to enable compilation."
+        LOGGER.warning(msg)
+
+    return version_req and has_triton
+
+
+def mark_for_compilation(model: BaseGraphModule, compile_config: DictConfig | None) -> BaseGraphModule:
+    """Marks modules within 'model' for compilation, according to 'compile_config'.
+
+    Modules are not compiled here. The compilation will occur
+    automatically before the first forward iteration.
+
+    returns an updated model, with modules marked for compilation
+    """
+    if compile_config is None:
+        return model
+
+    if not _meets_library_versions_for_compile():
+        return model
+
+    default_compile_options = {}
+    compiled_modules = []
+
+    # Loop through all modules
+    for name, module in model.named_modules():
+        entry = _get_compile_entry(module, compile_config)
+        # entry is 'None' if compilation was not requested for this module
+        if entry is not None:
+            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
+            # 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

@@ -0,0 +1,151 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import logging
+
+import torch
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+
+from anemoi.models.layers.attention import MultiHeadSelfAttention
+from anemoi.models.layers.normalization import ConditionalLayerNorm
+from anemoi.models.layers.utils import load_layer_kernels
+from anemoi.models.utils.compile import _get_compile_entry
+from anemoi.models.utils.compile import _meets_library_versions_for_compile
+from anemoi.models.utils.compile import mark_for_compilation
+
+LOGGER = logging.getLogger(__name__)
+
+
+def graphtransformer_compile_config() -> None:
+    return OmegaConf.create(
+        {
+            "compile": [
+                {
+                    "module": "anemoi.models.layers.conv.GraphTransformerConv",
+                },
+            ],
+        }
+    )
+
+
+def layer_kernel_compile_config() -> None:
+    return OmegaConf.create(
+        {
+            "compile": [
+                {
+                    "module": "torch.nn.Linear",
+                },
+            ],
+        }
+    )
+
+
+def graphtransformer_ens_compile_config() -> None:
+    return OmegaConf.create(
+        {
+            "compile": [
+                {
+                    "module": "anemoi.models.layers.conv.GraphTransformerConv",
+                },
+                {
+                    "module": "anemoi.models.layers.normalization.ConditionalLayerNorm",
+                    "options": {
+                        "dynamic": False,
+                    },
+                },
+            ],
+        }
+    )
+
+
+def test_compile_config_no_match() -> None:
+    """Tests that _get_compile_entry() returns None when no match is found."""
+    cfg = graphtransformer_compile_config()
+
+    num_channels = 64
+    cond_shape = 16
+    model = ConditionalLayerNorm(num_channels, condition_shape=cond_shape)
+    result = _get_compile_entry(model, cfg.compile)
+
+    assert result is None
+
+
+def test_compile_config_match() -> None:
+    """Tests that _get_compile_entry() returns a dict when a match is found."""
+    cfg = graphtransformer_ens_compile_config()
+
+    num_channels = 64
+    cond_shape = 16
+    model = ConditionalLayerNorm(num_channels, condition_shape=cond_shape)
+    result = _get_compile_entry(model, cfg.compile)
+
+    assert type(result) is DictConfig
+
+
+def test_compile() -> None:
+
+    # Skip this test if library versions aren't met
+    if not _meets_library_versions_for_compile():
+        LOGGER.warning("triton not installed. skipping 'test_compile.py::test_compile'")
+        return
+
+    num_channels = 64
+    cond_shape = 16
+    ln = ConditionalLayerNorm(num_channels, condition_shape=cond_shape)
+    x_in = torch.randn(num_channels)
+    cond = torch.randn(cond_shape)
+    result = ln.forward(x_in, cond)
+
+    cfg = graphtransformer_ens_compile_config()
+    ln_compiled = mark_for_compilation(ln, cfg.compile)
+
+    result_compiled = ln_compiled.forward(x_in, cond)
+
+    # check the function was compiled
+    assert hasattr(ln_compiled, "_compile_kwargs")
+
+    # check the result of the compiled function matches the uncompiled result
+    assert torch.allclose(result, result_compiled)
+
+
+def test_compile_layer_kernel() -> None:
+
+    # Skip this test if library versions aren't met
+    if not _meets_library_versions_for_compile():
+        LOGGER.warning("triton not installed. skipping 'test_compile.py::test_compile'")
+        return
+
+    cfg = layer_kernel_compile_config()
+    layer_kernels = load_layer_kernels(kernel_config={})
+
+    num_heads = 1
+    embed_dim = 64
+    dropout_p = 0.0
+    batch_size = 1
+    mhsa = MultiHeadSelfAttention(
+        num_heads,
+        embed_dim,
+        layer_kernels,
+        dropout_p=dropout_p,
+        attention_implementation="scaled_dot_product_attention",
+    )
+    mhsa_compiled = mark_for_compilation(mhsa, cfg.compile)
+
+    x = torch.randn(batch_size * 2, embed_dim, requires_grad=True)
+    shapes = [list(x.shape)]
+
+    result = mhsa.forward(x, shapes, batch_size)
+    result_compiled = mhsa_compiled.forward(x, shapes, batch_size)
+
+    # check the function was compiled
+    assert hasattr(mhsa_compiled.projection, "_compile_kwargs")
+
+    # check the result of the compiled function matches the uncompiled result
+    assert torch.allclose(result, result_compiled)

training/docs/user-guide/models.rst

@@ -191,3 +191,47 @@ configuration:
 This determines how many ensemble members are generated per device
 during training. Effective ensemble size is then the number of ensemble
 members per device times the number of GPUs per ensemble.
+
+*************
+ Compilation
+*************
+
+PyTorch supports JIT-compiliation of code. This can speed up execution
+and reduce peak memory usage. For more information, consult `the
+introduction to torch.compile
+<https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html>`__
+and `the official documentation
+<https://docs.pytorch.org/docs/stable/generated/torch.compile.html>`__.
+
+Compilation requires Triton. Normally Triton is pulled in as a
+dependancy when PyTorch is installed. Otherwise, Triton can be `built
+from source
+<https://github.com/triton-lang/triton?tab=readme-ov-file#install-from-source>`__
+. Compilation requires torch >= 2.6 and torch_geometric >= 2.6. If these
+versions are not met, or if Triton is not installed, then anemoi will
+run without compilation.
+
+Anemoi exposes 'torch.compile' at the module level through the model
+config. Below is an example:
+
+.. code:: yaml
+
+   #training/config/models/transformer_ens.yaml
+   compile:
+   - module: anemoi.models.layers.conv.GraphTransformerConv
+     options:
+       dynamic: false
+       mode: max-autotune
+   - module: anemoi.models.layers.normalization.ConditionalLayerNorm
+     options:
+       dynamic: false
+
+Under the 'compile' keyword, you provide a list of modules. These
+modules will be marked for compilation when the model is built. During
+their first forward pass, these modules will be compiled. No code
+modifications are required.
+
+You can optionally pass options to torch compile via the 'options'
+keyword. A full list of the possible options and their meanings can be
+found in the `torch.compile documentation
+<https://docs.pytorch.org/docs/stable/generated/torch.compile.html>`__.

training/src/anemoi/training/config/model/graphtransformer.yaml

@@ -68,6 +68,29 @@ trainable_parameters:
   hidden2data: 8
   hidden2hidden: 8 # GNN and GraphTransformer Processor only
 
+# Torch compile configuration
+# A list of modules present in the model, which will be compiled
+# You can optionally pass options to torch.compile with the 'options' key
+#
+# Below is an explanation of some common parameters to torch.compile
+# For a full list of possible parameters, consult the documenation for torch compile
+#       https://docs.pytorch.org/docs/stable/generated/torch.compile.html
+#
+# dynamic (bool): When True, it will try to avoid recompilation by generating
+#       as general a kernel as possible. But the performance of the general
+#       kernel might be worse. When False, it will generate a specific
+#       kernel for each input shape (until the configurable recompile
+#       limit has been hit), leading to possibly better performance but
+#       more recompilations
+# mode (str): Different compilation modes, allowing you to trade off
+#       compilation time versus potential performance. See the
+#       torch.compile documentation for list of possible modes
+# fullgraph (bool): When True, torch.compile will error when it hits a
+#       section of code it can't compile. When False, it will fallback to
+#       non-compiled ("eager") execution for those lines.
+# options (dict): a dict of further options which can be passed to torch.compile
+compile:
+  - module: anemoi.models.layers.conv.GraphTransformerConv
 
 attributes:
   edges:

training/src/anemoi/training/config/model/graphtransformer_diffusion.yaml

@@ -105,5 +105,35 @@ attributes:
   - edge_dirs
   nodes: []
 
+# Torch compile configuration
+# A list of modules present in the model, which will be compiled
+# You can optionally pass options to torch.compile with the 'options' key
+#
+# Below is an explanation of some common parameters to torch.compile
+# For a full list of possible parameters, consult the documenation for torch compile
+#       https://docs.pytorch.org/docs/stable/generated/torch.compile.html
+#
+# dynamic (bool): When True, it will try to avoid recompilation by generating
+#       as general a kernel as possible. But the performance of the general
+#       kernel might be worse. When False, it will generate a specific
+#       kernel for each input shape (until the configurable recompile
+#       limit has been hit), leading to possibly better performance but
+#       more recompilations
+# mode (str): Different compilation modes, allowing you to trade off
+#       compilation time versus potential performance. See the
+#       torch.compile documentation for list of possible modes
+# fullgraph (bool): When True, torch.compile will error when it hits a
+#       section of code it can't compile. When False, it will fallback to
+#       non-compiled ("eager") execution for those lines.
+# options (dict): a dict of further options which can be passed to torch.compile
+compile:
+  - module: anemoi.models.layers.conv.GraphTransformerConv
+    #options: # An example of setting torch.compile options
+      #dynamic: false
+      #mode: max-autotune
+  - module: anemoi.models.layers.normalization.ConditionalLayerNorm
+    options:
+      dynamic: false
+
 # Bounding configuration
 bounding: []