New logo

run_mlm_flax_stream.py

@@ -25,6 +25,7 @@ import json
 import os
 import shutil
 import sys
+import tempfile
 import time
 from collections import defaultdict
 from dataclasses import dataclass, field
@@ -60,6 +61,8 @@ from transformers import (
     TrainingArguments,
     is_tensorboard_available,
     set_seed,
+    FlaxRobertaForMaskedLM,
+    RobertaForMaskedLM,
 )
 
 
@@ -376,6 +379,27 @@ def rotate_checkpoints(path, max_checkpoints=5):
                 os.remove(path_to_delete)
 
 
+def to_f32(t):
+    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
+
+
+def convert(output_dir, destination_dir="./"):
+    shutil.copyfile(Path(output_dir) / "flax_model.msgpack", destination_dir)
+    shutil.copyfile(Path(output_dir) / "config.json", destination_dir)
+    # Saving extra files from config.json and tokenizer.json files
+    tokenizer = AutoTokenizer.from_pretrained(destination_dir)
+    tokenizer.save_pretrained(destination_dir)
+
+    # Temporary saving bfloat16 Flax model into float32
+    tmp = tempfile.mkdtemp()
+    flax_model = FlaxRobertaForMaskedLM.from_pretrained(destination_dir)
+    flax_model.params = to_f32(flax_model.params)
+    flax_model.save_pretrained(tmp)
+    # Converting float32 Flax to PyTorch
+    model = RobertaForMaskedLM.from_pretrained(tmp, from_flax=True)
+    model.save_pretrained(destination_dir, save_config=False)
+
+
 if __name__ == "__main__":
     # See all possible arguments in src/transformers/training_args.py
     # or by passing the --help flag to this script.
@@ -749,7 +773,8 @@ if __name__ == "__main__":
             eval_metrics = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics)
 
             # Update progress bar
-            steps.desc = f"Step... ({step + 1}/{num_train_steps} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"
+            steps.desc = f"Step... ({step}/{num_train_steps} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"
+            last_desc = steps.desc
 
             if has_tensorboard and jax.process_index() == 0:
                 write_eval_metric(summary_writer, eval_metrics, step)
@@ -762,8 +787,7 @@ if __name__ == "__main__":
             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}",
+                push_to_hub=False,
             )
             save_checkpoint_files(state, data_collator, training_args, training_args.output_dir)
             checkpoints_dir = Path(training_args.output_dir) / "checkpoints" / f"checkpoint-{step}"
@@ -774,6 +798,34 @@ if __name__ == "__main__":
                 Path(training_args.output_dir) / "checkpoints",
                 max_checkpoints=training_args.save_total_limit
             )
+            convert(training_args.output_dir, "./")
+            model.save_pretrained(
+                training_args.output_dir,
+                params=params,
+                push_to_hub=training_args.push_to_hub,
+                commit_message=last_desc,
+            )
 
         # update tqdm bar
         steps.update(1)
+
+    if jax.process_index() == 0:
+        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,
+            params=params,
+            push_to_hub=False,
+        )
+        save_checkpoint_files(state, data_collator, training_args, training_args.output_dir)
+        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)
+        save_checkpoint_files(state, data_collator, training_args, checkpoints_dir)
+        convert(training_args.output_dir, "./")
+        model.save_pretrained(
+            training_args.output_dir,
+            params=params,
+            push_to_hub=training_args.push_to_hub,
+            commit_message=last_desc,
+        )