fix quantizer dtype

Author: lxd-cumt

transformer_engine/plugin/core/backends/vendor/cuda/cuda.py

@@ -166,6 +166,17 @@ def quantize(
         noop: Optional[torch.Tensor] = None,
     ) -> Any:
         tex = self._get_tex()
+        # Quantizer.dtype may originate from a different TE extension module
+        # (e.g. transformer_engine_torch vs transformer_engine_torch_nv) which leads
+        # to a pybind enum cast failure. Normalize it to this backend's `tex.DType`.
+        try:
+            if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
+                qdtype = quantizer.dtype
+                if qdtype is not None:
+                    quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
+
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(

transformer_engine/plugin/core/backends/vendor/hygon/hygon.py

@@ -145,6 +145,14 @@ def quantize(
         noop: Optional[torch.Tensor] = None,
     ) -> Any:
         tex = self._get_tex()
+        try:
+            if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
+                qdtype = quantizer.dtype
+                if qdtype is not None:
+                    if qdtype.__class__ is not tex.DType:
+                        quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(

transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py

@@ -171,6 +171,14 @@ def quantize(
         noop: Optional[torch.Tensor] = None,
     ) -> Any:
         tex = self._get_tex()
+        try:
+            if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
+                qdtype = quantizer.dtype
+                if qdtype is not None:
+                    if qdtype.__class__ is not tex.DType:
+                        quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(

transformer_engine/plugin/core/backends/vendor/metax/metax.py

@@ -126,6 +126,14 @@ def quantize(
         noop: Optional[torch.Tensor] = None,
     ) -> Any:
         tex = self._get_tex()
+        try:
+            if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
+                qdtype = quantizer.dtype
+                if qdtype is not None:
+                    if qdtype.__class__ is not tex.DType:
+                        quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(

transformer_engine/plugin/core/backends/vendor/musa/musa.py

@@ -138,6 +138,14 @@ def quantize(
         noop: Optional[torch.Tensor] = None,
     ) -> Any:
         tex = self._get_tex()
+        try:
+            if quantizer is not None and hasattr(quantizer, "dtype") and hasattr(tex, "DType"):
+                qdtype = quantizer.dtype
+                if qdtype is not None:
+                    if qdtype.__class__ is not tex.DType:
+                        quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(