Add mfu and mbu

backend-cli.py

@@ -473,6 +473,7 @@ if __name__ == "__main__":
         precisions = args.precision.split(",")
         print(f"debug_model_names: {debug_model_names}, debug_task_name: {debug_task_name}, precisions: {precisions}")
         task_lst = TASKS_HARNESS.copy()
+        RESULTS_REPO = DEBUG_RESULTS_REPO
         for precision in precisions:
             for debug_model_name in debug_model_names:
                 for task in task_lst:

src/backend/hflm_with_measurement.py

@@ -5,6 +5,7 @@ import sys
 from time import time
 from pathlib import Path
 from typing import List, Literal, Optional, Tuple, Union
+from calflops import calculate_flops
 
 import torch
 import torch.nn.functional as F
@@ -37,6 +38,9 @@ from lm_eval.models.utils import (
     stop_sequences_criteria,
 )
 from lm_eval.models.huggingface import HFLM
+from src.utils import get_gpu_number, get_gpu_details, get_peak_bw, transfer_precision2bytes, get_peak_flops
+from src.submission.check_validity import get_model_size
+from src.envs import API
 
 
 class StopWatch(TextStreamer):
@@ -67,6 +71,10 @@ class StopWatch(TextStreamer):
 class HFLMWithMeasurement(HFLM):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        self.pretrained = kwargs.get("pretrained", None)
+        self.revision = kwargs.get("revision", None)
+        self.precision = kwargs.get("dtype", None)
+        self.total_flops = 0
 
     def _loglikelihood_tokens(
         self,
@@ -306,6 +314,7 @@ class HFLMWithMeasurement(HFLM):
             generation_kwargs.pop("temperature")
         
         generation_kwargs.pop("is_gsm8k")
+        context_length = context.shape[1]
 
         if not is_gsm8k:
         # build stopping criteria
@@ -341,12 +350,71 @@ class HFLMWithMeasurement(HFLM):
 
         batch_size = context.shape[0]
         output_length = stop_watch.decoding_iterations
+        
+        precision_bytes = transfer_precision2bytes(self.precision)
+        
+        model_info = API.model_info(repo_id=self.pretrained, revision=self.revision)
+        model_size_param = get_model_size(model_info=model_info, precision=self.precision)
+        model_size = model_size_param * precision_bytes
+        
+        model_config = self.model.config
+
+        n_layers = model_config.num_hidden_layers if hasattr(model_config, "num_hidden_layers") else model_config.num_layers
+        d_model = model_config.hidden_size if hasattr(model_config, "hidden_size") else model_config.d_model
+        
+        if hasattr(model_config, "num_experts_per_tok"):
+            n_experts_per_tok = model_config.num_experts_per_tok
+        elif hasattr(model_config, "num_selected_experts"):
+            n_experts_per_tok = model_config.num_selected_experts
+        else:
+            n_experts_per_tok = 1
+        
+        if hasattr(model_config, "ffn_dim"):
+            d_ff = model_config.ffn_dim
+        elif hasattr(model_config, "intermediate_size"):
+            d_ff = model_config.intermediate_size
+        elif hasattr(model_config, "d_ff"):
+            d_ff = model_config.d_ff
+        else:
+            raise ValueError("Unknown ffn dim model configuration")
+        
+        if hasattr(model_config, "num_local_experts"):
+            num_experts = model_config.num_local_experts
+        elif hasattr(model_config, "num_experts"):
+            num_experts = model_config.num_experts
+        else:
+            num_experts = 1
+            
+        ffn_params = n_layers * d_ff * 2 * d_model
+        
+        shared_params = model_size_param * 1e9 - num_experts * ffn_params
+
+        model_size = shared_params + n_experts_per_tok * ffn_params
+
+        per_token_kv_size = 2 * n_layers * d_model * precision_bytes
+        
+        peak_bw_single = get_peak_bw(get_gpu_details())
+        peak_bw = peak_bw_single * get_gpu_number()
+        
+        kv_size = (output_length - 1) * per_token_kv_size / 1e9
 
         end_to_end_time = (end - start) / batch_size
         prefilling_time = stop_watch.prefilling_time / batch_size
         decoding_time = stop_watch.decoding_time / batch_size
         token_per_sec = output_length / decoding_time
-        return res, end_to_end_time, prefilling_time, token_per_sec
+        ach_mem_bw = (model_size / 1e9 + kv_size) * token_per_sec
+        
+        flops_per_token = 2 * model_size + 2 * n_layers * context_length * d_model
+        peak_flops_single = get_peak_flops(get_gpu_details(), self.precision)
+        peak_flops = peak_flops_single * get_gpu_number()
+        
+        ## TODO only support llama-type decoder only models and moe models of switch transformer and mixtrial
+        mfu = token_per_sec * flops_per_token / peak_flops
+        mbu = ach_mem_bw / peak_bw
+        
+        # print(f"mfu: {mfu}, mbu: {mbu}")
+        
+        return res, end_to_end_time, prefilling_time, token_per_sec, mfu, mbu
 
     def generate_until(
         self, requests: List[Instance], disable_tqdm: bool = False
@@ -461,7 +529,7 @@ class HFLMWithMeasurement(HFLM):
                 kwargs["max_length"] = context_enc.shape[1] + max_gen_toks
 
             # perform batched generation
-            cont, end_to_end_time, prefilling_time, token_per_sec = self._model_generate(
+            cont, end_to_end_time, prefilling_time, token_per_sec, mfu, mbu = self._model_generate(
                 context=context_enc,
                 attention_mask=attn_masks,
                 stop=until,
@@ -476,6 +544,8 @@ class HFLMWithMeasurement(HFLM):
                     cont_toks = cont_toks[context_enc.shape[1] :]
                 
                 s = self.tok_decode(cont_toks)
+                
+                # print(s)
 
                 # use secondary stop seqs to cut off should-have-been-stopped content post-hoc
                 if not is_gsm8k:
@@ -485,7 +555,7 @@ class HFLMWithMeasurement(HFLM):
                             # for seq2seq case where self.tok_decode(self.eot_token_id) = ''
                             s = s.split(term)[0]
 
-                res.append((s, end_to_end_time, prefilling_time, token_per_sec))
+                res.append((s, end_to_end_time, prefilling_time, token_per_sec, mfu, mbu))
 
                 self.cache_hook.add_partial("generate_until", (context, gen_kwargs), s)
                 pbar.update(1)

src/backend/run_eval_suite.py

@@ -17,12 +17,16 @@ def process_results_decorator(func):
         end_to_end_time = sum([r[1] for r in results]) / len(results)
         prefilling_time = sum([r[2] for r in results]) / len(results)
         decoding_throughput = sum([r[3] for r in results]) / len(results)
+        mfu = sum([r[4] for r in results]) / len(results)
+        mbu = sum([r[5] for r in results]) / len(results)
         # print(f"end_to_end_time: {end_to_end_time}, prefilling_time: {prefilling_time}, decoding_throughput: {decoding_throughput}")
 
         result_dict = func(self, doc, processed_results, *args, **kwargs)
         result_dict["end_to_end_time"] = end_to_end_time
         result_dict["prefilling_time"] = prefilling_time
         result_dict["decoding_throughput"] = decoding_throughput
+        result_dict["mfu"] = mfu * 100
+        result_dict["mbu"] = mbu * 100
         return result_dict
     return wrapper
 ConfigurableTask.process_results = process_results_decorator(orig_process_results)
@@ -33,6 +37,8 @@ def aggregation_decorator(func):
         aggregation_list["end_to_end_time"] = mean
         aggregation_list["prefilling_time"] = mean
         aggregation_list["decoding_throughput"] = mean
+        aggregation_list["mfu"] = mean
+        aggregation_list["mbu"] = mean
         return aggregation_list
     return wrapper
 ConfigurableTask.aggregation = aggregation_decorator(orig_aggregation)
@@ -43,6 +49,8 @@ def higher_is_better_decorator(func):
         higher_is_better_dict["end_to_end_time"] = False
         higher_is_better_dict["prefilling_time"] = False
         higher_is_better_dict["decoding_throughput"] = True
+        higher_is_better_dict["mfu"] = True
+        higher_is_better_dict["mbu"] = True
         return higher_is_better_dict
     return wrapper
 ConfigurableTask.higher_is_better = higher_is_better_decorator(orig_higher_is_better)

src/display/utils.py

@@ -18,12 +18,16 @@ GPU_Power = 'Power(W)'
 GPU_Mem = 'Mem(G)'
 GPU_Name = "GPU"
 GPU_Util = 'Util(%)'
+MFU = 'MFU(%)'
+MBU = 'MBU(%)'
 BATCH_SIZE = 'bs'
 PRECISION = "Precision"
 system_metrics_to_name_map = {
     "end_to_end_time": f"{E2Es}",
     "prefilling_time": f"{PREs}",
     "decoding_throughput": f"{TS}",
+    "mfu": f"{MFU}",
+    "mbu": f"{MBU}"
 }
 
 gpu_metrics_to_name_map = {

src/submission/check_validity.py

@@ -74,7 +74,7 @@ def is_model_on_hub(
 
 
 def get_model_size(model_info: ModelInfo, precision: str):
-    size_pattern = size_pattern = re.compile(r"(\d\.)?\d+(b|m)")
+    size_pattern = re.compile(r"(\d\.)?\d+(b|m)")
     try:
         model_size = round(model_info.safetensors["total"] / 1e9, 3)
     except (AttributeError, TypeError):