Fix wrong initialization of lr scheduler

README.md

@@ -34,7 +34,7 @@ Nanotron is a library for pretraining transformer models. It provides a simple a
 ## Installation
 
 ```bash
-# Requirements: Python>=3.10
+# Requirements: Python>=3.10,<3.12
 git clone https://github.com/huggingface/nanotron
 cd nanotron
 pip install --upgrade pip

pyproject.toml

@@ -16,7 +16,7 @@ classifiers = [
 dependencies = [
     "torch>=1.13.1",
     "pyyaml",
-    "numpy",
+    "numpy<2",
     "packaging",
     "safetensors",
     "dacite",
@@ -45,7 +45,7 @@ test = [
 ]
 
 fast-modeling = [
-    "flash-attn>=2.5.0",
+    "flash-attn>=2.5.0,<2.7.0",
 ]
 
 nanosets = [

src/nanotron/config/config.py

@@ -159,6 +159,8 @@ class CheckpointsArgs:
     save_initial_state: Optional[bool] = False
     save_final_state: Optional[bool] = False
     resume_checkpoint_path: Optional[xPath] = None
+    load_lr_scheduler: Optional[bool] = True
+    load_optimizer: Optional[bool] = True
     checkpoints_path_is_shared_file_system: Optional[bool] = False
 
     def __post_init__(self):

src/nanotron/optim/zero.py

@@ -6,7 +6,12 @@
 
 import numpy as np
 import torch.optim
-from functorch.dim import tree_map
+
+try:
+    from functorch.dim import tree_map
+except:
+    from torch.utils._pytree import tree_map
+
 from torch import nn
 from tqdm import tqdm
 

src/nanotron/serialize/main.py

@@ -60,6 +60,7 @@ def save(
     should_save_lr_scheduler: bool = True,
     sanity_checks: bool = True,
 ) -> None:
+
     assert isinstance(training_metadata, TrainingMetadata)
 
     try:
@@ -107,6 +108,7 @@ def save(
                 lr_scheduler=lr_scheduler,
                 parallel_context=parallel_context,
                 root_folder=root_folder,
+                is_zero=config.optimizer.zero_stage
             )
     except Exception as e:
         log_rank(

src/nanotron/trainer.py

@@ -190,6 +190,13 @@ def __init__(
             optimizer_args=self.config.optimizer,
             parallel_context=self.parallel_context,
         )
+
+        # Init learning rate scheduler
+        self.lr_scheduler = lr_scheduler_builder(
+            optimizer=self.optimizer,
+            lr_scheduler_args=self.config.optimizer.learning_rate_scheduler,
+            total_training_steps=self.config.tokens.train_steps,
+        )
         if self.init_checkpoint_path is not None:
             load_optimizer(
                 optimizer=self.optimizer,
@@ -199,13 +206,6 @@ def __init__(
                 model=self.unwrapped_model,
                 map_location="cpu",
             )
-
-        # Init learning rate scheduler
-        self.lr_scheduler = lr_scheduler_builder(
-            optimizer=self.optimizer,
-            lr_scheduler_args=self.config.optimizer.learning_rate_scheduler,
-            total_training_steps=self.config.tokens.train_steps,
-        )
         if self.init_checkpoint_path is not None:
             load_lr_scheduler(
                 lr_scheduler=self.lr_scheduler,
@@ -215,7 +215,7 @@ def __init__(
             )
 
         # Define iteration start state
-        if self.init_checkpoint_path is not None:
+        if self.init_checkpoint_path is not None and self.config.checkpoints.load_lr_scheduler:
             checkpoint_metadata = load_meta(
                 parallel_context=self.parallel_context, root_folder=self.init_checkpoint_path
             )
@@ -553,7 +553,11 @@ def training_step(
             handle = None
 
         # Move optimizer states back to GPU before optimizer step
-        if self.init_checkpoint_path is not None and self.iteration_step == self.initial_iter_step:
+        if (
+            self.init_checkpoint_path is not None
+            and self.config.checkpoints.load_optimizer
+            and self.iteration_step == self.initial_iter_step
+        ):
             state_dict_to_device(self.optimizer.state_dict(), "cuda")
 
         before_optim_step_sanity_checks(