Sophiex/dev/ema fix inference (ecmwf#1094)

* Save current state

* Save current state

* Barebone FSDP2 prototype TODO save checkpoints

* First version of saving model

* Fix save_model

* Log everything and log to files

* Remove redundant path creation

* Allow for both slurm and torchrun + fewer log files

* Cleaning up init_ddp

* Ruff

* Attempt to avoid duplicate logging

* FSDP2 with mixed precision policy

* Ruff

* Clean up and logging

* Try to get loggers to behave as we want

* Makes ruff unhappy but works

* Fixed ruff issue

* Fixed problems with multi-node training.

* Fix for interactive/non-DDP runs

* No idea why, but this seems to work so far

Committing simply so it is saved, obviously needs cleanup

* Still works! So which is it memory or the grad scaler?

* Also still works, I now strongly suspect the amp.gradscaler

* This still works, I have no clue anymore why but whatever it works
now....

* Enable loading model from absolute paths

* Enable loading for 1 GPU only

* Fix 1 GPU train continue

* Appease ruff

* Fix saving the model more regularly and perf logging

* Fixed problem when training with 2 nodes.

* Fix data loader seed

* Appease ruff

* Shouldn't overwrite with_fsdp like this

* Potential fix for FSDP2 issue

with different ranks using different model parts

* Fix loss scaling and logging of dummy data loss

* Clean up

* Appease ruff

* Start implementing EMA, works for 1 GPU

* Make EMA model multi-gpu compatible

* Fix linting issues

* Address comments on PR

* Enforce the ema model to strictly be the same as model

Note this likely needs to be changed for student-teacher training

* Rename variable

* Fix inference

* One more thing

* Fixed linting

* Fix problem with checkpoint saving when validate_with_ema is false or not specified

---------

Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>

Author: sophie-xhonneux

src/weathergen/train/trainer.py

@@ -104,6 +104,7 @@ def inference(self, cf, devices, run_id_trained, epoch):
         cf = self.cf
         self.device_type = torch.accelerator.current_accelerator()
         self.device = torch.device(f"{self.device_type}:{cf.local_rank}")
+        self.ema_model = None
 
         # !! modifies config: adds config.streams[i].<stage>_source_channels
         # and config.streams[i].<stage>_target_channels !!
@@ -298,6 +299,7 @@ def run(self, cf, devices, run_id_contd=None, epoch_contd=None):
             self.model = torch.compile(self.model, dynamic=True)
 
         self.validate_with_ema = cf.get("validate_with_ema", False)
+        self.ema_model = None
         if self.validate_with_ema:
             meta_ema_model = self.init_model_and_shard(cf, devices)[0]
             self.ema_model = EMAModel(
@@ -642,7 +644,12 @@ def validate(self, epoch):
                         dtype=self.mixed_precision_dtype,
                         enabled=cf.with_mixed_precision,
                     ):
-                        preds, _ = self.ema_model.forward_eval(
+                        model_forward = (
+                            self.model.forward
+                            if self.ema_model is None
+                            else self.ema_model.forward_eval
+                        )
+                        preds, _ = model_forward(
                             self.model_params, batch, cf.forecast_offset, forecast_steps
                         )