Fixing restore checkpoint step

mc4/mc4.py

@@ -1,11 +1,11 @@
-"""mC4 dataset based on Common Crawl."""
+"""Perplexity Sampled mC4 dataset based on Common Crawl."""
 
 
 import gzip
 import json
 
 import datasets
-import kenlm
+import kenlm  # pip install https://github.com/kpu/kenlm/archive/master.zip
 import numpy as np
 from numpy.random import default_rng
 
@@ -289,6 +289,7 @@ class Mc4(datasets.GeneratorBasedBuilder):
         self.sampling_factor = kwargs.pop("sampling_factor", None)
         self.boundaries = kwargs.pop("boundaries", None)
         self.seed = kwargs.pop("seed", None)
+        self.kwargs = kwargs
         if self.sampling_method:
             if self.seed is not None:
                 self.rng = default_rng(self.seed)
@@ -316,7 +317,7 @@ class Mc4(datasets.GeneratorBasedBuilder):
             doc_length += length
         return 10.0 ** (-doc_log_score / doc_length)
 
-    def _should_keep_doc_step(self, doc, factor=1.5e5, boundaries=None):
+    def _should_keep_doc_step(self, doc, factor=1.5e5, boundaries=None, **kwargs):
         perplexity = self.get_perplexity(doc)
         if boundaries is None:
             boundaries = [536394.99320948, 662247.50212365, 919250.87225178]
@@ -331,17 +332,18 @@ class Mc4(datasets.GeneratorBasedBuilder):
         probability = factor / quartile_range
         return self.rng.uniform() < probability
 
-    def _should_keep_doc_gaussian(self, doc, factor=0.78, boundaries=None):
+    def _should_keep_doc_gaussian(self, doc, factor=0.78, boundaries=None, **kwargs):
+        width = kwargs.get("width", 9 / 2)  # width (spread) of the exponential curve
         perplexity = self.get_perplexity(doc)
         if boundaries is not None:
             m = boundaries[1]
         else:
             m = 662247.50212365
-        exponential = np.exp(-9/2 * ((perplexity - m) / m) ** 2)
+        exponential = np.exp((-1 / width) * ((perplexity - m) / m) ** 2)
         weighted_perplexity = factor * exponential
         return self.rng.uniform() < weighted_perplexity
 
-    def _should_keep_doc_random(self, doc, factor=None, boundaries=None):
+    def _should_keep_doc_random(self, doc, factor=None, boundaries=None, **kwargs):
         if factor is None:
             factor = 0.5
         return self.rng.uniform() <= factor
@@ -415,7 +417,8 @@ class Mc4(datasets.GeneratorBasedBuilder):
                                 if self.should_keep_doc(
                                     example["text"],
                                     factor=self.sampling_factor,
-                                    boundaries=self.boundaries):
+                                    boundaries=self.boundaries
+                                    **self.kwargs):
                                     yield id_, example
                                     id_ += 1
                     else:

run_mlm_flax_stream.py

@@ -348,6 +348,24 @@ def save_checkpoint_files(state, data_collator, training_args, save_dir):
         json.dump({"step": unreplicated_state.step.item()}, f)
 
 
+def restore_checkpoint(save_dir, state):
+    logger.info(f"Restoring checkpoint from {save_dir}")
+    with open(os.path.join(save_dir, "flax_model.msgpack"), "rb") as f:
+        params = from_bytes(state.params, f.read())
+
+    with open(os.path.join(save_dir, "optimizer_state.msgpack"), "rb") as f:
+        opt_state = from_bytes(state.opt_state, f.read())
+
+    args = joblib.load(os.path.join(save_dir, "training_args.joblib"))
+    data_collator = joblib.load(os.path.join(save_dir, "data_collator.joblib"))
+
+    with open(os.path.join(save_dir, "training_state.json"), "r") as f:
+        training_state = json.load(f)
+    step = training_state["step"]
+
+    return params, opt_state, step, args, data_collator
+
+
 def rotate_checkpoints(path, max_checkpoints=5):
     paths = sorted(Path(path).iterdir(), key=os.path.getmtime)[::-1]
     if len(paths) > max_checkpoints:
@@ -484,8 +502,6 @@ if __name__ == "__main__":
     has_tensorboard = is_tensorboard_available()
     if has_tensorboard and jax.process_index() == 0:
         try:
-            from flax.metrics.tensorboard import SummaryWriter
-            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
             # Enable Weight&Biases
             import wandb
             wandb.init(
@@ -496,6 +512,8 @@ if __name__ == "__main__":
             wandb.config.update(training_args)
             wandb.config.update(model_args)
             wandb.config.update(data_args)
+            from flax.metrics.tensorboard import SummaryWriter
+            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
         except ImportError as ie:
             has_tensorboard = False
             logger.warning(
@@ -569,6 +587,42 @@ if __name__ == "__main__":
 
     # Setup train state
     state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw)
+    saved_step = 0
+    if "checkpoint" in model_args.model_name_or_path:
+        params, opt_state, saved_step, args, data_collator = restore_checkpoint(model_args.model_name_or_path, state)
+        # Create learning rate schedule
+        warmup_fn = optax.linear_schedule(
+            init_value=0.0, end_value=args.learning_rate, transition_steps=args.warmup_steps
+        )
+        decay_fn = optax.linear_schedule(
+            init_value=args.learning_rate,
+            end_value=0,
+            transition_steps=data_args.num_train_steps - args.warmup_steps,
+        )
+        linear_decay_lr_schedule_fn = optax.join_schedules(
+            schedules=[warmup_fn, decay_fn], boundaries=[args.warmup_steps]
+        )
+        # create adam optimizer
+        adamw = optax.adamw(
+            learning_rate=linear_decay_lr_schedule_fn,
+            b1=training_args.adam_beta1,
+            b2=training_args.adam_beta2,
+            eps=training_args.adam_epsilon,
+            weight_decay=args.weight_decay,
+            mask=decay_mask_fn,
+        )
+        state = train_state.TrainState(
+            step=saved_step,
+            apply_fn=model.__call__,
+            params=params,
+            tx=adamw,
+            opt_state=opt_state,
+        )
+        # self.args = args
+        # data_collator = data_collator
+        # scheduler_fn = args.learning_rate
+        model.params = params
+
 
     # Define gradient update step fn
     def train_step(state, batch, dropout_rng):
@@ -637,7 +691,10 @@ if __name__ == "__main__":
     eval_samples = advance_iter_and_group_samples(training_iter, data_args.num_eval_samples, max_seq_length)
 
     steps = tqdm(range(num_train_steps), desc="Training...", position=0)
-    for step in range(num_train_steps):
+    for step in range(saved_step, num_train_steps):
+        if step < saved_step:
+            steps.update(1)
+            continue
         # ======================== Training ================================
         try:
             samples = advance_iter_and_group_samples(training_iter, train_batch_size, max_seq_length)
@@ -700,7 +757,7 @@ if __name__ == "__main__":
 
         # save checkpoint after eval_steps
         if step % training_args.save_steps == 0 and step > 0 and jax.process_index() == 0:
-            logger.info(f"Saving checkpoint at {step + 1} steps")
+            logger.info(f"Saving checkpoint at {step} steps")
             params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
             model.save_pretrained(
                 training_args.output_dir,
@@ -709,9 +766,9 @@ if __name__ == "__main__":
                 commit_message=f"Saving weights and logs of step {step + 1}",
             )
             save_checkpoint_files(state, data_collator, training_args, training_args.output_dir)
-            checkpoints_dir = Path(training_args.output_dir) / "checkpoints" / f"checkpoint-{step + 1}"
+            checkpoints_dir = Path(training_args.output_dir) / "checkpoints" / f"checkpoint-{step}"
             checkpoints_dir.mkdir(parents=True, exist_ok=True)
-            model.save_pretrained(checkpoints_dir, params=params,)
+            model.save_pretrained(checkpoints_dir, params=params)
             save_checkpoint_files(state, data_collator, training_args, checkpoints_dir)
             rotate_checkpoints(
                 Path(training_args.output_dir) / "checkpoints",