Merge pull request #90 from borisdayma/feat-new

dev/seq2seq/run_seq2seq_flax.py

@@ -100,12 +100,6 @@ class ModelArguments:
             "help": "Pretrained config name or path if not the same as model_name"
         },
     )
-    tokenizer_name: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": "Pretrained tokenizer name or path if not the same as model_name"
-        },
-    )
     cache_dir: Optional[str] = field(
         default=None,
         metadata={
@@ -422,7 +416,7 @@ def wandb_log(metrics, step=None, prefix=None):
             f"{prefix}/{k}" if prefix is not None else k: v for k, v in metrics.items()
         }
         if step is not None:
-            log_metrics["train/step"] = unreplicate(step)
+            log_metrics["train/step"] = step
         wandb.log(log_metrics)
 
 
@@ -534,11 +528,6 @@ def main():
             )
 
     else:
-        base_model = FlaxAutoModelForSeq2SeqLM.from_pretrained(
-            model_args.model_name_or_path,
-            seed=training_args.seed,
-            dtype=getattr(jnp, model_args.dtype),
-        )
         # Set up our new model config
         config = BartConfig.from_pretrained(model_args.model_name_or_path)
         config.tie_word_embeddings = False
@@ -563,11 +552,6 @@ def main():
             config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
         )
 
-        # Use pre-trained weights for encoder
-        model.params["model"]["encoder"] = base_model.params["model"]["encoder"]
-        model.params["model"]["shared"] = base_model.params["model"]["shared"]
-        del base_model
-
     # Load tokenizer if it has not been set
     if tokenizer is None:
         tokenizer = AutoTokenizer.from_pretrained(
@@ -862,7 +846,7 @@ def main():
         f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}"
     )
     logger.info(
-        f"  Total train batch size (w. parallel & distributed) = {batch_size_per_update}"
+        f"  Total train batch size (w. parallel, distributed & gradient accumulation) = {batch_size_per_update}"
     )
     logger.info(f"  Total global steps = {total_steps}")
     logger.info(f"  Total optimization steps = {total_optimization_steps}")
@@ -870,7 +854,7 @@ def main():
     epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
 
     # set default x-axis as 'train/step'
-    wandb_log({}, step=state.step)
+    wandb_log({}, step=unreplicate(state.step))
     wandb.define_metric("*", step_metric="train/step")
 
     # add interesting config parameters
@@ -909,7 +893,7 @@ def main():
             eval_metrics = jax.tree_map(jnp.mean, eval_metrics)
 
             # log metrics
-            wandb_log(eval_metrics, step=state.step, prefix="eval")
+            wandb_log(eval_metrics, step=unreplicate(state.step), prefix="eval")
 
             # Print metrics and update progress bar
             desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']})"
@@ -943,6 +927,10 @@ def main():
 
             # save to W&B
             if data_args.log_model:
+                # save some space
+                c = wandb.wandb_sdk.wandb_artifacts.get_artifacts_cache()
+                c.cleanup(wandb.util.from_human_size("5GB"))
+
                 metadata = {"step": step, "epoch": epoch}
                 if eval_metrics is not None:
                     metadata["eval/loss"] = eval_metrics["loss"]
@@ -970,11 +958,8 @@ def main():
                 artifact.add_file(
                     str(Path(training_args.output_dir) / "training_state.json")
                 )
-                wandb.run.log_artifact(artifact)
 
-                # save some space
-                c = wandb.wandb_sdk.wandb_artifacts.get_artifacts_cache()
-                c.cleanup(wandb.util.from_human_size("5GB"))
+                wandb.run.log_artifact(artifact)
 
             # save to the hub
             if training_args.push_to_hub:
@@ -988,7 +973,8 @@ def main():
 
     for epoch in epochs:
         # ======================== Training ================================
-        wandb_log({"train/epoch": epoch}, step=state.step)
+        step = unreplicate(state.step)
+        wandb_log({"train/epoch": epoch}, step=step)
 
         # Create sampling rng
         rng, input_rng = jax.random.split(rng)
@@ -1010,19 +996,20 @@ def main():
             total=steps_per_epoch,
         ):
             state, train_metric = p_train_step(state, batch)
+            step = unreplicate(state.step)
 
-            if state.step % data_args.log_interval == 0 and jax.process_index() == 0:
+            if step % data_args.log_interval == 0 and jax.process_index() == 0:
                 # log metrics
-                wandb_log(unreplicate(train_metric), step=state.step, prefix="train")
+                wandb_log(unreplicate(train_metric), step=step, prefix="train")
 
-            if training_args.eval_steps and state.step % training_args.eval_steps == 0:
+            if training_args.eval_steps and step % training_args.eval_steps == 0:
                 run_evaluation()
 
-            if state.step % data_args.save_model_steps == 0:
-                run_save_model(state, state.step, epoch)
+            if step % data_args.save_model_steps == 0:
+                run_save_model(state, step, epoch)
 
         # log final train metrics
-        wandb_log(unreplicate(train_metric), step=state.step, prefix="train")
+        wandb_log(unreplicate(train_metric), step=step, prefix="train")
 
         train_metric = unreplicate(train_metric)
         epochs.write(