Update optimizers

EasyLM/data.py

@@ -153,6 +153,7 @@ class HuggingfaceDataset(object):
         config.start_seek_loc = 0
         config.tokens_count_at_start = 0
         config.batch_token_dtype = 'i4'
+        config.reset_dataset_loc = False
 
         if updates is not None:
             config.update(ConfigDict(updates).copy_and_resolve_references())
@@ -173,6 +174,8 @@ class HuggingfaceDataset(object):
         self._dataset_loc = self.config.start_seek_loc
         self._total_tokens = self.config.tokens_count_at_start
         self._index = 0
+        self.reset_dataset_loc = self.config.reset_dataset_loc
+
 
     def __iter__(self):
         if not self._eval_dataset and self._train_epochs > 0:
@@ -236,6 +239,10 @@ class HuggingfaceDataset(object):
         self._dataset_loc = state_dict.get('dataset_loc', self.config.start_seek_loc)
         self._total_tokens = state_dict.get('total_tokens', self.config.tokens_count_at_start)
         self._train_epochs = state_dict.get('epochs', 0)
+        if self.reset_dataset_loc:
+            self._dataset_loc = 0
+            self._train_epochs = 0
+
 
     @property
     def seq_length(self):

EasyLM/optimizers.py

@@ -205,8 +205,9 @@ class LionOptimizerFactory(object):
         config.init_lr = 0.0
         config.end_lr = 0.0001
         config.lr = 0.001
-        config.lr_warmup_steps = 2000
-        config.lr_decay_steps = 500000
+        config.lr_warmup_steps = 60000
+        config.lr_constant_steps = 840000
+        config.lr_decay_steps = 100000
         config.b1 = 0.9
         config.b2 = 0.98
         config.clip_gradient = 1.0
@@ -243,6 +244,43 @@ class LionOptimizerFactory(object):
                 ],
                 [config.lr_warmup_steps],
             )
+        elif config.lr_schedule_type == "warmup_constant_linear_decay":
+            learning_rate_schedule = optax.join_schedules(
+                [
+                    optax.linear_schedule(
+                        init_value=config.init_lr,
+                        end_value=config.lr,
+                        transition_steps=config.lr_warmup_steps,
+                    ),
+                    optax.constant_schedule(config.lr),
+                    optax.linear_schedule(
+                        init_value=config.lr,
+                        end_value=config.end_lr,
+                        transition_steps=config.lr_decay_steps,
+                    )
+                ],
+                [config.lr_warmup_steps, config.lr_constant_steps],
+            )
+        elif config.lr_schedule_type == "warmup_constant_exponential_decay":
+            learning_rate_schedule = optax.join_schedules(
+                [
+                    optax.linear_schedule(
+                        init_value=config.init_lr,
+                        end_value=config.lr,
+                        transition_steps=config.lr_warmup_steps,
+                    ),
+                    optax.constant_schedule(config.lr),
+                    optax.exponential_decay(
+                        init_value=config.lr, 
+                        transition_steps=config.lr_decay_steps, 
+                        decay_rate=config.lr_decay_rate, 
+                        transition_begin=0, 
+                        staircase=False, 
+                        end_value=config.end_lr,
+                    )
+                ],
+                [config.lr_warmup_steps, config.lr_constant_steps],
+            )
         elif config.lr_schedule_type == "exponential_decay":
             learning_rate_schedule = optax.exponential_decay(
                         init_value=config.lr, 
@@ -252,8 +290,14 @@ class LionOptimizerFactory(object):
                         staircase=False, 
                         end_value=config.end_lr,
             )
+        elif config.lr_schedule_type == "linear_decay":
+            learning_rate_schedule = optax.linear_schedule(
+                        init_value=config.lr,
+                        end_value=config.end_lr,
+                        transition_steps=config.lr_decay_steps,
+            )
         else:
-            raise ValueError('config.lr_schedule_type must be "warmup_cosine_decay_schedule", "warmup_constant", or "exponential_decay"')
+            raise ValueError('config.lr_schedule_type must be "warmup_cosine_decay_schedule", "warmup_constant", "warmup_constant_linear_decay", "warmup_constant_exponential_decay", "exponential_decay" or "linear_decay"')
 
         optimizer_info = dict(
             learning_rate_schedule=learning_rate_schedule,

pretrain_llama_3b.sh

@@ -23,10 +23,11 @@ python3 -m EasyLM.models.llama.llama_train \
     --tokenizer.vocab_file='tokenizer.model' \
     --optimizer.type='lion' \
     --optimizer.lion_optimizer.weight_decay=1.0 \
-    --optimizer.lion_optimizer.lr_schedule_type='warmup_constant' \
+    --optimizer.lion_optimizer.lr_schedule_type='warmup_constant_linear_decay' \
     --optimizer.lion_optimizer.lr=1e-4 \
     --optimizer.lion_optimizer.end_lr=1e-5 \
     --optimizer.lion_optimizer.lr_warmup_steps=60000 \
+    --optimizer.lion_optimizer.lr_constant_steps=900000 \
     --optimizer.lion_optimizer.lr_decay_steps=100000 \
     --optimizer.lion_optimizer.bf16_momentum=True \
     --train_dataset.type='huggingface' \