In GatedFeedforward, add support for a single gating einsum parameter.

PiperOrigin-RevId: 762130610

Author: haozha111

ai_edge_torch/generative/layers/builder.py

@@ -15,9 +15,9 @@
 # Builder class for individual components.
 from typing import Callable
 
+from ai_edge_torch.generative.layers import normalization
 import ai_edge_torch.generative.layers.feed_forward as feed_forward
 import ai_edge_torch.generative.layers.model_config as cfg
-import ai_edge_torch.generative.layers.normalization as normalization
 import torch
 from torch import nn
 import torch.nn.functional as F
@@ -74,6 +74,8 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
         dim,
         eps=config.epsilon,
         zero_centered_gamma=config.zero_centered,
+        with_scale=config.with_scale,
+        scale_shift=config.scale_shift,
         enable_hlfb=config.enable_hlfb,
     )
   elif config.type == cfg.NormalizationType.LAYER_NORM:
@@ -107,20 +109,13 @@ def build_ff(dim: int, config: cfg.FeedForwardConfig):
   else:
     raise ValueError("Unsupported feedforward type.")
 
-  activation = get_activation(config.activation)
-
   pre_ff_norm = build_norm(dim, config.pre_ff_norm_config)
   post_ff_norm = build_norm(dim, config.post_ff_norm_config)
 
   return ff_module(
       dim=dim,
-      hidden_dim=config.intermediate_size,
-      activation=activation,
-      use_bias=config.use_bias,
-      use_glu=(
-          config.activation.type == cfg.ActivationType.GE_GLU
-          or config.activation.type == cfg.ActivationType.SILU_GLU
-      ),
+      activation=get_activation(config.activation),
+      config=config,
       pre_ff_norm=pre_ff_norm,
       post_ff_norm=post_ff_norm,
   )

ai_edge_torch/generative/layers/feed_forward.py

@@ -14,45 +14,69 @@
 # ==============================================================================
 # Common building blocks for FeedForward layers.
 
-from typing import Callable, Optional
+import abc
+from typing import Callable
 
+import ai_edge_torch.generative.layers.model_config as cfg
 import torch
 from torch import nn
 
 
-class SequentialFeedForward(nn.Module):
+class FeedForwardBase(nn.Module):
+  """Base class for feedforward layer."""
+
+  def __init__(
+      self,
+      dim: int,
+      activation: Callable[[torch.Tensor], torch.Tensor],
+      config: cfg.FeedForwardConfig,
+      pre_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
+      post_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
+  ):
+    super().__init__()
+    self.dim = dim
+    self.act = activation
+    self.config = config
+    self.hidden_dim = config.intermediate_size
+    self.use_bias = config.use_bias
+    self.use_glu = (
+        config.activation.type == cfg.ActivationType.GE_GLU
+        or config.activation.type == cfg.ActivationType.SILU_GLU
+    )
+    self.pre_ff_norm = pre_ff_norm
+    self.post_ff_norm = post_ff_norm
+
+  @abc.abstractmethod
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    raise NotImplementedError()
+
+
+class SequentialFeedForward(FeedForwardBase):
   """Vanilla sequential Feedforward with customizable activation."""
 
   def __init__(
       self,
       dim: int,
-      hidden_dim: int,
       activation: Callable[[torch.Tensor], torch.Tensor],
-      use_bias=False,
-      use_glu=False,
-      pre_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
-      post_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+      config: cfg.FeedForwardConfig,
+      pre_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
+      post_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
   ):
     """Init function for feedforward layer.
 
     Args:
       dim (int): embedding size.
-      hidden_dim (int): hidden dim size of the feedforward layer.
       activation (Callable): activation function used in this block.
-      use_bias (Boolean): whether to use bias. Default is false.
-      use_glu (Boolean): whether to use glu in activation. Default is false.
-      pre_ff_norm (Callable): pre feedforward norm. Default is None.
-      post_ff_norm (Callable): post feedforward norm. Default is None.
+      config (cfg.FeedForwardConfig): feedforward layer configuration.
+      pre_ff_norm (Callable): pre feedforward norm. Default is identity.
+      post_ff_norm (Callable): post feedforward norm. Default is identity.
     """
-    super().__init__()
-    self.act = activation
-    if use_glu:
-      self.w1 = nn.Linear(dim, hidden_dim * 2, bias=use_bias)
+    super().__init__(dim, activation, config, pre_ff_norm, post_ff_norm)
+    if self.use_glu:
+      self.w1 = nn.Linear(dim, self.hidden_dim * 2, bias=self.use_bias)
     else:
-      self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
-    self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
-    self.pre_ff_norm = pre_ff_norm if pre_ff_norm else lambda x: x
-    self.post_ff_norm = post_ff_norm if post_ff_norm else lambda x: x
+      self.w1 = nn.Linear(dim, self.hidden_dim, bias=self.use_bias)
+    self.w2 = nn.Linear(self.hidden_dim, dim, bias=self.use_bias)
 
   def forward(self, x):
     """Forward pass for Feedforward layer.
@@ -68,7 +92,7 @@ def forward(self, x):
     return self.post_ff_norm(out)
 
 
-class GatedFeedForward(nn.Module):
+class GatedFeedForward(FeedForwardBase):
   """Gated Feedforward with customizable activation.
 
   https://arxiv.org/pdf/2002.05202v1.pdf
@@ -77,34 +101,48 @@ class GatedFeedForward(nn.Module):
   def __init__(
       self,
       dim: int,
-      hidden_dim: int,
       activation: Callable[[torch.Tensor], torch.Tensor],
-      use_bias=False,
-      use_glu=False,
-      pre_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
-      post_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+      config: cfg.FeedForwardConfig,
+      pre_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
+      post_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
   ):
     """Init function for feedforward layer.
 
     Args:
       dim (int): embedding size.
-      hidden_dim (int): hidden dim size of the feedforward layer.
       activation (Callable): activation function used in this block.
-      use_bias (Boolean): whether to use bias. Default is false.
-      use_glu (Boolean): whether to use glu in activation. Default is false.
-      pre_ff_norm (Callable): pre feedforward norm. Default is None.
-      post_ff_norm (Callable): post feedforward norm. Default is None.
+      pre_ff_norm (Callable): pre feedforward norm. Default is identity.
+      post_ff_norm (Callable): post feedforward norm. Default is identity.
+      config (cfg.FeedForwardConfig): feedforward layer configuration.
     """
-    super().__init__()
-    self.act = activation
-    if use_glu:
-      self.w1 = nn.Linear(dim, hidden_dim * 2, bias=use_bias)
+    super().__init__(dim, activation, config, pre_ff_norm, post_ff_norm)
+
+    if self.use_glu:
+      assert (
+          self.config.use_separate_gating
+      ), 'use_separate_gating must be True for GE_GLU | SILU_GLU activation.'
+
+    if self.config.use_separate_gating:
+      if self.use_glu:
+        self.w1 = nn.Linear(dim, self.hidden_dim * 2, bias=self.use_bias)
+      else:
+        self.w1 = nn.Linear(dim, self.hidden_dim, bias=self.use_bias)
+      self.w3 = nn.Linear(dim, self.hidden_dim, bias=self.use_bias)
     else:
-      self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
-    self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
-    self.w3 = nn.Linear(dim, hidden_dim, bias=use_bias)
-    self.pre_ff_norm = pre_ff_norm if pre_ff_norm else lambda x: x
-    self.post_ff_norm = post_ff_norm if post_ff_norm else lambda x: x
+      self.w_gating = nn.Parameter(
+          torch.ones((2, dim, self.hidden_dim), dtype=torch.float32),
+          requires_grad=False,
+      )
+      self.gating_bias = (
+          nn.Parameter(
+              torch.zeros((2, self.hidden_dim), dtype=torch.float32),
+              requires_grad=False,
+          )
+          if self.use_bias
+          else torch.zeros((2, self.hidden_dim), dtype=torch.float32)
+      )
+
+    self.w2 = nn.Linear(self.hidden_dim, dim, bias=self.use_bias)
 
   def forward(self, x):
     """Forward pass for Feedforward layer.
@@ -116,5 +154,12 @@ def forward(self, x):
       torch.Tensor: output tensor after feedforward.
     """
     x_norm = self.pre_ff_norm(x)
-    out = self.w2(self.act(self.w1(x_norm)) * self.w3(x_norm))
+    if self.config.use_separate_gating:
+      out = self.w2(self.act(self.w1(x_norm)) * self.w3(x_norm))
+    else:
+      out = self.w2(
+          self.act(torch.matmul(x_norm, self.w_gating[0]) + self.gating_bias[0])
+          * (torch.matmul(x_norm, self.w_gating[1]) + self.gating_bias[1])
+      )
+
     return self.post_ff_norm(out)

ai_edge_torch/generative/layers/feed_forward_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 from ai_edge_torch.generative.layers import feed_forward
+from ai_edge_torch.generative.layers import model_config as cfg
 import torch
 import torch.nn.functional as F
 from absl.testing import absltest as googletest
@@ -22,28 +23,32 @@
 class FeedForwardTest(googletest.TestCase):
 
   def test_sequential_feed_forward(self):
+    ff_config = cfg.FeedForwardConfig(
+        type=cfg.FeedForwardType.SEQUENTIAL,
+        activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+        intermediate_size=10,
+        use_bias=True,
+    )
     ff = feed_forward.SequentialFeedForward(
         dim=10,
-        hidden_dim=10,
         activation=F.silu,
-        use_bias=True,
-        use_glu=False,
-        pre_ff_norm=torch.nn.Identity(),
-        post_ff_norm=torch.nn.Identity(),
+        config=ff_config,
     )
     x = torch.ones((1, 10))
     out = ff(x)
     self.assertEqual(out.shape, (1, 10))
 
   def test_gated_feed_forward(self):
+    ff_config = cfg.FeedForwardConfig(
+        type=cfg.FeedForwardType.GATED,
+        activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+        intermediate_size=10,
+        use_bias=True,
+    )
     ff = feed_forward.GatedFeedForward(
         dim=10,
-        hidden_dim=10,
         activation=F.silu,
-        use_bias=True,
-        use_glu=False,
-        pre_ff_norm=torch.nn.Identity(),
-        post_ff_norm=torch.nn.Identity(),
+        config=ff_config,
     )
     x = torch.ones((1, 10))
     out = ff(x)