Add support for saving intermediate models

Author: balvisio

bionemo-recipes/recipes/codonfm_native_te/hydra_config/defaults.yaml

@@ -60,6 +60,7 @@ lr_scheduler_kwargs:
 checkpoint:
   ckpt_dir: ???
   save_final_model: true
+  save_final_model_with_checkpoint: false
   resume_from_checkpoint: true
   save_every_n_steps: 1_000
   max_checkpoints: 5

bionemo-recipes/recipes/codonfm_native_te/train_ddp.py

@@ -241,6 +241,13 @@ def main(args: DictConfig) -> float | None:
                             dist_config=dist_config,
                             max_checkpoints=args.checkpoint.max_checkpoints,
                         )
+                        if args.checkpoint.save_final_model_with_checkpoint:
+                            save_final_model_ddp(
+                                model=model,
+                                config=config,
+                                save_directory=ckpt_path / f"step_{step}" / "final_model",
+                                dist_config=dist_config,
+                            )
 
                     if val_dataloader is not None and step > 0 and step % args.validation.eval_interval == 0:
                         model.eval()

bionemo-recipes/recipes/codonfm_native_te/train_fsdp2.py

@@ -273,6 +273,13 @@ def main(args: DictConfig) -> float | None:
                             dist_config=dist_config,
                             max_checkpoints=args.checkpoint.max_checkpoints,
                         )
+                        if args.checkpoint.save_final_model_with_checkpoint:
+                            save_final_model_fsdp2(
+                                model=model,
+                                config=config,
+                                save_directory=ckpt_path / f"step_{step}" / "final_model",
+                                dist_config=dist_config,
+                            )
 
                     if val_dataloader is not None and step > 0 and step % args.validation.eval_interval == 0:
                         model.eval()