Move args writer to the beginning of training

megatron/initialize.py

@@ -79,8 +79,6 @@ def initialize_megatron(extra_args_provider=None, args_defaults={},
         # Autoresume.
         _init_autoresume()
         
-        # Write arguments to tensorboard.
-        _write_args_to_tensorboard()
         # No continuation function
         return None
         
@@ -154,13 +152,14 @@ def _set_random_seed(seed_):
         raise ValueError('Seed ({}) should be a positive integer.'.format(seed))
 
 
-def _write_args_to_tensorboard():
+def write_args_to_tensorboard():
     """Write arguments to tensorboard."""
     args = get_args()
     writer = get_tensorboard_writer()
     if writer:
         for arg in vars(args):
-            writer.add_text(arg, str(getattr(args, arg)))
+            writer.add_text(arg, str(getattr(args, arg)),
+                            global_step=args.iteration)
 
 
 def _initialize_mem_buffs():

megatron/training.py

@@ -41,6 +41,7 @@ from megatron.checkpointing import save_checkpoint
 from megatron.fp16 import FP16_Module
 from megatron.fp16 import FP16_Optimizer
 from megatron.initialize import initialize_megatron
+from megatron.initialize import write_args_to_tensorboard
 from megatron.learning_rates import AnnealingLR
 from megatron.model import DistributedDataParallel as LocalDDP
 from megatron.model import get_params_for_weight_decay_optimization
@@ -811,6 +812,9 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
     args = get_args()
     timers = get_timers()
 
+    # Write args to tensorboard
+    write_args_to_tensorboard()
+
     # Turn on training mode which enables dropout.
     model.train()