feat: add AMD Instinct peak FLOPS for MFU on ROCm

src/zeroband/utils/__init__.py

@@ -30,7 +30,7 @@ def get_sharding_strategy(sharding_strategy: str) -> ShardingStrategy:
 ### code above inspired and copied from https://github.com/pytorch/torchtitan/blob/4b3f2e41a084bf79a8540068ed525539d1244edd/torchtitan/utils.py#L119
 
 
-# hardcoded BF16 type peak flops for NVIDIA A100 and H100 GPU
+# hardcoded BF16 type peak flops for NVIDIA A100/H100 and AMD Instinct MI GPUs
 def get_peak_flops(device_name: str) -> int:
     if "A100" in device_name:
         # data from https://www.nvidia.com/en-us/data-center/a100/
@@ -44,6 +44,20 @@ def get_peak_flops(device_name: str) -> int:
             return 756e12
         else:  # for H100 SXM and other variants
             return 989e12
+    elif "MI300" in device_name:
+        # AMD Instinct MI300X/MI300A dense matrix peak (BF16/FP16)
+        # https://www.amd.com/en/products/accelerators/instinct/mi300/mi300x.html
+        return 1300e12
+    elif "MI250" in device_name:
+        # AMD Instinct MI250X dense matrix peak (BF16/FP16)
+        # https://www.amd.com/en/products/accelerators/instinct/mi250.html
+        return 383e12
+    elif "MI210" in device_name:
+        # AMD Instinct MI210 dense matrix peak (BF16/FP16)
+        return 181e12
+    elif "MI100" in device_name:
+        # AMD Instinct MI100 dense matrix peak (FP16)
+        return 184.6e12
     else:  # for other GPU types, assume A100
         return 312e12
 

tests/test_peak_flops.py

@@ -0,0 +1,25 @@
+import pytest
+
+from zeroband.utils import get_peak_flops
+
+
+@pytest.mark.parametrize(
+    ("device_name", "expected"),
+    [
+        ("NVIDIA A100-SXM4-80GB", 312e12),
+        ("NVIDIA H100 80GB HBM3", 989e12),
+        ("NVIDIA H100 PCIe", 756e12),
+        ("NVIDIA H100 NVL", 835e12),
+        ("AMD Instinct MI300X", 1300e12),
+        ("AMD Instinct MI300A", 1300e12),
+        ("AMD Instinct MI250X", 383e12),
+        ("AMD Instinct MI210", 181e12),
+        ("AMD Instinct MI100", 184.6e12),
+    ],
+)
+def test_get_peak_flops_known_devices(device_name: str, expected: float):
+    assert get_peak_flops(device_name) == expected
+
+
+def test_get_peak_flops_unknown_defaults_to_a100():
+    assert get_peak_flops("Some Unknown GPU") == 312e12