UT

Author: wenchenvincent

csrc/extensions/include/hip_f8_impl.h

@@ -264,4 +264,4 @@ T cast_from_f8(uint8_t x) {
 }
 
 
-} // namespace hip_f8_impl
+} // namespace hip_f8_impl

csrc/extensions/include/hip_float8.h

@@ -378,4 +378,3 @@ __device__ hip_float32x4 mfma_f32_16x16x32(hip_f8x8<T_A> a, hip_f8x8<T_B> b, hip
 
 template<hip_f8_type T_A, hip_f8_type T_B>
 __device__ hip_float32x16 mfma_f32_32x32x16(hip_f8x8<T_A> a, hip_f8x8<T_B> b, hip_float32x16 c);
-

csrc/extensions/include/vectorized_pointwise.h

@@ -165,7 +165,13 @@ class VectorizedStorer : public VectorizedAccessor<DType, nvec, aligned> {
 
 
 constexpr int unary_kernel_threads = 512;
+
+#ifndef __HIP_PLATFORM_AMD__
 constexpr float e4m3_max = 448.0;
+#else
+constexpr float e4m3_max = 240.0;
+#endif
+
 constexpr float e5m2_max = 57344.0;
 
 extern __device__ msamp::DeviceSyncer device_syncer;

msamp/common/dtype/floating.py

@@ -6,6 +6,7 @@
 import torch
 import numpy as np
 
+from msamp.common.utils import Device, GPUType
 from msamp.common.dtype import Dtypes
 
 
@@ -26,6 +27,8 @@ def _get_fp_max(exp, man, inf_existed=True):
         Return:
             value (float): The float point value.
         """
+        if exp == 4 and man == 3 and Device.get_gpu_type() == GPUType.AMD:
+            return 240.0
         e_bias = np.power(2., exp - 1) - 1
         if inf_existed:
             max_value_exp = (np.power(2.0, exp) - 1) - e_bias - man - 1

msamp/common/utils/__init__.py

@@ -6,8 +6,8 @@
 from msamp.common.utils.logging import MsAmpLogger
 from msamp.common.utils.lazy_import import LazyImport
 from msamp.common.utils.dist import DistUtil
-from msamp.common.utils.device import Device
+from msamp.common.utils.device import Device, GPUType
 
 TransformerEngineWrapper = LazyImport('msamp.common.utils.transformer_engine_wrapper', 'TransformerEngineWrapper')
 
-__all__ = ['MsAmpLogger', 'TransformerEngineWrapper', 'DistUtil', 'Device']
+__all__ = ['MsAmpLogger', 'TransformerEngineWrapper', 'DistUtil', 'Device', 'GPUType']

msamp/common/utils/device.py

@@ -3,8 +3,14 @@
 
 """Device module."""
 
+from enum import Enum
+
 import torch
 
+class GPUType(Enum):
+    NVIDIA=1
+    AMD=2
+    UNKNOW=3
 
 class Device:
     """Device class for different hardwares."""
@@ -20,3 +26,14 @@ def is_fp8_supported():
             return True
 
         return False
+
+    @staticmethod
+    def get_gpu_type():
+        """Get the GPU type."""
+        if torch.cuda.device_count() > 0:
+            device_name = torch.cuda.get_device_name(0)
+            if "NVIDIA" in device_name:
+                return GPUType.NVIDIA
+            elif "AMD" in device_name:
+                return GPUType.AMD
+        return GPUType.UNKNOW

tests/common/dtype/test_floating.py

@@ -4,21 +4,25 @@
 """Tests for floating module."""
 
 import torch
+
+from msamp.common.utils import Device, GPUType
 from msamp.common.dtype import Floating
 from msamp.common.dtype import Dtypes
 
 
 def test_fp_max():
     """Test fp_max in Floating class."""
-    expect_fp_maxs = {torch.fp8e4m3: 448, torch.fp8e5m2: 57344, torch.float16: 65504}
+    fp8e4m3_max = Device.get_gpu_type() == GPUType.AMD and 240.0 or 448.0
+    expect_fp_maxs = {torch.fp8e4m3: fp8e4m3_max, torch.fp8e5m2: 57344, torch.float16: 65504}
 
     for k, v in expect_fp_maxs.items():
         assert Floating.fp_maxs[k] == v
 
 
 def test_qfp_max():
     """Test fp_max in Floating class."""
-    expected_qfp_maxs = {Dtypes.kfloat8_e4m3: 448, Dtypes.kfloat8_e5m2: 57344, Dtypes.kfloat16: 65504}
+    fp8e4m3_max = Device.get_gpu_type() == GPUType.AMD and 240.0 or 448.0
+    expected_qfp_maxs = {Dtypes.kfloat8_e4m3: fp8e4m3_max, Dtypes.kfloat8_e5m2: 57344, Dtypes.kfloat16: 65504}
 
     for k, v in expected_qfp_maxs.items():
         assert Floating.qfp_max[k] == v

tests/common/tensor/test_meta.py

@@ -6,6 +6,7 @@
 import torch
 import unittest
 
+from msamp.common.utils import Device, GPUType
 from msamp.common.dtype import Dtypes
 from msamp.common.dtype import Floating
 from msamp.common.tensor import ScalingMeta
@@ -29,13 +30,15 @@ def test_compute_scaling_factor(self):
         amax = torch.tensor(10, device='cuda')
         scale = torch.ones((), device='cuda')
         scale.copy_(ScalingMeta.compute_scaling_factor(amax, scale, fp_max, margin))
-        assert scale.item() == 32
+        expected_scale = 32 if Device.get_gpu_type() == GPUType.NVIDIA else 16
+        assert scale.item() == expected_scale
 
         # 1/(2^abs(floor(log2(448.0/10000))))
         amax = torch.tensor(10000, device='cuda')
         scale = torch.ones((), device='cuda')
         scale.copy_(ScalingMeta.compute_scaling_factor(amax, scale, fp_max, margin))
-        assert scale.item() == 1.0 / 32
+        expected_scale = 1.0 / 32 if Device.get_gpu_type() == GPUType.NVIDIA else 1.0 / 64
+        assert scale.item() == expected_scale
 
     def test_iswarmup_intime(self):
         """Test is_warmup and is_in_time_scaling i ScalingMeta."""

tests/fake/test_fake.py

@@ -26,9 +26,3 @@ def tearDown(self):
     def test_add(self):
         """Test add."""
         self.assertEqual(1 + 2, 3)
-
-    @decorator.cuda_test
-    def test_nvidia_smi(self):
-        """Test nvidia-smi."""
-        smi = os.popen('nvidia-smi').read().strip()
-        self.assertIn('NVIDIA-SMI', smi)