fix other vendor , flagos , torch func mismatch

Author: lihongyang1990

transformer_engine/plugin/core/backends/flagos/flagos.py

@@ -7,7 +7,7 @@
 
 import torch
 
-from ...ops import TEFLBackendBase, FP8TensorMeta, NVTE_Fused_Attn_Backend
+from ...ops import *
 
 from .impl import (
     rmsnorm_fwd_fl, rmsnorm_bwd_fl,
@@ -20,7 +20,6 @@
 def _check_flagos_available() -> bool:
     return True
 
-
 class FlagOSBackend(TEFLBackendBase):
     @staticmethod
     def check_available() -> bool:
@@ -29,10 +28,6 @@ def check_available() -> bool:
     def is_available(self) -> bool:
         return _check_flagos_available()
 
-    def get_flash_attention_class(self):
-        from .attention.dot_product_attention.backends import FlashAttentionFL
-        return FlashAttentionFL
-
     def get_attention_backend(self, attention_params=None):
         from packaging.version import Version as PkgVersion
         from ...logger_manager import get_logger
@@ -65,48 +60,50 @@ def get_attention_backend(self, attention_params=None):
             available_backends,
         )
 
+##### transformer_engine/pytorch/csrc/extensions/pybind.cpp #####
     def generic_gemm(
         self,
-        A: torch.Tensor,
-        transA: bool,
-        B: torch.Tensor,
-        transB: bool,
-        D: torch.Tensor,
+        A: Any,
+        transa: bool,
+        B: Any,
+        transb: bool,
+        D: Any,
         quantizer: Any,
-        output_dtype: torch.dtype,
+        out_dtype: Optional[DType],
         bias: Optional[torch.Tensor],
-        bias_type: Any,
+        bias_type: DType,
         gelu: bool,
         gelu_in: Optional[torch.Tensor],
         grad: bool,
         workspace: torch.Tensor,
-        workspace_size: int,
+        workspaceSize: int,
         accumulate: bool,
         use_split_accumulator: bool,
         comm_overlap: Optional[Any] = None,
-        comm_type: Optional[Any] = None,
+        comm_type: Optional[CommOverlapType] = None,
         extra_output: Optional[torch.Tensor] = None,
         bulk_overlap: bool = False,
         alpha: float = 1.0,
         beta: Optional[float] = None,
-    ) -> Any:
+    ) -> List[Any]:
         return generic_gemm_fl(
-            A, transA, B, transB, D, quantizer, output_dtype,
+            A, transa, B, transb, D, quantizer, out_dtype,
             bias, bias_type, gelu, gelu_in, grad,
-            workspace, workspace_size, accumulate, use_split_accumulator,
+            workspace, workspaceSize, accumulate, use_split_accumulator,
             comm_overlap=comm_overlap, comm_type=comm_type,
             extra_output=extra_output, bulk_overlap=bulk_overlap,
             alpha=alpha, beta=beta
         )
 
+    # Other granular functions
     def rmsnorm_fwd(
         self,
         input: torch.Tensor,
         weight: torch.Tensor,
         eps: float,
         ln_out: Optional[torch.Tensor],
         quantizer: Any,
-        otype: torch.dtype,
+        otype: DType,
         sm_margin: int,
         zero_centered_gamma: bool,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
@@ -115,22 +112,23 @@ def rmsnorm_fwd(
             quantizer=quantizer, odtype=otype,
             sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma,
         )
-
     def rmsnorm_bwd(
         self,
-        dy: torch.Tensor,
+        dz: torch.Tensor,
         x: torch.Tensor,
         rsigma: torch.Tensor,
         gamma: torch.Tensor,
         sm_margin: int = 0,
         zero_centered_gamma: bool = False,
-        eps: float = 1e-5,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         return rmsnorm_bwd_fl(
-            dy=dy, x=x, rsigma=rsigma, gamma=gamma,
-            sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma, eps=eps,
+            dy=dz, x=x, rsigma=rsigma, gamma=gamma,
+            sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma
         )
+    def get_fused_attn_backend(self, *args, **kwargs) -> int:
+        return NVTE_Fused_Attn_Backend.NVTE_No_Backend
 
+    # multi-tensor functions
     def multi_tensor_scale(
         self,
         chunk_size: int,
@@ -139,73 +137,66 @@ def multi_tensor_scale(
         scale: float,
     ) -> None:
         return multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale)
-
     def multi_tensor_l2norm(
         self,
         chunk_size: int,
         noop_flag: torch.Tensor,
         tensor_lists: List[List[torch.Tensor]],
-        per_tensor: bool = False,
-    ) -> Union[torch.Tensor, List[torch.Tensor]]:
-        result, _ = multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor)
-        return result
-
+        per_tensor: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor)
     def multi_tensor_adam(
         self,
-        chunk_size: int = None,
-        noop_flag: torch.Tensor = None,
-        tensor_lists: List[List[torch.Tensor]] = None,
-        lr: float = None,
-        beta1: float = None,
-        beta2: float = None,
-        eps: float = None,
-        step: int = None,
-        mode: int = None,
-        bias_correction: int = None,
-        weight_decay: float = None,
-    ):
+        chunk_size: int,
+        noop_flag: torch.Tensor,
+        tensor_lists: List[List[torch.Tensor]],
+        lr: float,
+        beta1: float,
+        beta2: float,
+        epsilon: float,
+        step: int,
+        mode: int,
+        bias_correction: int,
+        weight_decay: float,
+    ) -> None:
         if chunk_size is None:
             return multi_tensor_adam_fl
         return multi_tensor_adam_fl(
             chunk_size=chunk_size, noop_flag=noop_flag, tensor_lists=tensor_lists,
-            lr=lr, beta1=beta1, beta2=beta2, eps=eps,
+            lr=lr, beta1=beta1, beta2=beta2, eps=epsilon,
             step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
         )
-
     def multi_tensor_adam_param_remainder(
         self,
-        chunk_size: int = None,
-        noop_flag: torch.Tensor = None,
-        tensor_lists: List[List[torch.Tensor]] = None,
-        lr: float = None,
-        beta1: float = None,
-        beta2: float = None,
-        eps: float = None,
-        step: int = None,
-        mode: int = None,
-        bias_correction: int = None,
-        weight_decay: float = None,
-    ):
+        chunk_size: int,
+        noop_flag: torch.Tensor,
+        tensor_lists: List[List[torch.Tensor]],
+        lr: float,
+        beta1: float,
+        beta2: float,
+        epsilon: float,
+        step: int,
+        mode: int,
+        bias_correction: int,
+        weight_decay: float,
+    ) -> None:
         if chunk_size is None:
             return multi_tensor_adam_param_remainder_fl
         return multi_tensor_adam_param_remainder_fl(
             chunk_size=chunk_size, noop_flag=noop_flag, tensor_lists=tensor_lists,
-            lr=lr, beta1=beta1, beta2=beta2, eps=eps,
+            lr=lr, beta1=beta1, beta2=beta2, eps=epsilon,
             step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
         )
 
+    # Misc
     def get_cublasLt_version(self) -> int:
         return 110000
-
     def get_cudnn_version(self) -> int:
         return 90000
-
     def get_num_cublas_streams(self) -> int:
         return 0
 
-    def get_fused_attn_backend(self, *args, **kwargs) -> int:
-        return NVTE_Fused_Attn_Backend.NVTE_No_Backend
-
-    def create_fp8_tensor_meta(self) -> FP8TensorMeta:
-        return FP8TensorMeta()
-
+############## class func #################################
+    def get_flash_attention_class(self):
+        from .attention.dot_product_attention.backends import FlashAttentionFL
+        return FlashAttentionFL

transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py

@@ -2,25 +2,60 @@
 #
 # See LICENSE for license information.
 
+from typing import List, Tuple
 import torch
-from torch.distributed._tensor import DTensor
 import flag_gems
 
 
-def multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor, *args):
+def multi_tensor_l2_norm_fl(
+    _chunk_size: int,
+    noop_flag: torch.Tensor,
+    tensor_lists: List[List[torch.Tensor]],
+    per_tensor: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute L2 norm of tensors using flag_gems.
+
+    Returns:
+        Tuple of (total_norm, per_tensor_norms_or_dummy)
+        - total_norm: The combined L2 norm of all tensors
+        - per_tensor_norms_or_dummy: Per-tensor norms stacked if per_tensor=True, else dummy tensor
+    """
+    device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu'
+
+    if noop_flag.item() != 0:
+        return torch.tensor(0.0, device=device), torch.tensor(0.0, device=device)
 
     tensors = tensor_lists[0]
 
+    # Compute per-tensor norms
+    per_tensor_norms = []
+    total_norm_sq = torch.tensor(0.0, device=device)
+
+    for tensor in tensors:
+        norm_sq = flag_gems.sum(tensor.float() ** 2)
+        total_norm_sq = total_norm_sq + norm_sq
+        if per_tensor:
+            per_tensor_norms.append(flag_gems.sqrt(norm_sq))
+
+    total_norm = flag_gems.sqrt(total_norm_sq)
+
     if per_tensor:
-        norms = [torch.norm(t.float(), p=2) for t in tensors]
-        return norms, None
+        per_tensor_result = torch.stack(per_tensor_norms)
     else:
-        total_norm_sq = sum(flag_gems.sum(flag_gems.pow_func(t.float(), 2)) for t in tensors)
-        total_norm = flag_gems.sqrt(total_norm_sq)
-        return total_norm, None
+        per_tensor_result = torch.tensor(0.0, device=device)
+
+    return total_norm, per_tensor_result
 
 
-def multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale):
+def multi_tensor_scale_fl(
+    _chunk_size: int,
+    noop_flag: torch.Tensor,
+    tensor_lists: List[List[torch.Tensor]],
+    scale: float,
+) -> None:
+    if noop_flag.item() != 0:
+        return
 
     for src, dst in zip(tensor_lists[0], tensor_lists[1]):
         flag_gems.copy_(dst, src * scale)

transformer_engine/plugin/core/backends/flagos/impl/rmsnorm.py

@@ -42,7 +42,7 @@ def rmsnorm_bwd_fl(
     gamma,
     sm_margin,
     zero_centered_gamma,
-    eps,
+    eps=1e-5,
 ):
     # When zero_centered_gamma is True, forward uses (1 + gamma) as weight
     # So backward needs to use (1 + gamma) for computing dx

transformer_engine/plugin/core/backends/reference/impl/normalization.py

@@ -5,12 +5,37 @@
 from typing import Any, Optional, Tuple
 import torch
 import torch.nn.functional as F
+from ....ops import DType
 
 __all__ = [
     "layernorm_fwd_torch",
     "layernorm_bwd_torch",
 ]
 
+# Mapping from DType enum to torch.dtype
+_DTYPE_TO_TORCH_DTYPE = {
+    DType.kByte: torch.uint8,
+    DType.kInt16: torch.int16,
+    DType.kInt32: torch.int32,
+    DType.kInt64: torch.int64,
+    DType.kFloat32: torch.float32,
+    DType.kFloat16: torch.float16,
+    DType.kBFloat16: torch.bfloat16,
+    DType.kFloat8E4M3: torch.float8_e4m3fn,
+    DType.kFloat8E5M2: torch.float8_e5m2,
+}
+
+def _to_torch_dtype(dtype):
+    """Convert DType enum to torch.dtype."""
+    if dtype is None:
+        return None
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    if isinstance(dtype, (int, DType)):
+        dtype_enum = DType(dtype)
+        if dtype_enum in _DTYPE_TO_TORCH_DTYPE:
+            return _DTYPE_TO_TORCH_DTYPE[dtype_enum]
+    raise ValueError(f"Unsupported dtype: {dtype}")
 
 def layernorm_fwd_torch(
     input: torch.Tensor,
@@ -19,10 +44,11 @@ def layernorm_fwd_torch(
     eps: float,
     ln_out: Optional[torch.Tensor],
     quantizer: Any,
-    odtype: torch.dtype,
+    odtype: DType,
     sm_margin: int,
     zero_centered_gamma: bool,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    odtype = _to_torch_dtype(odtype)
     mean = input.mean(dim=-1, keepdim=True)
     var = input.var(dim=-1, keepdim=True, unbiased=False)
     rsigma = torch.rsqrt(var + eps)
@@ -45,7 +71,6 @@ def layernorm_fwd_torch(
 
     return output, mean, rsigma
 
-
 def layernorm_bwd_torch(
     dy: torch.Tensor,
     x: torch.Tensor,