Adding checkpoint after traning ends

eliebak · May 7, 2024 · b1872e1 · b1872e1
1 parent 6629dfa
commit b1872e1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/nanotron/config/config.py b/src/nanotron/config/config.py
@@ -129,6 +129,7 @@ class CheckpointsArgs:
     checkpoints_path: Path
     checkpoint_interval: int
     save_initial_state: Optional[bool] = False
+    save_final_state: Optional[bool] = False
     resume_checkpoint_path: Optional[Path] = None
     checkpoints_path_is_shared_file_system: Optional[bool] = False
 

diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
@@ -442,7 +442,10 @@ def train(
                     self.save_checkpoint()
 
         dist.barrier()  # let's wait for everyone before leaving
-
+
+        if self.config.checkpoints.save_final_state:
+            self.save_checkpoint()
+
         self.post_training()
 
     def training_step(