fuse t5 layernorm without gamma

Author: nihui

tools/pnnx/src/pass_level5/fuse_rmsnorm.cpp

@@ -54,6 +54,71 @@ pnnx.Output             output      1 0 out
     }
 };
 
+class fuse_rmsnorm_pass_without_gamma : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input       0 1 input
+pnnx.Expression         op_0        1 1 input sq expr=pow(@0,2)
+torch.mean              op_1        1 1 sq sqmean dim=(-1) keepdim=True
+pnnx.Expression         op_2        2 1 input sqmean out expr=mul(@0,rsqrt(add(@1,%eps)))
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "nn.RMSNorm";
+    }
+
+    const char* name_str() const
+    {
+        return "t5ln";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& /*captured_params*/, const std::map<std::string, Attribute>& /*captured_attrs*/) const
+    {
+        const Operator* op_0 = matched_operators.at("op_0");
+        const std::vector<int>& shape = op_0->inputs[0]->shape;
+        if (shape.empty())
+        {
+            // unknown normalized_shape
+            return false;
+        }
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        const std::vector<int>& shape = op->inputs[0]->shape;
+        const int c = shape[shape.size() - 1];
+
+        op->params["elementwise_affine"] = false;
+        op->params["eps"] = captured_params.at("eps");
+        op->params["normalized_shape"] = std::vector<int>{c};
+    }
+};
+
+class fuse_rmsnorm_pass_without_gamma_1 : public fuse_rmsnorm_pass_without_gamma
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input       0 1 input
+pnnx.Expression         op_0        1 1 input sq expr=pow(@0,2)
+torch.mean              op_1        1 1 sq sqmean dim=(-1) keepdim=True
+pnnx.Expression         op_2        2 1 input sqmean out expr=mul(@0,reciprocal(sqrt(add(@1,%eps))))
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
 class fuse_rmsnorm_pass_onnx : public fuse_rmsnorm_pass
 {
 public:
@@ -75,11 +140,15 @@ void fuse_rmsnorm(Graph& graph)
 {
     fuse_rmsnorm_pass a;
     fuse_rmsnorm_pass_1 a1;
+    fuse_rmsnorm_pass_without_gamma a2;
+    fuse_rmsnorm_pass_without_gamma_1 a3;
     fuse_rmsnorm_pass_onnx b;
     int opindex = 0;
 
     pnnx_graph_rewrite(graph, &a, opindex);
     pnnx_graph_rewrite(graph, &a1, opindex);
+    pnnx_graph_rewrite(graph, &a2, opindex);
+    pnnx_graph_rewrite(graph, &a3, opindex);
     pnnx_graph_rewrite(graph, &b, opindex);
 }
 

tools/pnnx/tests/test_F_rms_norm.py

@@ -6,6 +6,32 @@
 import torch.nn.functional as F
 from packaging import version
 
+class T5LayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.rand(hidden_size))
+        self.variance_epsilon = 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.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+class T5LayerNorm_without_gamma(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = 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.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
 class Model(nn.Module):
     def __init__(self):
         super(Model, self).__init__()
@@ -14,7 +40,10 @@ def __init__(self):
         self.w4 = nn.Parameter(torch.rand(12, 16))
         self.w5 = nn.Parameter(torch.rand(24))
 
-    def forward(self, x, y, z, w0, w1, w2):
+        self.rmsnorm = T5LayerNorm(66)
+        self.rmsnorm_2 = T5LayerNorm_without_gamma(66)
+
+    def forward(self, x, y, z, w0, w1, w2, x2):
         x = F.rms_norm(x, (24,), w0)
         x = F.rms_norm(x, (12,24), None)
         x = F.rms_norm(x, (24,), self.w3)
@@ -26,7 +55,10 @@ def forward(self, x, y, z, w0, w1, w2):
         z = F.rms_norm(z, (24,), w2)
         z = F.rms_norm(z, (12,16,24), None, eps=1e-2)
         z = F.rms_norm(z, (24,), self.w5)
-        return x, y, z
+
+        x2 = self.rmsnorm(x2)
+        x2 = self.rmsnorm_2(x2)
+        return x, y, z, x2
 
 def test():
     if version.parse(torch.__version__) < version.parse('2.4'):
@@ -42,22 +74,26 @@ def test():
     w0 = torch.rand(24)
     w1 = torch.rand(12, 16)
     w2 = torch.rand(24)
+    x2 = torch.rand(3, 22, 66)
 
-    a0, a1, a2 = net(x, y, z, w0, w1, w2)
+    a = net(x, y, z, w0, w1, w2, x2)
 
     # export torchscript
-    mod = torch.jit.trace(net, (x, y, z, w0, w1, w2))
+    mod = torch.jit.trace(net, (x, y, z, w0, w1, w2, x2))
     mod.save("test_F_rms_norm.pt")
 
     # torchscript to pnnx
     import os
-    os.system("../src/pnnx test_F_rms_norm.pt inputshape=[1,12,24],[2,3,12,16],[1,10,12,16,24],[24],[12,16],[24]")
+    os.system("../src/pnnx test_F_rms_norm.pt inputshape=[1,12,24],[2,3,12,16],[1,10,12,16,24],[24],[12,16],[24],[3,22,66]")
 
     # pnnx inference
     import test_F_rms_norm_pnnx
-    b0, b1, b2 = test_F_rms_norm_pnnx.test_inference()
+    b = test_F_rms_norm_pnnx.test_inference()
 
-    return torch.equal(a0, b0) and torch.equal(a1, b1) and torch.equal(a2, b2)
+    for a0, b0 in zip(a, b):
+        if not torch.equal(a0, b0):
+            return False
+    return True
 
 if __name__ == "__main__":
     if test():