Fix: Resolve parameter mismatch between TE_FL and NVTE functions

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,17 +60,18 @@ def get_attention_backend(self, attention_params=None):
             available_backends,
         )
 
+##### transformer_engine/pytorch/csrc/extensions/pybind.cpp #####
     def generic_gemm(
         self,
-        A: torch.Tensor,
+        A: Any,
         transA: bool,
-        B: torch.Tensor,
+        B: Any,
         transB: bool,
-        D: torch.Tensor,
+        D: Any,
         quantizer: Any,
-        output_dtype: torch.dtype,
+        output_dtype: Optional[DType],
         bias: Optional[torch.Tensor],
-        bias_type: Any,
+        bias_type: DType,
         gelu: bool,
         gelu_in: Optional[torch.Tensor],
         grad: bool,
@@ -84,53 +80,53 @@ def generic_gemm(
         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,
-            bias, bias_type, gelu, gelu_in, grad,
-            workspace, workspace_size, accumulate, use_split_accumulator,
-            comm_overlap=comm_overlap, comm_type=comm_type,
-            extra_output=extra_output, bulk_overlap=bulk_overlap,
-            alpha=alpha, beta=beta
+            bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size,
+            accumulate, use_split_accumulator, comm_overlap, comm_type,
+            extra_output, bulk_overlap, alpha, beta
         )
 
+    # Other granular functions
     def rmsnorm_fwd(
         self,
-        input: torch.Tensor,
-        weight: torch.Tensor,
+        input: Any,
+        weight: Any,
         eps: float,
-        ln_out: Optional[torch.Tensor],
+        ln_out: Any,
         quantizer: Any,
-        otype: torch.dtype,
+        otype: DType,
         sm_margin: int,
         zero_centered_gamma: bool,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
+    ) -> List[Any]:
         return rmsnorm_fwd_fl(
             input=input, weight=weight, eps=eps, ln_out=ln_out,
             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]:
+        sm_margin: int,
+        zero_centered_gamma: bool,
+    ) -> List[Any]:
         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 +135,61 @@ 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,
-    ):
-        if chunk_size is None:
-            return multi_tensor_adam_fl
+        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:
         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,
-            step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
+            chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, epsilon,
+            step, mode, bias_correction, 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,
-    ):
-        if chunk_size is None:
-            return multi_tensor_adam_param_remainder_fl
+        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:
         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,
-            step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay,
+            chunk_size, noop_flag, tensor_lists,
+            lr, beta1, beta2, epsilon,
+            step, mode, bias_correction, 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/fused_adam.py

@@ -187,4 +187,4 @@ def multi_tensor_adam_param_remainder_fl(
 
         # Write back
         flag_gems.copy_(p, param_bf16)
-        flag_gems.copy_(p_remainder, remainder_int16)
+        flag_gems.copy_(p_remainder, remainder_int16)

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

@@ -23,4 +23,4 @@ def multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor, *ar
 def multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale):
 
     for src, dst in zip(tensor_lists[0], tensor_lists[1]):
-        flag_gems.copy_(dst, src * scale)
+        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,

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

@@ -310,4 +310,4 @@ def multi_tensor_compute_scale_and_scale_inv_torch(
 
         # Update scale and scale_inv
         scale.copy_(computed_scale)
-        scale_inv.copy_(1.0 / computed_scale)
+        scale_inv.copy_(1.0 / computed_scale)

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

@@ -43,7 +43,6 @@ def rmsnorm_bwd_torch(
     gamma,
     sm_margin,
     zero_centered_gamma,
-    eps,
 ):
     inv_rms = rsigma.unsqueeze(-1)