Introduce configuration flags for gradient checkpointing

[javak87]

Description

Replaces #1168 . I opened a fresh PR because resolving conflicts in the original became messy. The change set is the same intent; this PR is the clean continuation.

Issue Number

Closes #1141

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

Performance gain

Regarding the performance gain: when we disable all gradient checkpointing on a single node using the default config, we see about a 20% improvement with no CUDA out-of-memory (OOM) issues, compared to the develop branch (commit df59a9a). Here’s how I ran the code on both the develop branch and this branch:

../WeatherGenerator-private/hpc/launch-slurm.py --time 180 --nodes=1
The experiments have been uploaded:
Develop branch test ID: cka04xbw
javad/dev/cond_checkpoint_all-1141-ver1 branch test ID: s8p7iec0

If a different configuration runs into CUDA OOM, you can set some of the flags (not all of them) with the checkpoint_enabled option to True. This enables gradient checkpointing so activations are recomputed during the backward pass, which helps avoid CUDA OOM errors.