Merge branch 'flagos-ai:main' into te_groups_gemm

Author: chai-xiaonan

transformer_engine/__init__.py

@@ -21,9 +21,7 @@
     from .plugin.core.backends.vendor.musa.patches import apply_patch as _musa_apply_patch
 
     _musa_apply_patch()
-    print("[TE-FL] MUSA patches applied")
 except Exception as e:
-    print(f"[TE-FL] MUSA patches not applied: {e}")
     pass
 
 

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

@@ -79,7 +79,6 @@ def try_load_lib(name, search_patterns):
                     return True
         return False
     except Exception as e:
-        print(f"[CUDA] Failed to load CUDA libs: {e}")
         return False
 
 
@@ -90,6 +89,8 @@ def _ensure_cuda_libs():
     global _cuda_libs_loaded
     if not _cuda_libs_loaded:
         _cuda_libs_loaded = _load_cuda_libs()
+        if _cuda_libs_loaded:
+            print(f"[CUDA] Successfully loaded CUDA libs")
     return _cuda_libs_loaded
 
 
@@ -166,6 +167,15 @@ def quantize(
         noop: Optional[torch.Tensor] = None,
     ) -> Any:
         tex = self._get_tex()
+        # Normalize quantizer.dtype 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(
@@ -183,6 +193,16 @@ def bgrad_quantize(
         quantizer: Any,
     ) -> List[Any]:
         tex = self._get_tex()
+
+        # Normalize quantizer.dtype 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.bgrad_quantize(input, quantizer)
 
     def generic_gemm(

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

@@ -54,7 +54,6 @@ def _get_sys_extension() -> str:
         spec.loader.exec_module(solib)
         return True
     except Exception as e:
-        print(f"[HYGON] Failed to load hygon libs: {e}")
         return False
 
 
@@ -65,6 +64,8 @@ def _ensure_hygon_libs():
     global _hygon_libs_loaded
     if not _hygon_libs_loaded:
         _hygon_libs_loaded = _load_hygon_libs()
+        if _hygon_libs_loaded:
+            print(f"[HYGON] Successfully loaded HYGON libs")
     return _hygon_libs_loaded
 
 
@@ -145,6 +146,13 @@ 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:
+                    quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(
@@ -162,6 +170,16 @@ def bgrad_quantize(
         quantizer: Any,
     ) -> List[Any]:
         tex = self._get_tex()
+
+        # Normalize quantizer.dtype 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.bgrad_quantize(input, quantizer)
 
     def generic_gemm(

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

@@ -77,7 +77,6 @@ def try_load_lib(name, search_patterns):
                     return True
         return False
     except Exception as e:
-        print(f"[ILUVATAR] Failed to load ILUVATAR libs: {e}")
         return False
 
 
@@ -88,6 +87,8 @@ def _ensure_iluvatar_libs():
     global _iluvatar_libs_loaded
     if not _iluvatar_libs_loaded:
         _iluvatar_libs_loaded = _load_iluvatar_libs()
+        if _iluvatar_libs_loaded:
+            print(f"[ILUVATAR] Successfully loaded ILUVATAR libs")
     return _iluvatar_libs_loaded
 
 
@@ -171,6 +172,13 @@ 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:
+                    quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(
@@ -188,6 +196,16 @@ def bgrad_quantize(
         quantizer: Any,
     ) -> List[Any]:
         tex = self._get_tex()
+
+        # Normalize quantizer.dtype 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.bgrad_quantize(input, quantizer)
 
     def generic_gemm(

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

@@ -37,7 +37,6 @@ def get_ext():
                     return True
         return False
     except Exception as e:
-        print(f"[Metax] Failed to load Metax libs: {e}")
         return False
 
 
@@ -48,6 +47,8 @@ def _ensure_metax_libs():
     global _metax_libs_loaded
     if not _metax_libs_loaded:
         _metax_libs_loaded = _load_metax_libs()
+        if _metax_libs_loaded:
+            print(f"[Metax] Successfully loaded Metax libs")
     return _metax_libs_loaded
 
 
@@ -126,6 +127,13 @@ 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:
+                    quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(
@@ -143,6 +151,16 @@ def bgrad_quantize(
         quantizer: Any,
     ) -> List[Any]:
         tex = self._get_tex()
+
+        # Normalize quantizer.dtype 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.bgrad_quantize(input, quantizer)
 
     def generic_gemm(

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

@@ -66,7 +66,6 @@ def try_load_lib(name, search_patterns):
 
         return True
     except Exception as e:
-        print(f"[MUSA] Failed to load MUSA libs: {e}")
         return False
 
 
@@ -77,6 +76,8 @@ def _ensure_musa_libs():
     global _musa_libs_loaded
     if not _musa_libs_loaded:
         _musa_libs_loaded = _load_musa_libs()
+        if _musa_libs_loaded:
+            print(f"[MUSA] Successfully loaded MUSA libs")
     return _musa_libs_loaded
 
 
@@ -138,6 +139,13 @@ 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:
+                    quantizer.dtype = tex.DType(int(qdtype))
+        except Exception:
+            pass
         return tex.quantize(tensor, quantizer, output, noop)
 
     def dequantize(
@@ -155,6 +163,16 @@ def bgrad_quantize(
         quantizer: Any,
     ) -> List[Any]:
         tex = self._get_tex()
+
+        # Normalize quantizer.dtype 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.bgrad_quantize(input, quantizer)
 
     def generic_gemm(

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

@@ -70,3 +70,4 @@ def apply_patch() -> None:
         except Exception:
             # Best-effort: patching should never crash import/initialization.
             continue
+    print(f"[TE-FL] MUSA backend patches applied")