Use torch RMSNorm

modeling_flashppi.py

@@ -16,7 +16,6 @@ try:
     from flash_attn.layers.rotary import apply_rotary_emb_func
     from flash_attn import flash_attn_varlen_kvpacked_func
     from flash_attn.bert_padding import pad_input, unpad_input
-    from flash_attn.ops.triton.layer_norm import RMSNorm
     FLASH_ATTN_AVAILABLE = True
 except ImportError:
     FLASH_ATTN_AVAILABLE = False
@@ -26,21 +25,20 @@ except ImportError:
     def swiglu(x, y):
         return F.silu(x) * y
 
-    class RMSNorm(nn.Module):
-        """RMSNorm without variance_epsilon buffer for checkpoint compatibility."""
-        def __init__(self, dim, eps=1e-6):
-            super().__init__()
-            self.weight = nn.Parameter(torch.ones(dim))
-            self.eps = eps
-
-        def forward(self, hidden_states):
-            input_dtype = hidden_states.dtype
-            hidden_states = hidden_states.to(torch.float32)
-            variance = hidden_states.pow(2).mean(-1, keepdim=True)
-            hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-            return (self.weight * hidden_states).to(input_dtype)
-
-
+class RMSNorm(nn.Module):
+    """RMSNorm without variance_epsilon buffer for checkpoint compatibility."""
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight * hidden_states).to(input_dtype)
+    
 @dataclass
 class FlashPPIOutput(ModelOutput):
     """Output type for FlashPPI model.