Saving weights and logs of epoch 1

.gitattributes

@@ -16,3 +16,4 @@
 *.pth filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 nohup.out filter=lfs diff=lfs merge=lfs -text
+flax_model.msgpack filter=lfs diff=lfs merge=lfs -text

events.out.tfevents.1626286603.t1v-n-b95d739e-w-0.590614.3.v2 → events.out.tfevents.1626318482.t1v-n-b95d739e-w-0.622701.3.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:69d0b0c11510581415e3ad84919fcc5857dd72e276dfa98d90a601a31995e9d7
-size 4897718
+oid sha256:c0c8833b4f4649f58ab0d01d47c772f8c05080f371d5f9d57e7134d997e944a1
+size 157187

flax_model.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91922c046bb159618da797c8c5076e9684aa45c6ded263ff8c60dab3cb008059
+size 498796983

flax_to_torch.py

@@ -0,0 +1,7 @@
+from transformers import RobertaForMaskedLM, AutoTokenizer
+
+model = RobertaForMaskedLM.from_pretrained("./", from_flax=True)
+model.save_pretrained("./")
+
+tokenizer = AutoTokenizer.from_pretrained("./")
+tokenizer.save_pretrained("./")

nohup.out

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f2645f6739234c77a54e6320ab13a6dcdd86e2decd09fe90e309506975ad0b0b
-size 4470375
+oid sha256:99c46650710b372548e97ab2d4a123983e2b495c3ceb094847f500b4ac3a64f7
+size 193918

run.sh

@@ -1,5 +1,6 @@
 #!/usr/bin/env bash
 python3 run_mlm_flax.py \
+    --model_name_or_path="flax_model.msgpack" \
     --output_dir="./" \
     --model_type="roberta" \
     --config_name="./" \

run_mlm_flax.py

@@ -56,6 +56,24 @@ from transformers import (
 )
 
 
+# Cache the result
+has_tensorboard = is_tensorboard_available()
+if has_tensorboard:
+    try:
+        from flax.metrics.tensorboard import SummaryWriter
+    except ImportError as ie:
+        has_tensorboard = False
+        print(
+            f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
+        )
+
+else:
+    print(
+        "Unable to display metrics through TensorBoard because the package is not installed: "
+        "Please run pip install tensorboard to enable."
+    )
+
+
 MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_MASKED_LM_MAPPING.keys())
 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
@@ -156,7 +174,7 @@ class DataTrainingArguments:
         metadata={"help": "Overwrite the cached training and evaluation sets"},
     )
     validation_split_percentage: Optional[int] = field(
-        default=10,
+        default=5,
         metadata={
             "help": "The percentage of the train set used as validation set in case there's no validation split"
         },
@@ -314,7 +332,7 @@ def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndar
     return batch_idx
 
 
-def write_train_metric(summary_writer, train_metrics, train_time, step):
+def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
     summary_writer.scalar("train_time", train_time, step)
 
     train_metrics = get_metrics(train_metrics)
@@ -323,8 +341,6 @@ def write_train_metric(summary_writer, train_metrics, train_time, step):
         for i, val in enumerate(vals):
             summary_writer.scalar(tag, val, step - len(vals) + i + 1)
 
-
-def write_eval_metric(summary_writer, eval_metrics, step):
     for metric_name, value in eval_metrics.items():
         summary_writer.scalar(f"eval_{metric_name}", value, step)
 
@@ -366,6 +382,10 @@ if __name__ == "__main__":
 
     # Log on each process the small summary:
     logger = logging.getLogger(__name__)
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
 
     # Set the verbosity to info of the Transformers logger (on main process only):
     logger.info(f"Training/evaluation parameters {training_args}")
@@ -557,22 +577,8 @@ if __name__ == "__main__":
         )
 
     # Enable tensorboard only on the master node
-    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))
-        except ImportError as ie:
-            has_tensorboard = False
-            logger.warning(
-                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
-            )
-    else:
-        logger.warning(
-            "Unable to display metrics through TensorBoard because the package is not installed: "
-            "Please run pip install tensorboard to enable."
-        )
+        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
 
     # Data collator
     # This one will take care of randomly masking the tokens.
@@ -584,17 +590,9 @@ if __name__ == "__main__":
     rng = jax.random.PRNGKey(training_args.seed)
     dropout_rngs = jax.random.split(rng, jax.local_device_count())
 
-    if model_args.model_name_or_path:
-        model = FlaxAutoModelForMaskedLM.from_pretrained(
-            model_args.model_name_or_path,
-            config=config,
-            seed=training_args.seed,
-            dtype=getattr(jnp, model_args.dtype),
-        )
-    else:
-        model = FlaxAutoModelForMaskedLM.from_config(
-            config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
-        )
+    model = FlaxAutoModelForMaskedLM.from_config(
+        config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
+    )
 
     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
@@ -636,23 +634,18 @@ if __name__ == "__main__":
         return traverse_util.unflatten_dict(flat_mask)
 
     # create adam optimizer
-    if training_args.adafactor:
-        # We use the default parameters here to initialize adafactor,
-        # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
-        optimizer = optax.adafactor(learning_rate=linear_decay_lr_schedule_fn,)
-    else:
-        optimizer = 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=training_args.weight_decay,
-            mask=decay_mask_fn,
-        )
+    adamw = optax.adamw(
+        learning_rate=linear_decay_lr_schedule_fn,
+        b1=training_args.adam_beta1,
+        b2=training_args.adam_beta2,
+        eps=1e-8,
+        weight_decay=training_args.weight_decay,
+        mask=decay_mask_fn,
+    )
 
     # Setup train state
     state = train_state.TrainState.create(
-        apply_fn=model.__call__, params=model.params, tx=optimizer
+        apply_fn=model.__call__, params=model.params, tx=adamw
     )
 
     # Define gradient update step fn
@@ -742,7 +735,7 @@ if __name__ == "__main__":
         train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)
 
         # Gather the indexes for creating the batch and do a training step
-        for step, batch_idx in enumerate(
+        for i, batch_idx in enumerate(
             tqdm(train_batch_idx, desc="Training...", position=1)
         ):
             samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
@@ -755,69 +748,52 @@ if __name__ == "__main__":
             )
             train_metrics.append(train_metric)
 
-            cur_step = epoch * (num_train_samples // train_batch_size) + step
-
-            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
-                # Save metrics
-                train_metric = jax_utils.unreplicate(train_metric)
-                train_time += time.time() - train_start
-                if has_tensorboard and jax.process_index() == 0:
-                    write_train_metric(
-                        summary_writer, train_metrics, train_time, cur_step
-                    )
-
-                epochs.write(
-                    f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
-                )
-
-                train_metrics = []
-
-            if cur_step % training_args.eval_steps == 0 and cur_step > 0:
-                # ======================== Evaluating ==============================
-                num_eval_samples = len(tokenized_datasets["test"])
-                eval_samples_idx = jnp.arange(num_eval_samples)
-                eval_batch_idx = generate_batch_splits(
-                    eval_samples_idx, eval_batch_size
-                )
-
-                eval_metrics = []
-                for i, batch_idx in enumerate(
-                    tqdm(eval_batch_idx, desc="Evaluating ...", position=2)
-                ):
-                    samples = [
-                        tokenized_datasets["test"][int(idx)] for idx in batch_idx
-                    ]
-                    model_inputs = data_collator(samples, pad_to_multiple_of=16)
-
-                    # Model forward
-                    model_inputs = shard(model_inputs.data)
-                    metrics = p_eval_step(state.params, model_inputs)
-                    eval_metrics.append(metrics)
-
-                # normalize eval metrics
-                eval_metrics = get_metrics(eval_metrics)
-                eval_metrics = jax.tree_map(jnp.sum, eval_metrics)
-                eval_normalizer = eval_metrics.pop("normalizer")
-                eval_metrics = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics)
-
-                # Update progress bar
-                epochs.desc = f"Step... ({cur_step} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"
-
-                # Save metrics
-                if has_tensorboard and jax.process_index() == 0:
-                    cur_step = epoch * (
-                        len(tokenized_datasets["train"]) // train_batch_size
-                    )
-                    write_eval_metric(summary_writer, eval_metrics, cur_step)
-
-            if cur_step % training_args.save_steps == 0 and cur_step > 0:
-                # 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 {cur_step}",
-                    )
+        train_time += time.time() - train_start
+
+        epochs.write(
+            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
+        )
+
+        # ======================== Evaluating ==============================
+        num_eval_samples = len(tokenized_datasets["test"])
+        eval_samples_idx = jnp.arange(num_eval_samples)
+        eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
+
+        eval_metrics = []
+        for i, batch_idx in enumerate(
+            tqdm(eval_batch_idx, desc="Evaluating ...", position=2)
+        ):
+            samples = [tokenized_datasets["test"][int(idx)] for idx in batch_idx]
+            model_inputs = data_collator(samples, pad_to_multiple_of=16)
+
+            # Model forward
+            model_inputs = shard(model_inputs.data)
+            metrics = p_eval_step(state.params, model_inputs)
+            eval_metrics.append(metrics)
+
+        # normalize eval metrics
+        eval_metrics = get_metrics(eval_metrics)
+        eval_metrics = jax.tree_map(jnp.sum, eval_metrics)
+        eval_normalizer = eval_metrics.pop("normalizer")
+        eval_metrics = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics)
+
+        # Update progress bar
+        epochs.desc = f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"
+
+        # Save metrics
+        if has_tensorboard and jax.process_index() == 0:
+            cur_step = epoch * (len(tokenized_datasets["train"]) // train_batch_size)
+            write_metric(
+                summary_writer, train_metrics, eval_metrics, train_time, cur_step
+            )
+
+        # 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 epoch {epoch+1}",
+            )