Merging megatron with ICT

megatron/utils.py

@@ -19,12 +19,41 @@ import sys
 
 import torch
 
+from apex.multi_tensor_apply import multi_tensor_applier
+import amp_C
+
 from megatron import get_args
 from megatron import print_rank_0
 from megatron import get_adlr_autoresume
 from megatron import mpu
 from megatron.checkpointing import save_checkpoint
-from megatron.fp16 import FP16_Optimizer
+from megatron.model.module import param_is_not_shared
+from megatron.mpu.layers import param_is_not_tensor_parallel_duplicate
+
+
+def calc_params_l2_norm(model):
+    """Calculate l2 norm of parameters """
+    # Remove duplicate params.
+    params_data = []
+    for param in model.parameters():
+        is_not_shared = param_is_not_shared(param)
+        is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param)
+        if is_not_shared and is_not_tp_duplicate:
+            params_data.append(param.data)
+    # Calculate norm
+    dummy_overflow_buf = torch.cuda.IntTensor([0])
+    norm, _ = multi_tensor_applier(
+        amp_C.multi_tensor_l2norm,
+        dummy_overflow_buf,
+        [params_data],
+        False # no per-parameter norm
+    )
+    norm_2 = norm * norm
+    # Sum across all model-parallel GPUs.
+    torch.distributed.all_reduce(norm_2,
+                                 op=torch.distributed.ReduceOp.SUM,
+                                 group=mpu.get_model_parallel_group())
+    return norm_2.item() ** 0.5
 
 
 def average_losses_across_data_parallel_group(losses):
@@ -47,11 +76,13 @@ def report_memory(name):
         torch.cuda.memory_allocated() / mega_bytes)
     string += ' | max allocated: {}'.format(
         torch.cuda.max_memory_allocated() / mega_bytes)
-    string += ' | reserved: {}'.format(torch.cuda.memory_reserved() / mega_bytes)
+    string += ' | reserved: {}'.format(
+        torch.cuda.memory_reserved() / mega_bytes)
     string += ' | max reserved: {}'.format(
         torch.cuda.max_memory_reserved() / mega_bytes)
     if mpu.get_data_parallel_rank() == 0:
-        print("[Rank {}] {}".format(torch.distributed.get_rank(), string), flush=True)
+        print("[Rank {}] {}".format(torch.distributed.get_rank(), string),
+              flush=True)
 
 
 def print_params_min_max_norm(optimizer, iteration):
@@ -59,9 +90,7 @@ def print_params_min_max_norm(optimizer, iteration):
     index = 0
     rank = torch.distributed.get_rank()
     string = 'iteration, rank, index, tensor-model-parallel, min, max, norm\n'
-    optimizer_ = optimizer
-    if isinstance(optimizer, FP16_Optimizer):
-        optimizer_ = optimizer.optimizer
+    optimizer_ = optimizer.optimizer
     for param_group in optimizer_.param_groups:
         for param in param_group['params']:
             index += 1

pretrain_ict.py

@@ -149,6 +149,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
         short_seq_prob=args.short_seq_prob,
         seed=args.seed,
         skip_warmup=(not args.mmap_warmup),
+        binary_head=False,
         dataset_type='ict')
     print_rank_0("> finished creating BERT ICT datasets ...")