Add ckpt_path config to resume training (#67)

Minor update to the training script, allowing for easy resumption of training. Co-authored-by: Juncheng Liu <[email protected]>
SalesforceAIResearch · Jun 25, 2024 · fac7a93 · fac7a93
1 parent 92ddd74
commit fac7a93
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 0 deletions.
diff --git a/cli/conf/finetune/default.yaml b/cli/conf/finetune/default.yaml
@@ -10,6 +10,7 @@ run_name: ???
 seed: 0
 tf32: true
 compile: false  # set to mode: default, reduce-overhead, max-autotune
+ckpt_path: null
 trainer:
   _target_: lightning.Trainer
   accelerator: auto

diff --git a/cli/conf/pretrain/default.yaml b/cli/conf/pretrain/default.yaml
@@ -10,6 +10,7 @@ run_name: ???
 seed: 0
 tf32: true
 compile: false  # set to mode: default, reduce-overhead, max-autotune
+ckpt_path: null  # set to "last" to resume training
 trainer:
   _target_: lightning.Trainer
   accelerator: auto

diff --git a/cli/train.py b/cli/train.py
@@ -142,6 +142,7 @@ def main(cfg: DictConfig):
     trainer.fit(
         model,
         datamodule=DataModule(cfg, train_dataset, val_dataset),
+        ckpt_path=cfg.ckpt_path,
     )
-Original file line number
+Diff line change
@@ Expand Up / @@ -142,6 +142,7 @@ def main(cfg: DictConfig): @@
         trainer.fit(
             model,
             datamodule=DataModule(cfg, train_dataset, val_dataset),
+            ckpt_path=cfg.ckpt_path,
         )
@@ Expand Down @@