feature(models): Add model comm group to predict_step (#77)

* added optional model comm group to predict_step

* update changelog

* added **kwargs to predict step

models/CHANGELOG.md

@@ -18,6 +18,7 @@ Keep it human-readable, your future self will thank you!
 - Reduced memory usage when using chunking in the mapper [#84](https://github.com/ecmwf/anemoi-models/pull/84)
 - Added `supporting_arrays` argument, which contains arrays to store in checkpoints. [#97](https://github.com/ecmwf/anemoi-models/pull/97)
 - Add remappers, e.g. link functions to apply during training to facilitate learning of variables with a difficult distribution [#88](https://github.com/ecmwf/anemoi-models/pull/88)
+- 'predict\_step' can take an optional model comm group. [#77](https://github.com/ecmwf/anemoi-core/pull/77)
 
 ## [0.4.0](https://github.com/ecmwf/anemoi-models/compare/0.3.0...0.4.0) - Improvements to Model Design
 

models/src/anemoi/models/interface/__init__.py

@@ -8,9 +8,11 @@
 # nor does it submit to any jurisdiction.
 
 import uuid
+from typing import Optional
 
 import torch
 from hydra.utils import instantiate
+from torch.distributed.distributed_c10d import ProcessGroup
 from torch_geometric.data import HeteroData
 
 from anemoi.models.preprocessing import Processors
@@ -94,7 +96,9 @@ def _build_model(self) -> None:
         # Use the forward method of the model directly
         self.forward = self.model.forward
 
-    def predict_step(self, batch: torch.Tensor) -> torch.Tensor:
+    def predict_step(
+        self, batch: torch.Tensor, model_comm_group: Optional[ProcessGroup] = None, **kwargs
+    ) -> torch.Tensor:
         """Prediction step for the model.
 
         Parameters
@@ -118,6 +122,6 @@ def predict_step(self, batch: torch.Tensor) -> torch.Tensor:
             # batch, timesteps, horizonal space, variables
             x = batch[:, 0 : self.multi_step, None, ...]  # add dummy ensemble dimension as 3rd index
 
-            y_hat = self(x)
+            y_hat = self(x, model_comm_group)
 
         return self.post_processors(y_hat, in_place=False)