Manual revert training.py to main verion

Author: tjruwase

megatron/training.py

@@ -15,42 +15,45 @@
 
 """Pretrain utilities."""
 
-import deepspeed
-from megatron.utils import report_memory, flops_calculator
-from megatron.schedules import forward_backward_pipelining_with_interleaving
-from megatron.schedules import forward_backward_pipelining_without_interleaving
-from megatron.schedules import forward_backward_no_pipelining
-from megatron.utils import calc_params_l2_norm
-from megatron.data.data_samplers import build_pretraining_data_loader
-from megatron.utils import unwrap_model
-from megatron.utils import check_adlr_autoresume_termination
-from megatron.model import DistributedDataParallel as LocalDDP
-from megatron.learning_rates import AnnealingLR
-from megatron.initialize import write_args_to_tensorboard
-from megatron.initialize import initialize_megatron
-from megatron.optimizer import get_megatron_optimizer
-from megatron.model import Float16Module
-from megatron.checkpointing import save_checkpoint
-from megatron.checkpointing import load_checkpoint
-from megatron import print_rank_last
-from megatron import print_rank_0
-from megatron import mpu
-from megatron import update_num_microbatches
-from megatron import is_last_rank
-from megatron import get_num_microbatches
-from megatron import get_current_global_batch_size
-from megatron import get_tensorboard_writer
-from megatron import get_timers
-from megatron import get_args
-from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
-import torch
 from datetime import datetime
 import math
 import sys
 import time
 # The earliest we can measure the start time.
 _TRAIN_START_TIME = time.time()
 
+import torch
+from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
+
+from megatron import get_args
+from megatron import get_timers
+from megatron import get_tensorboard_writer
+from megatron import get_current_global_batch_size
+from megatron import get_num_microbatches
+from megatron import is_last_rank
+from megatron import update_num_microbatches
+from megatron import mpu
+from megatron import print_rank_0
+from megatron import print_rank_last
+from megatron.checkpointing import load_checkpoint
+from megatron.checkpointing import save_checkpoint
+from megatron.model import Float16Module
+from megatron.optimizer import get_megatron_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.utils import check_adlr_autoresume_termination
+from megatron.utils import unwrap_model
+from megatron.data.data_samplers import build_pretraining_data_loader
+from megatron.utils import calc_params_l2_norm
+from megatron.schedules import forward_backward_no_pipelining
+from megatron.schedules import forward_backward_pipelining_without_interleaving
+from megatron.schedules import forward_backward_pipelining_with_interleaving
+from megatron.utils import report_memory, flops_calculator
+
+import deepspeed
+
 
 def print_datetime(string):
     """Note that this call will sync across all ranks."""
@@ -159,7 +162,6 @@ def pretrain(train_valid_test_dataset_provider,
                                    test_data_iterator, model,
                                    0, True)
 
-
 def update_train_iters(args):
 
     # For iteration-based training, we don't need to do anything
@@ -184,7 +186,7 @@ def update_train_iters(args):
         # Constant phase
         # Note that we throw away any partial last batch.
         iterations += (args.train_samples - consumed_samples) // \
-            args.global_batch_size
+                      args.global_batch_size
         args.train_iters = iterations
 
     print_rank_0('setting training iterations to {}'.format(args.train_iters))
@@ -216,6 +218,7 @@ def get_model(model_provider_func):
             post_process=post_process
         )
 
+
     if not isinstance(model, list):
         model = [model]
 
@@ -231,10 +234,10 @@ def get_model(model_provider_func):
     if mpu.get_data_parallel_rank() == 0:
         print(' > number of parameters on (tensor, pipeline) '
               'model parallel rank ({}, {}): {}'.format(
-                  mpu.get_tensor_model_parallel_rank(),
-                  mpu.get_pipeline_model_parallel_rank(),
-                  sum([sum([p.ds_numel if hasattr(p, 'ds_id') else p.nelement() for p in model_module.parameters()])
-                       for model_module in model])), flush=True)
+            mpu.get_tensor_model_parallel_rank(),
+            mpu.get_pipeline_model_parallel_rank(),
+            sum([sum([p.ds_numel if hasattr(p,'ds_id') else p.nelement() for p in model_module.parameters()])
+                 for model_module in model])), flush=True)
 
     if args.deepspeed:
         return model
@@ -358,7 +361,7 @@ def setup_model_and_optimizer(model_provider_func):
 
     # get model without FP16 and/or TorchDDP wrappers
     if args.iteration == 0 and len(unwrapped_model) == 1 \
-            and hasattr(unwrapped_model[0], 'init_state_dict_from_bert'):
+        and hasattr(unwrapped_model[0], 'init_state_dict_from_bert'):
         print_rank_0("Initializing ICT from pretrained BERT model")
         unwrapped_model[0].init_state_dict_from_bert()
         if args.fp16:
@@ -379,7 +382,7 @@ def train_step(forward_step_func, data_iterator,
         skipped_iter = 0
         grad_norm = 0.
         num_zeros_in_grad = 0
-        return {'lm loss': loss}, skipped_iter, grad_norm, num_zeros_in_grad
+        return {'lm loss' : loss}, skipped_iter, grad_norm, num_zeros_in_grad
 
     # Set grad to zero.
     if not args.deepspeed:
@@ -439,8 +442,8 @@ def train_step(forward_step_func, data_iterator,
     timers('optimizer').start()
     if args.deepspeed:
         increment = get_num_microbatches() * \
-            args.micro_batch_size * \
-            args.data_parallel_size
+                    args.micro_batch_size * \
+                    args.data_parallel_size
         model[0].step(lr_kwargs={'increment': increment})
         update_successful = model[0].was_step_applied()
     else:
@@ -455,8 +458,8 @@ def train_step(forward_step_func, data_iterator,
     else:
         if update_successful:
             increment = get_num_microbatches() * \
-                args.micro_batch_size * \
-                args.data_parallel_size
+                        args.micro_batch_size * \
+                        args.data_parallel_size
             lr_scheduler.step(increment=increment)
             skipped_iter = 0
         else:
@@ -504,8 +507,8 @@ def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
         else:
             value = loss_dict[key].float().sum().item()
             is_nan = value == float('inf') or \
-                value == -float('inf') or \
-                value != value
+                     value == -float('inf') or \
+                     value != value
             got_nan = got_nan or is_nan
     total_loss_dict[nan_iters_key] = total_loss_dict.get(
         nan_iters_key, 0) + int(got_nan)
@@ -539,10 +542,10 @@ def add_to_logging(name):
         get_num_microbatches()
 
     total_iterations = total_loss_dict[advanced_iters_key] + \
-        total_loss_dict[skipped_iters_key]
+                       total_loss_dict[skipped_iters_key]
 
     # Tensorboard values.
-    if writer and (iteration % args.tensorboard_log_interval == 0) and \
+    if writer and (iteration % args.tensorboard_log_interval == 0 ) and \
        is_last_rank():
         if args.log_learning_rate_to_tensorboard:
             writer.add_scalar('learning-rate', learning_rate, iteration)
@@ -553,7 +556,7 @@ def add_to_logging(name):
             writer.add_scalar('batch-size vs samples', batch_size,
                               args.consumed_train_samples)
         for key in loss_dict:
-            writer.add_scalar(key, loss_dict[key], iteration)
+            writer.add_scalar(key , loss_dict[key], iteration)
             writer.add_scalar(key + ' vs samples', loss_dict[key],
                               args.consumed_train_samples)
         if args.log_loss_scale_to_tensorboard:
@@ -595,7 +598,7 @@ def add_to_logging(name):
             if key not in [advanced_iters_key, skipped_iters_key,
                            nan_iters_key]:
                 avg = total_loss_dict[key].item() / \
-                    float(max(1, total_loss_dict[advanced_iters_key]))
+                      float(max(1, total_loss_dict[advanced_iters_key]))
                 if avg > 0.0:
                     log_string += ' {}: {:.6E} |'.format(key, avg)
                 total_loss_dict[key] = torch.cuda.FloatTensor([0.0])
@@ -663,10 +666,11 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
         if args.deepspeed:
             # inform deepspeed of any batch size changes
             global_batch_size = mpu.get_data_parallel_world_size() * \
-                args.micro_batch_size * \
-                get_num_microbatches()
+                                args.micro_batch_size * \
+                                get_num_microbatches()
             model[0].set_train_batch_size(global_batch_size)
 
+
         loss_dict, skipped_iter, grad_norm, num_zeros_in_grad = \
             train_step(forward_step_func,
                        train_data_iterator,
@@ -675,8 +679,8 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
                        lr_scheduler)
         iteration += 1
         args.consumed_train_samples += mpu.get_data_parallel_world_size() * \
-            args.micro_batch_size * \
-            get_num_microbatches()
+                                       args.micro_batch_size * \
+                                       get_num_microbatches()
 
         # Logging.
         if args.deepspeed:
@@ -739,6 +743,7 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
             print_datetime('exiting program at iteration {}'.format(iteration))
             sys.exit()
 
+
     return iteration
 
 
@@ -767,17 +772,17 @@ def evaluate(forward_step_func, data_iterator, model, verbose=False):
                     forward_backward_func = forward_backward_pipelining_without_interleaving
             else:
                 forward_backward_func = forward_backward_no_pipelining
-
+            
             if args.deepspeed:
                 # DeepSpeed uses eval_batch() and already aggregates losses.
                 assert isinstance(model, list) and len(model) == 1
                 loss = model[0].eval_batch(data_iterator)
-                loss_dicts = [{'lm loss': loss}] * get_num_microbatches()
+                loss_dicts = [{'lm loss' : loss}] * get_num_microbatches()
             else:
                 loss_dicts = forward_backward_func(
                     forward_step_func, data_iterator, model, optimizer=None,
                     timers=None, forward_only=True)
-
+            
             if mpu.is_pipeline_last_stage(ignore_virtual=True):
                 # Reduce across processes.
                 for loss_dict in loss_dicts:
@@ -786,8 +791,8 @@ def evaluate(forward_step_func, data_iterator, model, verbose=False):
                             key, torch.cuda.FloatTensor([0.0])) + loss_dict[key]
 
             args.consumed_valid_samples += mpu.get_data_parallel_world_size() \
-                * args.micro_batch_size \
-                * get_num_microbatches()
+                                           * args.micro_batch_size \
+                                           * get_num_microbatches()
     # Move model back to the train mode.
     for model_module in model:
         model_module.train()
@@ -797,7 +802,6 @@ def evaluate(forward_step_func, data_iterator, model, verbose=False):
 
     return total_loss_dict
 
-
 def evaluate_and_print_results(prefix, forward_step_func,
                                data_iterator, model,
                                iteration, verbose=False):
@@ -835,7 +839,6 @@ def cyclic_iter(iter):
         for x in iter:
             yield x
 
-
 def build_train_valid_test_data_iterators(
         build_train_valid_test_datasets_provider):
     """XXX"""
@@ -865,7 +868,7 @@ def build_train_valid_test_data_iterators(
         else:
             train_samples = args.train_iters * args.global_batch_size
         eval_iters = (args.train_iters // args.eval_interval + 1) * \
-            args.eval_iters
+                     args.eval_iters
         test_iters = args.eval_iters
         train_val_test_num_samples = [train_samples,
                                       eval_iters * args.global_batch_size,
@@ -904,25 +907,26 @@ def build_train_valid_test_data_iterators(
     args.do_valid = flags[1].item()
     args.do_test = flags[2].item()
 
+
     # Build iterators.
     dl_type = args.dataloader_type
     assert dl_type in ['single', 'cyclic']
 
     if train_dataloader is not None:
         train_data_iterator = iter(train_dataloader) if dl_type == 'single' \
-            else iter(cyclic_iter(train_dataloader))
+                              else iter(cyclic_iter(train_dataloader))
     else:
         train_data_iterator = None
 
     if valid_dataloader is not None:
         valid_data_iterator = iter(valid_dataloader) if dl_type == 'single' \
-            else iter(cyclic_iter(valid_dataloader))
+                              else iter(cyclic_iter(valid_dataloader))
     else:
         valid_data_iterator = None
 
     if test_dataloader is not None:
         test_data_iterator = iter(test_dataloader) if dl_type == 'single' \
-            else iter(cyclic_iter(test_dataloader))
+                             else iter(cyclic_iter(test_dataloader))
     else:
         test_data_iterator = None