Merge branch 'tensorboard_queue_size_increase' into 'main'

added option to change tensorboard queue size

See merge request ADLR/megatron-lm!223

megatron/arguments.py

@@ -285,6 +285,10 @@ def _add_logging_args(parser):
                        help='If set, calculate and log parameters norm.')
     group.add_argument('--tensorboard-log-interval', type=int, default=1,
                        help='Report to tensorboard interval.')
+    group.add_argument('--tensorboard-queue-size', type=int, default=1000,
+                       help='Size of the tensorboard queue for pending events '
+                       'and summaries before one of the ‘add’ calls forces a '
+                       'flush to disk.')
     group.add_argument('--log-timers-to-tensorboard', action='store_true',
                        help='If set, write timers to tensorboard.')
     group.add_argument('--log-batch-size-to-tensorboard', action='store_true',

megatron/global_vars.py

@@ -137,7 +137,8 @@ def _set_tensorboard_writer(args):
             from torch.utils.tensorboard import SummaryWriter
             print('> setting tensorboard ...')
             _GLOBAL_TENSORBOARD_WRITER = SummaryWriter(
-                log_dir=args.tensorboard_dir)
+                log_dir=args.tensorboard_dir,
+                max_queue=args.tensorboard_queue_size)
         except ModuleNotFoundError:
             print('WARNING: TensorBoard writing requested but is not '
                   'available (are you using PyTorch 1.1.0 or later?), '