fix calculation error of mbu

src/backend/hflm_with_measurement.py

@@ -312,11 +312,14 @@ class HFLMWithMeasurement(HFLM):
         if do_sample is False and generation_kwargs.get("temperature") == 0.0:
             generation_kwargs.pop("temperature")
         
-        generation_kwargs.pop("is_gsm8k")
+        if is_gsm8k:
+            generation_kwargs.pop("is_gsm8k")
+            
         context_length = context.shape[1]
 
         if not is_gsm8k:
         # build stopping criteria
+            print("Using normal stopping criteria")
             stopping_criteria = stop_sequences_criteria(
                 self.tokenizer, stop, context.shape[1], context.shape[0]
             )
@@ -354,7 +357,6 @@ class HFLMWithMeasurement(HFLM):
         
         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
 
@@ -401,7 +403,7 @@ class HFLMWithMeasurement(HFLM):
         prefilling_time = stop_watch.prefilling_time / batch_size
         decoding_time = stop_watch.decoding_time / batch_size
         token_per_sec = output_length / decoding_time
-        ach_mem_bw = (model_size / 1e9 + kv_size) * token_per_sec
+        ach_mem_bw = (model_size * precision_bytes / 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)

src/utils.py

@@ -31,6 +31,12 @@ PEAK_FLOPS_DICT = {
         "NVIDIA-H100-PCIe-80GB": 1513e12,
         "NVIDIA-RTX-A5000-24GB": 444.4e12
     },
+    "bfloat16":{
+        "NVIDIA-A100-PCIe-80GB": 624e12,
+        "NVIDIA-A100-SXM-80GB": 624e12,
+        "NVIDIA-H100-PCIe-80GB": 1513e12,
+        "NVIDIA-RTX-A5000-24GB": 444.4e12
+    },
     "8bit":{
         "NVIDIA-A100-PCIe-80GB": 1248e12,
         "NVIDIA-A100-SXM-80GB": 1248e12,