feat(models): Triton GraphTransformer (#631)

## Description
Introduce custom Triton Kernel for a fast, memory efficient and
deterministic GraphTransformer.

Benchmarking results comparing the current GraphTransformer
implementation (PyG) against PyG compiled and the Triton kernel on
different graphs, measured in ms per fwd + bwd:

All runs use era = n320, h = lvl6-multi-scale, dtype = fp16.

A100

| Graph type |    PyG | PyG (compiled) | Triton   |
|:------------|-------:|---------------:|---------:|
| era_to_h    |  53.09 |          15.86 | **6.81** |
| h_to_h      |  20.32 |           4.45 | **1.82** |
| h_to_era    | 100.44 |          20.53 | **11.19** |
| dop_enc     |  60.44 |          17.90 | **9.08** |

H200

| Graph type |   PyG | PyG (compiled) |   Triton |
| :--------- | ----: | -------------: | -------: |
| era_to_h   | 23.97 |           6.02 | **2.77** |
| h_to_h     |  9.64 |           1.69 | **0.83** |
| h_to_era   | 46.99 |           8.16 | **4.87** |
| dop_enc    | 26.62 |           6.85 | **3.86** |

Exemplary impact on training for an o800 lvl7 GraphTransformer model in
bf16 on H200s:
pyg compiled: 3.90 s/it → Triton: 2.62 s/it (×1.5 speedup)

## What problem does this change solve?
Current PyG GraphTransformer suffers from high memory footprint due to
per-edge computations, multiple kernel launches and stability issues
with atomics in low precision.
[Compiling](#181) can help
adress some of these issues, but one can do better with a custom triton
kernel.

##  Additional notes ##
Checked outputs and gradients match pyg version + small training sanity
check:
[mlflow](https://mlflow.ecmwf.int/#/metric?runs=[%224803881259ff447cbbb4ff4436926823%22,%22b37aee50b3664afa9fc4aee6b557d71b%22]&metric=%22train_mse_loss_step%22&experiments=[%2245%22]&plot_metric_keys=%5B%22train_mse_loss_step%22%5D&plot_layout={%22autosize%22:true,%22xaxis%22:{%22autorange%22:true,%22type%22:%22linear%22,%22range%22:[-1035.9219986020375,960.4671241229228]},%22yaxis%22:{%22range%22:[-0.6010538821019717,6.501633332403755],%22autorange%22:true,%22type%22:%22linear%22}}&x_axis=step&y_axis_scale=linear&line_smoothness=1&show_point=false&deselected_curves=[]&last_linear_y_axis_range=[])

TODO
- [x] make GT backend configurable, pyg as fallback
- [x] add pytests

***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)

---------

Co-authored-by: Dieter Van den Bleeken <[email protected]>
Co-authored-by: Simon Lang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ana Prieto Nemesio <[email protected]>

Author: 5 people

.github/workflows/integration-tests-hpc.yml

@@ -43,6 +43,7 @@ jobs:
             pip install --upgrade pip
             pip install -e ./training[all,tests] -e ./models[all,tests] -e ./graphs[all,tests]
             python3 -m pytest -v training/tests/integration --slow
+            python3 -m pytest -v models/tests/integration --slow
             deactivate
             rm -rf $REPO_NAME
           sbatch_options: |

models/pytest.ini

@@ -3,5 +3,6 @@ markers =
     data_dependent: marks tests depending on data (deselect with '-m "not data_dependent"')
     auth: marks tests that require authentication (deselect with '-m "not auth"')
     gpu: marks tests that require a GPU (deselect with '-m "not gpu"')
+    slow: mark test as slow (skipped unless --slow is used)
 
 tmp_path_retention_policy = none

models/src/anemoi/models/layers/block.py

@@ -34,6 +34,8 @@
 from anemoi.models.layers.conv import GraphConv
 from anemoi.models.layers.conv import GraphTransformerConv
 from anemoi.models.layers.mlp import MLP
+from anemoi.models.triton.gt import GraphTransformerFunction
+from anemoi.models.triton.utils import edge_index_to_csc
 from anemoi.utils.config import DotDict
 
 LOGGER = logging.getLogger(__name__)
@@ -443,6 +445,7 @@ def __init__(
         qk_norm: bool = False,
         update_src_nodes: bool = False,
         layer_kernels: DotDict,
+        graph_attention_backend: str = "triton",
         **kwargs,
     ) -> None:
         """Initialize GraphTransformerBlock.
@@ -466,6 +469,8 @@ def __init__(
         layer_kernels : DotDict
             A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
             Defined in config/models/<model>.yaml
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
         super().__init__(**kwargs)
 
@@ -483,8 +488,6 @@ def __init__(
         self.lin_self = Linear(in_channels, num_heads * self.out_channels_conv, bias=bias)
         self.lin_edge = Linear(edge_dim, num_heads * self.out_channels_conv)  # , bias=False)
 
-        self.conv = GraphTransformerConv(out_channels=self.out_channels_conv)
-
         self.projection = Linear(out_channels, out_channels)
 
         if self.qk_norm:
@@ -499,6 +502,19 @@ def __init__(
             Linear(hidden_dim, out_channels),
         )
 
+        self.graph_attention_backend = graph_attention_backend
+        assert self.graph_attention_backend in [
+            "triton",
+            "pyg",
+        ], f"Backend {self.graph_attention_backend} not supported for GraphTransformerBlock, valid options are 'triton' and 'pyg'"
+
+        if self.graph_attention_backend == "triton":
+            LOGGER.info(f"{self.__class__.__name__} using triton graph attention backend.")
+            self.conv = GraphTransformerFunction.apply
+        else:
+            LOGGER.warning(f"{self.__class__.__name__} using pyg graph attention backend, consider using 'triton'.")
+            self.conv = GraphTransformerConv(out_channels=self.out_channels_conv)
+
     def run_node_dst_mlp(self, x, **layer_kwargs):
         return self.node_dst_mlp(self.layer_norm_mlp_dst(x, **layer_kwargs))
 
@@ -555,37 +571,50 @@ def shard_qkve_heads(
 
         return query, key, value, edges
 
-    def attention_block(
+    def apply_gt(
         self,
         query: Tensor,
         key: Tensor,
         value: Tensor,
         edges: Tensor,
         edge_index: Adj,
         size: Union[int, tuple[int, int]],
-        num_chunks: int,
     ) -> Tensor:
         # self.conv requires size to be a tuple
         conv_size = (size, size) if isinstance(size, int) else size
 
+        if self.graph_attention_backend == "triton":
+            csc, perm, reverse = edge_index_to_csc(edge_index, num_nodes=conv_size, reverse=True)
+            edges_csc = edges.index_select(0, perm)
+            args_conv = (edges_csc, csc, reverse)
+        else:
+            args_conv = (edges, edge_index, conv_size)
+
+        return self.conv(query, key, value, *args_conv)
+
+    def attention_block(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        edges: Tensor,
+        edge_index: Adj,
+        size: Union[int, tuple[int, int]],
+        num_chunks: int,
+    ) -> Tensor:
+        # split 1-hop edges into chunks, compute self.conv chunk-wise
         if num_chunks > 1:
-            # split 1-hop edges into chunks, compute self.conv chunk-wise
             edge_attr_list, edge_index_list = sort_edges_1hop_chunks(
                 num_nodes=size, edge_attr=edges, edge_index=edge_index, num_chunks=num_chunks
             )
             # shape: (num_nodes, num_heads, out_channels_conv)
             out = torch.zeros((*query.shape[:-1], self.out_channels_conv), device=query.device)
             for i in range(num_chunks):
-                out += self.conv(
-                    query=query,
-                    key=key,
-                    value=value,
-                    edge_attr=edge_attr_list[i],
-                    edge_index=edge_index_list[i],
-                    size=conv_size,
+                out += self.apply_gt(
+                    query=query, key=key, value=value, edges=edge_attr_list[i], edge_index=edge_index_list[i], size=size
                 )
         else:
-            out = self.conv(query=query, key=key, value=value, edge_attr=edges, edge_index=edge_index, size=conv_size)
+            out = self.apply_gt(query=query, key=key, value=value, edges=edges, edge_index=edge_index, size=size)
 
         return out
 
@@ -635,6 +664,7 @@ def __init__(
         update_src_nodes: bool = False,
         layer_kernels: DotDict,
         shard_strategy: str = "edges",
+        graph_attention_backend: str = "triton",
         **kwargs,
     ) -> None:
         """Initialize GraphTransformerBlock.
@@ -662,6 +692,8 @@ def __init__(
             Defined in config/models/<model>.yaml
         shard_strategy: str, by default "edges"
             Strategy to shard tensors
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
 
         super().__init__(
@@ -674,6 +706,7 @@ def __init__(
             bias=bias,
             qk_norm=qk_norm,
             update_src_nodes=update_src_nodes,
+            graph_attention_backend=graph_attention_backend,
             **kwargs,
         )
 
@@ -791,6 +824,7 @@ def __init__(
         qk_norm: bool = False,
         update_src_nodes: bool = False,
         layer_kernels: DotDict,
+        graph_attention_backend: str = "triton",
         **kwargs,
     ) -> None:
         """Initialize GraphTransformerBlock.
@@ -814,6 +848,8 @@ def __init__(
         layer_kernels : DotDict
             A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
             Defined in config/models/<model>.yaml
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
 
         super().__init__(
@@ -826,6 +862,7 @@ def __init__(
             bias=bias,
             qk_norm=qk_norm,
             update_src_nodes=update_src_nodes,
+            graph_attention_backend=graph_attention_backend,
             **kwargs,
         )
 

models/src/anemoi/models/layers/mapper.py

@@ -216,6 +216,7 @@ def __init__(
         cpu_offload: bool = False,
         layer_kernels: DotDict = None,
         shard_strategy: str = "edges",
+        graph_attention_backend: str = "triton",
     ) -> None:
         """Initialize GraphTransformerBaseMapper.
 
@@ -254,6 +255,8 @@ def __init__(
             Defined in config/models/<model>.yaml
         shard_strategy : str, optional
             Strategy to shard tensors, by default "edges"
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
         super().__init__(
             in_channels_src=in_channels_src,
@@ -282,6 +285,7 @@ def __init__(
             qk_norm=qk_norm,
             layer_kernels=self.layer_factory,
             shard_strategy=shard_strategy,
+            graph_attention_backend=graph_attention_backend,
         )
 
         self.offload_layers(cpu_offload)
@@ -539,6 +543,7 @@ def __init__(
         cpu_offload: bool = False,
         layer_kernels: DotDict = None,
         shard_strategy: str = "edges",
+        graph_attention_backend: str = "triton",
     ) -> None:
         """Initialize GraphTransformerForwardMapper.
 
@@ -574,6 +579,8 @@ def __init__(
             A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
         shard_strategy : str, optional
             Strategy to shard tensors, by default "edges"
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
         super().__init__(
             in_channels_src=in_channels_src,
@@ -592,6 +599,7 @@ def __init__(
             dst_grid_size=dst_grid_size,
             layer_kernels=layer_kernels,
             shard_strategy=shard_strategy,
+            graph_attention_backend=graph_attention_backend,
         )
 
         self.emb_nodes_src = self.layer_factory.Linear(self.in_channels_src, self.hidden_dim)
@@ -643,6 +651,7 @@ def __init__(
         cpu_offload: bool = False,
         layer_kernels: DotDict = None,
         shard_strategy: str = "edges",
+        graph_attention_backend: str = "triton",
     ) -> None:
         """Initialize GraphTransformerBackwardMapper.
 
@@ -683,6 +692,8 @@ def __init__(
             Defined in config/models/<model>.yaml
         shard_strategy : str, optional
             Strategy to shard tensors, by default "edges"
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
         super().__init__(
             in_channels_src=in_channels_src,
@@ -701,6 +712,7 @@ def __init__(
             dst_grid_size=dst_grid_size,
             layer_kernels=layer_kernels,
             shard_strategy=shard_strategy,
+            graph_attention_backend=graph_attention_backend,
         )
 
         self.node_data_extractor = nn.Sequential(

models/src/anemoi/models/layers/processor.py

@@ -408,6 +408,7 @@ def __init__(
         qk_norm: bool = False,
         cpu_offload: bool = False,
         layer_kernels: DotDict,
+        graph_attention_backend: str = "triton",
         **kwargs,
     ) -> None:
         """Initialize GraphTransformerProcessor.
@@ -441,6 +442,8 @@ def __init__(
         layer_kernels : DotDict
             A dict of layer implementations e.g. layer_kernels.Linear = "torch.nn.Linear"
             Defined in config/models/<model>.yaml
+        graph_attention_backend: str, by default "triton"
+            Backend to use for graph transformer conv, options are "triton" and "pyg"
         """
         super().__init__(
             num_channels=num_channels,
@@ -465,6 +468,7 @@ def __init__(
             edge_dim=self.edge_dim,
             layer_kernels=self.layer_factory,
             qk_norm=qk_norm,
+            graph_attention_backend=graph_attention_backend,
         )
 
         self.offload_layers(cpu_offload)

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

@@ -40,7 +40,7 @@ class GraphTransformerDecoderSchema(TransformerModelComponent):
     @model_validator(mode="after")
     def check_valid_extras(self) -> Any:
         # This is a check to allow backwards compatibilty of the configs, as the extra fields are not required.
-        allowed_extras = {"shard_strategy": str}
+        allowed_extras = {"shard_strategy": str, "graph_attention_backend": str}
         extras = getattr(self, "__pydantic_extra__", {}) or {}
         for extra_field, value in extras.items():
             if extra_field not in allowed_extras:

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

@@ -38,7 +38,7 @@ class GraphTransformerEncoderSchema(TransformerModelComponent):
     @model_validator(mode="after")
     def check_valid_extras(self) -> Any:
         # This is a check to allow backwards compatibilty of the configs, as the extra fields are not required.
-        allowed_extras = {"shard_strategy": str}
+        allowed_extras = {"shard_strategy": str, "graph_attention_backend": str}
         extras = getattr(self, "__pydantic_extra__", {}) or {}
         for extra_field, value in extras.items():
             if extra_field not in allowed_extras:

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

@@ -43,6 +43,21 @@ class GraphTransformerProcessorSchema(TransformerModelComponent):
     qk_norm: bool = Field(example=False)
     "Normalize the query and key vectors. Default to False."
 
+    @model_validator(mode="after")
+    def check_valid_extras(self) -> Any:
+        # This is a check to allow backwards compatibilty of the configs, as the extra fields are not required.
+        allowed_extras = {"graph_attention_backend": str}
+        extras = getattr(self, "__pydantic_extra__", {}) or {}
+        for extra_field, value in extras.items():
+            if extra_field not in allowed_extras:
+                msg = f"Extra field '{extra_field}' is not allowed. Allowed fields are: {list(allowed_extras.keys())}."
+                raise ValueError(msg)
+            if not isinstance(value, allowed_extras[extra_field]):
+                msg = f"Extra field '{extra_field}' must be of type {allowed_extras[extra_field].__name__}."
+                raise TypeError(msg)
+
+        return self
+
 
 class TransformerProcessorSchema(TransformerModelComponent):
     target_: Literal["anemoi.models.layers.processor.TransformerProcessor"] = Field(..., alias="_target_")