Preparing code for final runs

config.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 from transformers import RobertaConfig
 config = RobertaConfig.from_pretrained("roberta-large")
-config.save_pretrained("./")
+config.save_pretrained("./configs/large")
 
 config = RobertaConfig.from_pretrained("roberta-base")
-config.save_pretrained("./config-base.json")
+config.save_pretrained("./configs/base")

run_mlm_flax_stream.py

@@ -21,13 +21,16 @@ Here is the full list of checkpoints on the hub that can be fine-tuned by this s
 https://huggingface.co/models?filter=masked-lm
 """
 import logging
+import json
 import os
+import shutil
 import sys
 import time
 from collections import defaultdict
 from dataclasses import dataclass, field
 
 # You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.
+import joblib
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -39,9 +42,10 @@ from tqdm import tqdm
 import flax
 import jax
 import jax.numpy as jnp
-import kenlm
+import kenlm  # pip install https://github.com/kpu/kenlm/archive/master.zip
 import optax
 from flax import jax_utils, traverse_util
+from flax.serialization import from_bytes, to_bytes
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard
 from transformers import (
@@ -334,6 +338,26 @@ def write_eval_metric(summary_writer, eval_metrics, step):
         summary_writer.scalar(f"eval_{metric_name}", value, step)
 
 
+def save_checkpoint_files(state, data_collator, training_args, save_dir):
+    unreplicated_state = jax_utils.unreplicate(state)
+    with open(os.path.join(save_dir, "optimizer_state.msgpack"), "wb") as f:
+        f.write(to_bytes(unreplicated_state.opt_state))
+    joblib.dump(training_args, os.path.join(save_dir, "training_args.joblib"))
+    joblib.dump(data_collator, os.path.join(save_dir, "data_collator.joblib"))
+    with open(os.path.join(save_dir, "training_state.json"), "w") as f:
+        json.dump({"step": unreplicated_state.step.item()}, f)
+
+
+def rotate_checkpoints(path, max_checkpoints=5):
+    paths = sorted(Path(path).iterdir(), key=os.path.getmtime)[::-1]
+    if len(paths) > max_checkpoints:
+        for path_to_delete in paths[max_checkpoints:]:
+            try:
+                shutil.rmtree(path_to_delete)
+            except OSError:
+                os.remove(path_to_delete)
+
+
 if __name__ == "__main__":
     # See all possible arguments in src/transformers/training_args.py
     # or by passing the --help flag to this script.
@@ -391,19 +415,31 @@ if __name__ == "__main__":
             filepaths["train"] = data_args.train_file
         if data_args.validation_file:
             filepaths["validation"] = data_args.validation_file
-        dataset = load_dataset(
-            data_args.dataset_name,
-            data_args.dataset_config_name,
-            cache_dir=model_args.cache_dir,
-            streaming=True,
-            split="train",
-            sampling_method=sampling_args.sampling_method,
-            sampling_factor=sampling_args.sampling_factor,
-            boundaries=sampling_args.boundaries,
-            perplexity_model=sampling_args.perplexity_model,
-            seed=training_args.seed,
-            data_files=filepaths,
-        )
+        try:
+            dataset = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                cache_dir=model_args.cache_dir,
+                streaming=True,
+                split="train",
+                sampling_method=sampling_args.sampling_method,
+                sampling_factor=sampling_args.sampling_factor,
+                boundaries=sampling_args.boundaries,
+                perplexity_model=sampling_args.perplexity_model,
+                seed=training_args.seed,
+                data_files=filepaths,
+            )
+        except Exception as exc:
+            logger.warning(
+                f"Unable to load local dataset with perplexity sampling support. Using huggingface.co/datasets/{data_args.dataset_name}: {exc}"
+            )
+            dataset = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                cache_dir=model_args.cache_dir,
+                streaming=True,
+                split="train",
+            )
 
     if model_args.config_name:
         config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
@@ -662,15 +698,25 @@ if __name__ == "__main__":
                 write_eval_metric(summary_writer, eval_metrics, step)
             eval_metrics = []
 
-            # save checkpoint after each epoch and push checkpoint to the hub
-            if jax.process_index() == 0:
-                params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
-                model.save_pretrained(
-                    training_args.output_dir,
-                    params=params,
-                    push_to_hub=training_args.push_to_hub,
-                    commit_message=f"Saving weights and logs of step {step+1}",
-                )
+        # save checkpoint after eval_steps
+        if step % training_args.save_steps == 0 and step > 0 and jax.process_index() == 0:
+            print(f"Saving checkpoint at {step + 1} steps")
+            params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
+            model.save_pretrained(
+                training_args.output_dir,
+                params=params,
+                push_to_hub=training_args.push_to_hub,
+                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.mkdir(parents=True, exist_ok=True)
+            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",
+                max_checkpoints=training_args.save_total_limit
+            )
 
         # update tqdm bar
         steps.update(1)

run_stream.sh

@@ -4,9 +4,10 @@ python ./run_mlm_flax_stream.py \
     --output_dir="./outputs" \
     --model_type="roberta" \
     --config_name="./configs/base" \
-    --tokenizer_name="./" \
+    --tokenizer_name="./configs/base" \
     --dataset_name="./mc4" \
     --dataset_config_name="es" \
+    --train_file="path/to/mc4-es-train-50M-XXX.jsonl" \
     --max_seq_length="128" \
     --pad_to_max_length  \
     --per_device_train_batch_size="256" \
@@ -16,13 +17,11 @@ python ./run_mlm_flax_stream.py \
     --adam_epsilon="1e-6" \
     --learning_rate="6e-4" \
     --weight_decay="0.01" \
-    --save_strategy="steps" \
-    --save_steps="1000" \
+    --save_steps="10000" \
     --save_total_limit="5" \
     --warmup_steps="24000" \
     --overwrite_output_dir \
-    --num_train_steps="500000" \
-    --eval_steps="1000" \
+    --num_train_steps="250000" \
+    --eval_steps="10000" \
     --dtype="bfloat16" \
-    --sampling_method="steps" \
     --logging_steps="500" 2>&1 | tee run_stream.log