Merge pull request huggingface#165 from angegonzalez/feature/save-tra…

…in-ending-checkpoint Adding checkpoint after traning ends
eliebak · Aug 2, 2024 · 2793c92 · 2793c92
2 parents 6ad5994 + b1872e1
commit 2793c92
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
@@ -145,6 +145,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(