Minor refactoring

megatron/arguments.py

@@ -244,7 +244,8 @@ def _add_training_args(parser):
     group.add_argument('--global-batch-size', type=int, default=None,
                        help='Training batch size. If this value is None, then '
                        'use micro-batch-size * data-parallel-size as the '
-                       'global batch size')
+                       'global batch size. This choice will result in 1 for '
+                       'number of micro-batches.')
     group.add_argument('--rampup-batch-size', nargs='*', default=None,
                        help='Batch size ramp up with the following values:'
                        '  --rampup-batch-size <start batch size> '

megatron/global_vars.py

@@ -34,20 +34,12 @@ _GLOBAL_ADLR_AUTORESUME = None
 _GLOBAL_TIMERS = None
 
 
-
 def get_args():
     """Return arguments."""
     _ensure_var_is_initialized(_GLOBAL_ARGS, 'args')
     return _GLOBAL_ARGS
 
 
-def get_num_microbatches_calculator():
-    """Return num-microbatches calculator."""
-    _ensure_var_is_initialized(_GLOBAL_NUM_MICROBATCHES_CALCULATOR,
-                               'number of micro-batches calculator.')
-    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR
-
-
 def get_num_microbatches():
     return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get()
 
@@ -141,6 +133,7 @@ class NumMicroBatchesCalculator(ABC):
     def get(self):
         pass
 
+    @abstractmethod
     def update(self, consumed_samples):
         pass
 

megatron/training.py

@@ -223,7 +223,7 @@ def setup_model_and_optimizer(model_provider_func):
     else:
         args.iteration = 0
 
-    # Wrap model for distributed training."""
+    # We only support local DDP with multiple micro-batches.
     if get_num_microbatches() > 1:
         assert args.DDP_impl == 'local'