ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl

ai_edge_torch/generative/layers/attention_utils.py

@@ -0,0 +1,169 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common utility functions used with attention module.
+
+import math
+from typing import Tuple
+
+import torch
+
+
+def build_rope_cache(
+    size: int,
+    dim: int,
+    base: int = 10000,
+    condense_ratio: int = 1,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Precompute Rotary Positional Embedding Sin and Cos values for quick lookups
+  during the inference.
+
+  Args:
+      size (int): The size of the built cache.
+      dim (int): Each sequence's dimmension.
+      base (int, optional): Rope base value. Defaults to 10000.
+      condense_ratio (int, optional): The ratio by which sequence indicies are
+        condensed. Defaults to 1.
+      dtype (torch.dtype, optional): Output tensor's data type. Defaults to
+        torch.float32.
+      device (torch.device, optional): Output tensor's data type. Defaults to
+        None in which case "cpu" is used.
+
+  Returns:
+      Tuple[torch.Tensor, torch.Tensor]: Rope's Cosine and Sine waves.
+  """
+  if device is None:
+    device = torch.device('cpu')
+  theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+  seq_idx = torch.arange(size) / condense_ratio
+  idx_theta = torch.outer(seq_idx, theta)
+  cos = torch.cos(idx_theta).to(dtype=dtype, device=device)
+  sin = torch.sin(idx_theta).to(dtype=dtype, device=device)
+  return cos, sin
+
+
+def build_causal_mask_cache(
+    size: int,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device = None,
+) -> torch.Tensor:
+  """Build a cache for causal attention mask.
+
+  Args:
+      size (int): The size of the built mask cache.
+      dtype (torch.dtype, optional): Output tensor's data type. Defaults to
+        torch.float32.
+      device (torch.device, optional): Output tensor's data type. Defaults to
+        None in which case "cpu" is used.
+
+  Returns:
+      torch.Tensor: Causal attention mask.
+  """
+  if device is None:
+    device = torch.device('cpu')
+  mask = torch.full((size, size), float('-inf'), dtype=dtype, device=device)
+  return torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+
+
+def relative_position_bucket(
+    relative_position: torch.Tensor,
+    bidirectional: bool,
+    num_buckets: int,
+    max_distance: int,
+) -> torch.Tensor:
+  """
+  Adapted from Mesh Tensorflow:
+  https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
+
+  Translate relative position to a bucket number for relative attention. The relative position is defined as
+  memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to
+  position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for
+  small absolute relative_position and larger buckets for larger absolute relative_positions. All relative
+  positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.
+  This should allow for more graceful generalization to longer sequences than the model has been trained on
+
+  Args:
+      relative_position: an int32 Tensor
+      bidirectional: a boolean - whether the attention is bidirectional
+      num_buckets: an integer for number of buckets.
+      max_distance: an integer for max distance.
+
+  Returns:
+      a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
+  """
+  relative_buckets = 0
+  if bidirectional:
+    num_buckets //= 2
+    relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+    relative_position = torch.abs(relative_position)
+  else:
+    relative_position = -torch.min(
+        relative_position, torch.zeros_like(relative_position)
+    )
+  # now relative_position is in the range [0, inf)
+
+  # half of the buckets are for exact increments in positions
+  max_exact = num_buckets // 2
+  is_small = relative_position < max_exact
+
+  # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+  relative_position_if_large = max_exact + (
+      torch.log(relative_position.float() / max_exact)
+      / math.log(max_distance / max_exact)
+      * (num_buckets - max_exact)
+  ).to(torch.long)
+  relative_position_if_large = torch.min(
+      relative_position_if_large,
+      torch.full_like(relative_position_if_large, num_buckets - 1),
+  )
+
+  relative_buckets += torch.where(
+      is_small, relative_position, relative_position_if_large
+  )
+  return relative_buckets
+
+
+def build_relative_position_buckets(
+    query_length: int,
+    key_length: int,
+    bidirectional: bool = True,
+    num_buckets: int = 32,
+    max_distance: int = 128,
+) -> torch.Tensor:
+  """Relative position buckets for computing bias.
+
+  Args:
+    query_length: an integer of length of current query tensor.
+    key_length: an integer of length of current key tensor.
+    bidirectional: a boolean - whether the attention is bidirectional, default is True.
+    num_buckets: an integer for number of buckets, default is 32.
+    max_distance: an integer for max distance, default is 128.
+
+  Returns:
+    A torch.Tensor of computed relative position buckets.
+  """
+  context_position = torch.arange(query_length, dtype=torch.long)[:, None]
+  memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
+  relative_position = (
+      memory_position - context_position
+  )  # shape (query_length, key_length)
+  rel_pos_bucket = relative_position_bucket(
+      relative_position,  # shape (query_length, key_length)
+      bidirectional=bidirectional,
+      num_buckets=num_buckets,
+      max_distance=max_distance,
+  )
+  return rel_pos_bucket.unsqueeze(0).unsqueeze(0)

ai_edge_torch/generative/layers/builder.py

@@ -0,0 +1,131 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Builder class for individual components.
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+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
+
+
+class GeGLU(nn.Module):
+  """GeGLU is an activation function which is a variant of GELU.
+
+  GeGLU(x) = (xW+b) * GELU(xV+c)
+  See: https://arxiv.org/abs/2002.05202v1
+
+  """
+
+  def __init__(self, d_in: int, d_out: int):
+    super().__init__()
+    self.proj = nn.Linear(d_in, d_out * 2)
+
+  def forward(self, x: torch.Tensor):
+    x, gate = self.proj(x).chunk(2, dim=-1)
+    return x * F.gelu(gate)
+
+
+def build_norm(dim: int, config: cfg.NormalizationConfig):
+  """Builder function for normalizers.
+
+  Args:
+    dim (int): dimension of the input tensor.
+    config (`NormalizationConfig` object): the normalization configuration.
+
+  Returns:
+    The constructed `nn.Module` normalization layer.
+
+  Raises:
+    ValueError: If config's `layer_norm_type` is not supported.
+  """
+  if config.type == cfg.NormalizationType.NONE:
+    return lambda x: x
+  elif config.type == cfg.NormalizationType.RMS_NORM:
+    return normalization.RMSNorm(
+        dim,
+        eps=config.epsilon,
+        zero_centered_gamma=config.zero_centered,
+    )
+  elif config.type == cfg.NormalizationType.LAYER_NORM:
+    return nn.LayerNorm(dim, eps=config.epsilon)
+  elif config.type == cfg.NormalizationType.GROUP_NORM:
+    return nn.GroupNorm(config.group_num, dim, config.epsilon)
+  else:
+    raise ValueError("Unsupported norm type.")
+
+
+def build_ff(dim: int, config: cfg.FeedForwardConfig):
+  """Builder function for Feed Forward. Supports `Sequential` and `Gated`.
+
+  Args:
+    dim (int): dimension of the input tensor.
+    config (`ModelConfig` object): the model configuration.
+
+  Returns:
+    The constructed `nn.Module` feedforward layer.
+
+  Raises:
+    ValueError: If config's `ff_type` is not supported.
+  """
+  ff_type = config.type
+  if ff_type == cfg.FeedForwardType.SEQUENTIAL:
+    ff_module = feed_forward.SequentialFeedForward
+  elif ff_type == cfg.FeedForwardType.GATED:
+    ff_module = feed_forward.GatedFeedForward
+  else:
+    raise ValueError("Unsupported feedforward type.")
+
+  activation = get_activation(config.activation)
+
+  return ff_module(
+      dim=dim,
+      hidden_dim=config.intermediate_size,
+      activation=activation,
+      use_bias=config.use_bias,
+  )
+
+
+def get_activation(config: cfg.ActivationConfig):
+  """Get pytorch callable activation from the activation config.
+
+  Args:
+    config (cfg.ActivationConfig): activation config.
+
+  Returns:
+    Activation function.
+
+  Raises:
+    ValueError: If activation config is not supported.
+  """
+  if config.type == cfg.ActivationType.LINEAR:
+    return lambda x: x
+  elif config.type == cfg.ActivationType.SILU:
+    return F.silu
+  elif config.type == cfg.ActivationType.GELU:
+    return F.gelu
+  elif config.type == cfg.ActivationType.GELU_TANH:
+    return lambda x: F.gelu(x, approximate="tanh")
+  elif config.type == cfg.ActivationType.GELU_QUICK:
+    # GELU approximation that is fast but somewhat inaccurate.
+    # See: https://github.com/hendrycks/GELUs
+    return lambda x: x * F.sigmoid(1.702 * x)
+  elif config.type == cfg.ActivationType.GE_GLU:
+    return GeGLU(config.dim_in, config.dim_out)
+  elif config.type == cfg.ActivationType.RELU:
+    return F.relu
+  else:
+    raise ValueError("Unsupported activation type.")

ai_edge_torch/generative/layers/feed_forward.py

@@ -0,0 +1,95 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common building blocks for FeedForward layers.
+
+from typing import Callable
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class SequentialFeedForward(nn.Module):
+  """Vanilla sequential Feedforward with customizable activation."""
+
+  def __init__(
+      self,
+      dim: int,
+      hidden_dim: int,
+      activation: Callable[[torch.Tensor], torch.Tensor],
+      use_bias=False,
+  ):
+    """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.
+    """
+    super().__init__()
+    self.act = activation
+    self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
+    self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
+
+  def forward(self, x):
+    """Forward pass for Feedforward layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+
+    Returns:
+      torch.Tensor: output tensor after feedforward.
+    """
+    return self.w2(self.act(self.w1(x)))
+
+
+class GatedFeedForward(nn.Module):
+  """Gated Feedforward with customizable activation.
+
+  https://arxiv.org/pdf/2002.05202v1.pdf
+  """
+
+  def __init__(
+      self,
+      dim: int,
+      hidden_dim: int,
+      activation: Callable[[torch.Tensor], torch.Tensor],
+      use_bias=False,
+  ):
+    """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.
+    """
+    super().__init__()
+    self.act = activation
+    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)
+
+  def forward(self, x):
+    """Forward pass for Feedforward layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+
+    Returns:
+      torch.Tensor: output tensor after feedforward.
+    """
+    return self.w2(self.act(self.w1(x)) * self.w3(x))

ai_edge_torch/generative/layers/kv_cache.py

@@ -0,0 +1,83 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# `nn.Module` which implements a KV cache.
+
+import torch
+from torch import nn
+import torch_xla
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+
+class KVCache(nn.Module):
+
+  def __init__(self, batch_size, kv_cache_max, n_heads, head_dim, enable_hlfb=False):
+    """Initializes the KVCache layer.
+
+    Args:
+      batch_size (int): batch size. Currently only batch size 1 is supported.
+      kv_cache_max (int): the max length of KV cache.
+      n_heads (int): number of kv heads.
+      head_dim (int): the head dimension size.
+      enable_hlfb (bool): whether hlfb is enabled or not.
+    """
+    super().__init__()
+    cache_shape = (batch_size, kv_cache_max, n_heads, head_dim)
+    self.register_buffer("k_cache", torch.zeros(cache_shape), persistent=False)
+    self.register_buffer("v_cache", torch.zeros(cache_shape), persistent=False)
+    self.enable_hlfb = enable_hlfb
+    self.kv_cache_max = kv_cache_max
+
+  def update_cache(self, input_pos, k_val, v_val):
+    """Update an entry in the KV cache.
+
+    Args:
+      input_pos (torch.Tensor): the input position.
+      k_val (torch.Tensor): the new `key` value.
+      v_val (torch.Tensor): the new `value` value.
+
+    Returns:
+      The updated key and value tensor.
+    """
+    if self.enable_hlfb:
+      return self.update_cache_with_hlfb(input_pos, k_val, v_val)
+
+    updated_k = self.k_cache.index_copy_(1, input_pos, k_val)
+    updated_v = self.v_cache.index_copy_(1, input_pos, v_val)
+    # Here we need a clone otherwise dynamo export will fail.
+    return torch.clone(updated_k), torch.clone(updated_v)
+
+  def update_cache_with_hlfb(self, input_pos, k_val, v_val):
+    """Update an entry in the KV cache and enable high-level function boundary.
+
+    Args:
+      input_pos (torch.Tensor): the input position.
+      k_val (torch.Tensor): the new `key` value.
+      v_val (torch.Tensor): the new `value` value.
+
+    Returns:
+      The updated key and value tensor.
+    """
+
+    builder = StableHLOCompositeBuilder(
+        name="odml.update_kv_cache", attr={"kv_cache_max": self.kv_cache_max}
+    )
+    k_cache, v_cache, input_pos, k_val, v_val = builder.mark_inputs(
+        self.k_cache, self.v_cache, input_pos, k_val, v_val
+    )
+    updated_k = k_cache.index_copy_(1, input_pos, k_val)
+    updated_v = v_cache.index_copy_(1, input_pos, v_val)
+    updated_k, updated_v = builder.mark_outputs(updated_k, updated_v)
+    return updated_k, updated_v

ai_edge_torch/generative/layers/model_config.py

@@ -0,0 +1,158 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Model configuration class.
+from dataclasses import dataclass
+from dataclasses import field
+import enum
+from typing import Optional
+
+
+@enum.unique
+class ActivationType(enum.Enum):
+  """Different activation functions supported by the default builder."""
+
+  LINEAR = enum.auto()
+  SILU = enum.auto()
+  GELU = enum.auto()
+  GELU_TANH = enum.auto()
+  GELU_QUICK = enum.auto()
+  GE_GLU = enum.auto()
+  RELU = enum.auto()
+
+
+@enum.unique
+class NormalizationType(enum.Enum):
+  """Different normalization functions"""
+
+  # No normalization is applied.
+  NONE = enum.auto()
+  RMS_NORM = enum.auto()
+  LAYER_NORM = enum.auto()
+  GROUP_NORM = enum.auto()
+
+
+@enum.unique
+class FeedForwardType(enum.Enum):
+  """Different variations of the Feed Forward module."""
+
+  # `output = linear(act(linear(x)))`.
+  SEQUENTIAL = enum.auto()
+  # `output = linear_2(act(linear_1(x)) * lienar_3(x))`.
+  GATED = enum.auto()
+
+
+@dataclass
+class AttentionConfig:
+  """Attention moduel's parameters."""
+
+  num_heads: int
+  # Used to determine number of groups in grouped query attention (GQA)
+  # https://arxiv.org/pdf/2305.13245.pdf
+  num_query_groups: Optional[int]
+  # Percentage of Rotary Positional Embedding added Q and K projections.
+  rotary_percentage: Optional[float] = None
+  # Whether to transpose the query groups of qkv bundled tensor before
+  # splitting into separated tensors.
+  qkv_transpose_before_split: bool = False
+  # Whether to use bias with Query, Key, and Value projection.
+  qkv_use_bias: bool = False
+  # Whether the fused q, k, v projection weights interleaves q, k, v heads.
+  # If True, the projection weights are in format [q_head_0, k_head_0, v_head_0, q_head_1, k_head_1, v_head_1, ...]
+  # If False, the projection weights are in format [q_head_0, q_head_1, ..., k_head_0, k_head_1, ... v_head_0, v_head_1, ...]
+  qkv_fused_interleaved: bool = True
+  # Whether to use bias with attention output projection.
+  output_proj_use_bias: bool = False
+  enable_kv_cache: bool = True
+  relative_attention_num_buckets: int = 0
+  relative_attention_max_distance: int = 0
+
+
+@dataclass
+class ActivationConfig:
+  type: ActivationType = ActivationType.LINEAR
+  # Dimension of input and output, used in GeGLU.
+  dim_in: Optional[int] = None
+  dim_out: Optional[int] = None
+
+
+@dataclass
+class FeedForwardConfig:
+  """FeedForward module's parameters."""
+
+  type: FeedForwardType
+  activation: ActivationConfig
+  intermediate_size: int
+  use_bias: bool = False
+
+
+@dataclass
+class NormalizationConfig:
+  """Normalizater parameters."""
+
+  type: NormalizationType = NormalizationType.NONE
+  epsilon: float = 1e-5
+  zero_centered: bool = False
+  # Number of groups used in group normalization.
+  group_num: Optional[float] = None
+
+
+@dataclass
+class ModelConfig:
+  """Base configurations for building a transformer architecture."""
+
+  vocab_size: int
+  num_layers: int
+  max_seq_len: int
+  embedding_dim: int
+
+  attn_config: AttentionConfig
+  ff_config: FeedForwardConfig
+  # The normalization applied to attention's input.
+  pre_attention_norm_config: NormalizationConfig = field(
+      default_factory=NormalizationConfig
+  )
+  # The normalization applied to feed forward's input.
+  pre_ff_norm_config: NormalizationConfig = field(default_factory=NormalizationConfig)
+  # The normalization applied before LM head.
+  final_norm_config: NormalizationConfig = field(default_factory=NormalizationConfig)
+
+  # If set to True, only pre_attention_norm is applied to the input and the
+  # decode's output is computed as `output = input + attn_out + ff_out` where
+  # attention and feed forward are called with pre_attention_norm's output.
+  parallel_residual: bool = False
+  # Use bias term within LLM's HEAD.
+  lm_head_use_bias: bool = False
+  # Whether to turn on high-level function boundary.
+  enable_hlfb: bool = False
+
+  # The maximum sequence length of the KV cache. Should not exceed max_seq_len.
+  kv_cache_max_len: int = 0
+
+  # The Attention computation will include relative positional bias.
+  relative_attention: bool = False
+
+  # Default batch size of the exported model. Default value is 1.
+  batch_size: int = 1
+
+  @property
+  def kv_cache_max(self) -> int:
+    if self.kv_cache_max_len > 0:
+      return self.kv_cache_max_len
+    else:
+      return self.max_seq_len
+
+  @property
+  def head_dim(self) -> int:
+    return self.embedding_dim // self.attn_config.num_heads

ai_edge_torch/generative/layers/normalization.py

@@ -0,0 +1,62 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Common normalization layers.
+
+import torch
+
+
+# Implementation for RMSNorm from: https://arxiv.org/abs/1910.07467
+class RMSNorm(torch.nn.Module):
+
+  def __init__(self, dim: int, eps: float = 1e-6, zero_centered_gamma=False):
+    """
+    Initialize the RMSNorm layer.
+
+    Args:
+      dim (int): dimension of the input tensor.
+      eps (float): A small float value to ensure numerical stability (default: 1e-6).
+    """
+    super().__init__()
+    self.eps = eps
+    self.weight = torch.nn.Parameter(torch.ones(dim))
+    self.zero_centered_gamma = zero_centered_gamma
+
+  def _norm(self, x):
+    """
+    Apply RMSNorm normalization.
+
+    Args:
+      x (torch.Tensor): input tensor.
+
+    Returns:
+      torch.Tensor: The normalized output tensor.
+    """
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+  def forward(self, x):
+    """
+    Running the forward pass of RMSNorm layer.
+
+    Args:
+      x (torch.Tensor): input tensor.
+
+    Returns:
+      torch.Tensor: output tensor after applying RMSNorm.
+    """
+    output = self._norm(x.float()).type_as(x)
+    if self.zero_centered_gamma:
+      return output * (1 + self.weight)
+    else:
+      return output * self.weight