Perf statistics for save_checkpoint (#64)

* save_checkpoint perf monitoring

* Disable checkpoint save on exit

Author: tjruwase

Megatron-LM-v1.1.5-ZeRO3/megatron/training.py

@@ -38,7 +38,7 @@
 from megatron.model.realm_model import ICTBertModel
 from megatron.utils import check_adlr_autoresume_termination
 from megatron.utils import make_data_loader
-from megatron.utils import report_memory, flops_calculator
+from megatron.utils import report_memory, flops_calculator, throughput_calculator, checkpoint_throughput_calculator
 
 import deepspeed
 from deepspeed.runtime.utils import see_memory_usage
@@ -106,8 +106,8 @@ def pretrain(train_valid_test_dataset_provider, model_provider,
                                    valid_data_iterator, model,
                                    iteration, False)
 
-    if args.save and iteration != 0:
-        save_checkpoint(iteration, model, optimizer, lr_scheduler)
+#    if args.save and iteration != 0:
+#        save_checkpoint(iteration, model, optimizer, lr_scheduler)
 
     if args.do_test:
         # Run on test data.
@@ -175,8 +175,8 @@ def get_optimizer(model):
                                        weight_decay=args.weight_decay)
     else:
         # Use torch Adam instead of Fused Adam from NVIDIA which seems to have some issue.
-        #optimizer = Adam(param_groups,
-        optimizer = torch.optim.AdamW(param_groups,
+        optimizer = Adam(param_groups,
+        #optimizer = torch.optim.AdamW(param_groups,
                          lr=args.lr,
                          weight_decay=args.weight_decay,
                          betas=(args.adam_beta1, args.adam_beta2),
@@ -384,6 +384,7 @@ def add_to_logging(name):
     add_to_logging('backward-clip-grad')
     add_to_logging('optimizer')
     add_to_logging('batch generator')
+    add_to_logging('save checkpoint')
 
     # Tensorboard values.
     if writer and torch.distributed.get_rank() == 0:
@@ -423,12 +424,14 @@ def add_to_logging(name):
             total_loss_dict[got_nan_key])
         total_loss_dict[skipped_iters_key] = 0
         total_loss_dict[got_nan_key] = 0
+        timers.log(timers_to_log, normalizer=args.log_interval)
         print_rank_0(log_string)
         if report_memory_flag:
             report_memory('after {} iterations'.format(iteration))
             report_memory_flag = False
-        timers.log(timers_to_log, normalizer=args.log_interval)
+
         flops_calculator(model, args, elapsed_time)
+        throughput_calculator(model, args, elapsed_time)      
 
     return report_memory_flag
 
@@ -462,11 +465,6 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
         loss_scale = None
         if args.fp16:
             loss_scale = optimizer.cur_scale if args.deepspeed else optimizer.loss_scale
-        report_memory_flag = training_log(loss_dict, total_loss_dict,
-                                          optimizer.param_groups[0]['lr'],
-                                          iteration, loss_scale,
-                                          report_memory_flag, skipped_iter,
-                                          model=model)
 
         # Autoresume
         if args.adlr_autoresume and \
@@ -475,9 +473,21 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
                                               lr_scheduler)
 
         # Checkpointing
-        if args.save and args.save_interval and \
-           iteration % args.save_interval == 0:
+        should_save_checkpoint = args.save and args.save_interval and \
+           iteration % args.save_interval == 0           
+        timers('save checkpoint').start()
+        if should_save_checkpoint:
             save_checkpoint(iteration, model, optimizer, lr_scheduler)
+        timers('save checkpoint').stop()
+
+        if should_save_checkpoint:
+            checkpoint_throughput_calculator(model, args, timers('save checkpoint').elapsed(reset=False))
+
+        report_memory_flag = training_log(loss_dict, total_loss_dict,
+                                          optimizer.param_groups[0]['lr'],
+                                          iteration, loss_scale,
+                                          report_memory_flag, skipped_iter,
+                                          model=model)
 
         # Evaluation
         # XXX temporarily disabled for ZeRO-3

Megatron-LM-v1.1.5-ZeRO3/megatron/utils.py

@@ -194,3 +194,22 @@ def flops_calculator(model, args, iteration_time):
     effective_tera_flops_per_gpu = giga_flops_per_model_per_train_step / (iteration_time * 1000.0 * gpus_per_model)
 
     print_rank_0(f"Effective Tera Flops per GPU: {round(effective_tera_flops_per_gpu, 2)} and total parameters {round(approx_parameters_in_billions, 3)} B")
+
+
+def throughput_calculator(model, args, iteration_time):
+    gpus_per_model = torch.distributed.get_world_size(group = mpu.get_model_parallel_group())
+    samples_per_model = args.batch_size * args.seq_length
+    model_replica_count = torch.distributed.get_world_size() / gpus_per_model
+    approx_parameters_in_billions = get_parameters_in_billions(model)
+    samples_per_second = samples_per_model * model_replica_count / (iteration_time * 1000.0)
+
+    print_rank_0(f'Samples per second: {round(samples_per_second, 2)} and total parameters {round(approx_parameters_in_billions, 3)} B')
+
+
+def checkpoint_throughput_calculator(model, args, latency_sec):
+    approx_parameters_in_billions = get_parameters_in_billions(model)
+    checkpoint_multiplier = 12 # fp16 weights (2), fp32 weights (4), fp32 momentum (4), fp32 variance (4)
+    checkpoint_giga_bytes = approx_parameters_in_billions * checkpoint_multiplier
+    giga_bytes_per_second = checkpoint_giga_bytes / latency_sec
+    
+    print_rank_0(f'Checkpoint Save GB: {round(checkpoint_giga_bytes, 3)}, GB_PerSec: {round(giga_bytes_per_second, 2)}, Latency(secs): {round(latency_sec, 3)}')