Training GenCast

[gbruno16]

Training Gencast

Description

This PR is currently a draft and implements several modifications:

1. Registering Graph Tensors as Buffers:

   • The graph's tensors are registered as buffers during initialization. This allows them to move along with the model in a multi-device setting.
   • Since persistent=False, these tensors will not be part of the state_dict.

2. Wrapping the Model with PyTorch Lightning:

   • The model is wrapped using PyTorch Lightning to facilitate easier training with various strategies (e.g., DDP, fp16).
   • This includes support for different schedulers (e.g., cosine annealing), loggers (e.g., wandb) and callbacks.

3. Normalize first in sparse transformers:

   • The flag norm_first in sparse transformers move the conditional layer normalization before the multi-head attention. This seems to improve training performance.

4. Noise generation with torch_harmonics

   • torch_harmonics supports noise generation on gpu with torch backend.

[jacobbieker]

Really good work here! This looks great.