nvidia_energy_llm_fix.patch

diff --git a/megatron/training.py b/megatron/training.py
index cebe085b..be272522 100644
--- a/megatron/training.py
+++ b/megatron/training.py
@@ -5,6 +5,7 @@
 from datetime import datetime
 import math
 import sys
+import os
 import time
 # The earliest we can measure the start time.
 _TRAIN_START_TIME = time.time()
@@ -41,6 +42,54 @@ from megatron.core.pipeline_parallel import get_forward_backward_func
 from megatron.utils import report_memory
 from megatron.model.vision.knn_monitor import compute_feature_bank
 
+accelerator_list = ["JEDI", "GH200", "H100", "WAIH100", "A100"]
+def num_floating_point_operations(args, batch_size):
+    # Attention projection size.
+    query_projection_size = args.kv_channels * args.num_attention_heads
+    query_projection_to_hidden_size_ratio = query_projection_size / args.hidden_size
+    # Group Query Attention.
+    if not args.group_query_attention:
+        args.num_query_groups = args.num_attention_heads
+    # MoE.
+    num_experts_routed_to = 1 if args.num_experts is None else args.moe_router_topk
+    gated_linear_multiplier = 3 / 2 if args.swiglu else 1
+
+    # The 12x term below comes from the following factors; for more details, see
+    # "APPENDIX: FLOATING-POINT OPERATIONS" in https://arxiv.org/abs/2104.04473.
+    # - 3x: Each GEMM in the model needs to be performed 3 times (forward pass,
+    #       backward wgrad [weight gradient], backward dgrad [data gradient]).
+    # - 2x: GEMMs of a particular size are stacked twice in the standard Transformer model
+    #       architectures implemented in this codebase (e.g., h->ffn_h GEMM and ffn_h->h GEMM
+    #       in MLP layer).
+    # - 2x: A GEMM of a m*n tensor with a n*k tensor requires 2mnk floating-point operations.
+    expansion_factor = 3 * 2 * 2
+
+    return (
+        expansion_factor
+        * batch_size
+        * args.seq_length
+        * args.num_layers
+        * args.hidden_size
+        * args.hidden_size
+        * (
+            # Attention.
+            (
+                (
+                    1
+                    + (args.num_query_groups / args.num_attention_heads)
+                    + (args.seq_length / args.hidden_size)
+                ) * query_projection_to_hidden_size_ratio
+            )
+            # MLP.
+            + (
+                (args.ffn_hidden_size / args.hidden_size)
+                * num_experts_routed_to
+                * gated_linear_multiplier
+            )
+            # Logit.
+            + (args.padded_vocab_size / (2 * args.num_layers * args.hidden_size))
+        )
+    )
 
 def print_datetime(string):
     """Note that this call will sync across all ranks."""
@@ -155,7 +204,9 @@ def pretrain(train_valid_test_dataset_provider,
                               model, optimizer, opt_param_scheduler,
                               train_data_iterator, valid_data_iterator,
                               process_non_loss_data_func, config)
-
+            if iteration == -1:
+                print_datetime("Exiting code....")
+                return
         print_datetime('after training is done')
 
         if args.save and iteration != 0:
@@ -596,18 +647,78 @@ def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
             )
 
     if iteration % args.log_interval == 0:
+        # benchmark_tokens = args.consumed_train_samples * args.seq_length
+        # if benchmark_tokens <= 100*10**6:
         elapsed_time = timers('interval-time').elapsed(barrier=True)
         elapsed_time_per_iteration = elapsed_time / total_iterations
+
+        # Not checking for variable sequence length to compute throughput
+        seq_len = args.seq_length
+        hidden_size = args.hidden_size
+        num_layers = args.num_layers
+        vocab_size = args.padded_vocab_size
+
+        samples_per_sec = batch_size / elapsed_time_per_iteration
+        samples_per_sec_per_replica = samples_per_sec / args.data_parallel_size
+        tokens_per_sec = samples_per_sec * seq_len
+        tokens_per_sec_per_replica = tokens_per_sec / args.data_parallel_size
+
+        # General TFLOPs formula (borrowed from Equation 3 in Section 5.1 of
+        # https://arxiv.org/pdf/2104.04473.pdf).
+        # The factor of 4 is when used with activation check-pointing,
+        # otherwise it will be 3
+        checkpoint_activations_factor = 4 if args.recompute_granularity else 3
+
+        # GLU activations double the hidden states in the upscaling
+        # feed-forward in each transformer layer
+        # This leads to 16bsh^2 instead of 8bsh^2 per first feed-forward
+        # layer in MLP, thus we increase the coefficient by 8.
+        # Refer to
+        # https://github.com/bigscience-workshop/Megatron-DeepSpeed/pull/283#issue-1260805063
+        # for more details.
+
+        coefficient = 32 if args.swiglu else 24
+
+        flops_per_iteration = (
+            coefficient * checkpoint_activations_factor * batch_size
+            * seq_len * num_layers * (hidden_size ** 2)
+        ) * (
+            1.
+            + (seq_len / (6. * hidden_size))
+            + (vocab_size / (16. * num_layers * hidden_size))
+        )
+        tflops = (
+            flops_per_iteration
+            / (elapsed_time_per_iteration * args.world_size * (10 ** 12))
+        )
+
+        tflops_megatron = num_floating_point_operations(args, batch_size) / (
+            elapsed_time_per_iteration * 10**12 * args.world_size)
+
         if writer:
             if args.log_timers_to_tensorboard:
                 writer.add_scalar('iteration-time',
                                   elapsed_time_per_iteration, iteration)
+                writer.add_scalar('iteration-time vs samples',
+                                  elapsed_time_per_iteration, args.consumed_train_samples)
+                writer.add_scalar('iteration-time vs tokens',
+                                  elapsed_time_per_iteration, args.consumed_train_samples * args.seq_length)
+                writer.add_scalar('samples per second',
+                                  samples_per_sec, iteration)
+                writer.add_scalar('samples per second per replica',
+                                  samples_per_sec_per_replica, iteration)
+                writer.add_scalar('tokens per second',
+                                  tokens_per_sec, iteration)
+                writer.add_scalar('tokens per second per replica',
+                                  tokens_per_sec_per_replica, iteration)
+                writer.add_scalar('TFLOPs per gpu (estimated)',
+                                  tflops, iteration)
         log_string = ' iteration {:8d}/{:8d} |'.format(
             iteration, args.train_iters)
         log_string += ' consumed samples: {:12d} |'.format(
             args.consumed_train_samples)
-        log_string += ' elapsed time per iteration (ms): {:.1f} |'.format(
-            elapsed_time_per_iteration * 1000.0)
+        log_string += ' elapsed time per iteration (s): {:.4f} |'.format(
+            elapsed_time_per_iteration)
         log_string += ' learning rate: {:.3E} |'.format(learning_rate)
         log_string += ' global batch size: {:5d} |'.format(batch_size)
         for key in total_loss_dict:
@@ -629,6 +740,9 @@ def training_log(loss_dict, total_loss_dict, learning_rate, iteration,
             total_loss_dict[skipped_iters_key])
         log_string += ' number of nan iterations: {:3d} |'.format(
             total_loss_dict[nan_iters_key])
+        log_string += ' samples per second: {:.3f} |'.format(samples_per_sec)
+        log_string += ' TFLOPs: {:.2f} |'.format(tflops)
+        log_string += ' TFLOPs(megatron): {:.2f} |'.format(tflops_megatron)
         total_loss_dict[advanced_iters_key] = 0
         total_loss_dict[skipped_iters_key] = 0
         total_loss_dict[nan_iters_key] = 0
@@ -743,7 +857,10 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler,
                 save_checkpoint_and_time(iteration, model, optimizer,
                                          opt_param_scheduler)
                 print_datetime('exiting program after receiving SIGTERM.')
-                sys.exit()
+                if os.getenv("ACCELERATOR") in accelerator_list:
+                    return -1
+                else:
+                    sys.exit()
 
         if args.save and args.save_interval and \
            iteration % args.save_interval == 0:
@@ -764,7 +881,10 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler,
                     save_checkpoint_and_time(iteration, model, optimizer,
                                              opt_param_scheduler)
                 print_datetime('exiting program after {} minutes'.format(train_time))
-                sys.exit()
+                if os.getenv("ACCELERATOR") in accelerator_list:
+                    return -1
+                else:
+                    sys.exit()
 
         # Exiting based on iterations
         if args.exit_interval and iteration % args.exit_interval == 0:
@@ -773,7 +893,10 @@ def train(forward_step_func, model, optimizer, opt_param_scheduler,
                                          opt_param_scheduler)
             torch.distributed.barrier()
             print_datetime('exiting program at iteration {}'.format(iteration))
-            sys.exit()
+            if os.getenv("ACCELERATOR") in accelerator_list:
+                    return -1
+            else:
+                    sys.exit()
 
         if args.profile and \
            iteration == args.profile_step_end and \
diff --git a/pretrain_gpt.py b/pretrain_gpt.py
index 09e0710a..7237b2ca 100644
--- a/pretrain_gpt.py
+++ b/pretrain_gpt.py
@@ -1,7 +1,25 @@
 # Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
 
 """Pretrain GPT"""
+import unicodedata
+import re
+def slugify(value, allow_unicode=False):
+    """
+    Taken from https://github.com/django/django/blob/master/django/utils/text.py
+    Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated
+    dashes to single dashes. Remove characters that aren't alphanumerics,
+    underscores, or hyphens. Convert to lowercase. Also strip leading and
+    trailing whitespace, dashes, and underscores.
+    """
+    value = str(value)
+    if allow_unicode:
+        value = unicodedata.normalize('NFKC', value)
+    else:
+        value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')
+    value = re.sub(r'[^\w\s-]', '', value.lower())
+    return re.sub(r'[-\s]+', '-', value).strip('-_')
 
+import sys
 import os
 import torch
 from functools import partial
@@ -126,9 +144,59 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
 
 
 if __name__ == "__main__":
+    from jpwr.ctxmgr import get_power
 
-    pretrain(train_valid_test_datasets_provider,
-             model_provider,
-             ModelType.encoder_or_decoder,
-             forward_step,
-             args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})
+    methods = set()
+    if not torch.cuda.is_available():
+        print("Not running on GPU")
+    else:
+        for i in range(torch.cuda.device_count()):
+            device_name = torch.cuda.get_device_name(i)
+            if "AMD" in device_name:
+                methods.add("rocm")
+            if "NVIDIA" in device_name:
+                methods.add("pynvml")
+            if "GH200" in device_name:
+                methods.add("gh")
+
+    power_methods = []
+    for m in methods:
+        if "rocm" == m:
+            from jpwr.gpu.rocm import power
+            power_methods.append(power())
+        if "pynvml" == m:
+            from jpwr.gpu.pynvml import power
+            power_methods.append(power())
+        if "gh" == m:
+            from jpwr.sys.gh import power
+            power_methods.append(power())
+
+
+    with get_power(power_methods, 100) as measured_scope:
+        pretrain(train_valid_test_datasets_provider,
+                 model_provider,
+                 ModelType.encoder_or_decoder,
+                 forward_step,
+                 args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})
+    energy_df, additional_data = measured_scope.energy()
+    import platform
+    nodename  = platform.node()
+    rankid    = int(os.getenv("RANK"))
+    power_file_base = os.getenv("ENERGY_PATH")
+    power_file = power_file_base.replace("csv", f"{rankid}.csv")
+    measured_scope.df["nodename"] = nodename
+    measured_scope.df["rank"] = rankid
+    if not os.path.exists(power_file):
+        measured_scope.df.to_csv(power_file)
+    energy_df["nodename"] = nodename
+    energy_df["rank"] = rankid
+    energy_file = power_file.replace("csv", f"energy.csv")
+    if not os.path.exists(energy_file):
+        energy_df.to_csv(energy_file)
+    print(f"Host: {nodename}")
+    print(f"Energy-per-GPU-list integrated(Wh): \n{energy_df.to_string()}")
+    for k,v in additional_data.items():
+        additional_path = power_file.replace("csv", f"{slugify(k)}.csv")
+        print(f"Writing {k} df to {additional_path}")
+        v.T.to_csv(additional_path)
+        print(f"Energy-per-GPU-list from {k}(Wh): {v.to_string()}")