Merge pull request #377 from Modalities/feat--support-NVIDIA-GH200-120GB

feat: supports now gh200 mfu calculations

Author: le1nux

src/modalities/utils/mfu.py

@@ -14,10 +14,7 @@
 #       https://www.nvidia.com/en-us/data-center/h100/
 #
 # NOTE: These values are valid for fp16 and bf16 only
-PEAK_PERFORMANCE = {
-    "A100": 312e12,
-    "H100": 989e12,
-}
+PEAK_PERFORMANCE = {"A100": 312e12, "H100": 989e12, "GH200": 989e12}
 
 
 class MFUCalculatorABC:
@@ -128,6 +125,10 @@ def _get_theoretical_gpu_peak_performance(wrapped_model: FSDPX, world_size: int)
             single_gpu_peak_performance = MFUCalculatorABC._get_theoretical_gpu_peak_performance_single(
                 precision, "H100"
             )
+        elif device_name.startswith("NVIDIA GH200"):
+            single_gpu_peak_performance = MFUCalculatorABC._get_theoretical_gpu_peak_performance_single(
+                precision, "GH200"
+            )
         else:
             warnings.warn(f"Could not get theoretical GPU peak performance for unknown device = {device_name}.")
             return None