add pre-emption callback to esm2 train (#433)

Adds the callback to trigger checkpoint saving to the ESM2 pretraining
script. Now we just need to figure out how to add the `sbatch
--signal=TERM@300` to our virtualpipe script

sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py

@@ -190,6 +190,7 @@ def main(
         PerplexityLoggingCallback(log_train=False, log_val=True),
         RichModelSummary(max_depth=4),
         LearningRateMonitor(),
+        nl_callbacks.PreemptionCallback(),
     ]
     if nsys_profiling:
         if nsys_end_step is None: