wip

huggingface · Nov 2, 2024 · 2e8f1b0 · 2e8f1b0
1 parent 51ca40b
commit 2e8f1b0
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/nanotron/trainer.py b/src/nanotron/trainer.py
@@ -87,6 +87,7 @@
 from nanotron.serialize import (
     load_lr_scheduler,
     load_meta,
+    load_random_states,
     load_weights,
     parse_ckpt_path,
     save,
@@ -170,6 +171,11 @@ def __init__(
         self.random_states = init_random_states(
             parallel_config=self.config.parallelism, tp_pg=self.parallel_context.tp_pg
         )
+        if self.init_checkpoint_path is not None:
+            self.random_states = load_random_states(
+                parallel_context=self.parallel_context,
+                root_folder=self.init_checkpoint_path,
+            )
         self.model = self.init_model()  # Defines self.model
         self.unwrapped_model: NanotronModel = (
             self.model.module if isinstance(self.model, DistributedDataParallel) else self.model