ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/layers/model_config.py

@@ -0,0 +1,238 @@
+# 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."""
+
+import dataclasses
+import enum
+from typing import Callable, Optional, Sequence, Union
+from ai_edge_torch.generative.layers import rotary_position_embedding
+
+@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()
+  SILU_GLU = 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()
+
+
+class AttentionType(enum.Enum):
+  GLOBAL = enum.auto()
+  LOCAL_SLIDING = enum.auto()
+
+
+@dataclasses.dataclass
+class NormalizationConfig:
+  """Normalizater parameters."""
+
+  type: NormalizationType = NormalizationType.NONE
+  enable_hlfb: bool = False
+  epsilon: float = 1e-5
+  zero_centered: bool = False
+  # Number of groups used in group normalization.
+  group_num: Optional[float] = None
+
+
+@dataclasses.dataclass
+class AttentionConfig:
+  """Attention model's parameters."""
+
+  num_heads: int
+  head_dim: int
+  # Used to determine number of groups in grouped query attention (GQA)
+  # https://arxiv.org/pdf/2305.13245.pdf
+  num_query_groups: Optional[int]
+  # Base of rotary positional embedding.
+  rotary_base: int = 10_000
+  # 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
+  # The normalization applied to query projection's output.
+  query_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+  # The normalization applied to key projection's output.
+  key_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+  relative_attention_num_buckets: int = 0
+  relative_attention_max_distance: int = 0
+  # Softcap on the output logits.
+  logit_softcap: Optional[float] = None
+  # The type of attention.
+  attn_type: Optional[AttentionType] = None
+  # The size of the sliding window used for local attention.
+  sliding_window_size: Optional[int] = None
+
+
+@dataclasses.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
+
+
+@dataclasses.dataclass
+class FeedForwardConfig:
+  """FeedForward module's parameters."""
+
+  type: FeedForwardType
+  activation: ActivationConfig
+  intermediate_size: int
+  use_bias: bool = False
+  # The normalization applied to feed forward's input.
+  pre_ff_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+  # The normalization applied to feed forward's output.
+  post_ff_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+
+
+@dataclasses.dataclass
+class TransformerBlockConfig:
+  """TransformerBlock module's parameters."""
+
+  attn_config: AttentionConfig
+  ff_config: FeedForwardConfig
+  # The normalization applied to attention's input.
+  pre_attention_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+  # The normalization applied to attentions's output.
+  post_attention_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+  # If set to True, only attn_config.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
+  # The Attention computation will include relative positional bias.
+  relative_attention: bool = False
+
+
+@dataclasses.dataclass
+class ImageEmbeddingConfig:
+  """Image embedding parameters."""
+
+  channels: int
+  # All images should be normalized to the size of [image_size * image_size].
+  image_size: int
+  patch_size: int
+
+
+@dataclasses.dataclass
+class ModelConfig:
+  """Base configurations for building a transformer architecture."""
+
+  vocab_size: int
+  num_layers: int
+  max_seq_len: int
+  embedding_dim: int
+
+  # TransformerBlockConfig for each layer block. If a single
+  # TransformerBlockConfig is provided, it will be used for all layers.
+  block_configs: Union[TransformerBlockConfig, Sequence[TransformerBlockConfig]]
+
+  # The normalization applied before LM head.
+  final_norm_config: NormalizationConfig = dataclasses.field(
+      default_factory=NormalizationConfig
+  )
+
+  # Scale factor of the embedding.
+  embedding_scale: Optional[float] = None
+  # Use bias term within embedding.
+  embedding_use_bias: bool = False
+  # Image embedding parameters.
+  image_embedding: Optional[ImageEmbeddingConfig] = None
+
+  # Use bias term within LLM's HEAD.
+  lm_head_use_bias: bool = False
+  # Whether LLM's HEAD shares the weight of the embedding.
+  lm_head_share_weight_with_embedding: bool = True
+
+  # 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
+
+  # Default batch size of the exported model. Default value is 1.
+  batch_size: int = 1
+
+  # Softcap on the model output logits.
+  final_logit_softcap: Optional[float] = None
+
+  # The function to call to create the RoPE sin and cos vectors during the
+  # forward pass. Defaults to a standard implementation.
+  build_rope: Callable = rotary_position_embedding.build_rope
+
+  @property
+  def kv_cache_max(self) -> int:
+    if self.kv_cache_max_len > 0:
+      return self.kv_cache_max_len
+    return self.max_seq_len
+
+  def block_config(self, idx: int) -> TransformerBlockConfig:
+    if isinstance(self.block_configs, TransformerBlockConfig):
+      return self.block_configs
+    if idx < 0 or idx >= len(self.block_configs):
+      raise ValueError(
+          f"Index {idx} is out of range for layer configs: {self.block_configs}"
+      )
+    return self.block_configs[idx]

ai_edge_torch/generative/layers/normalization.py

@@ -0,0 +1,222 @@
+# 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.
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+# 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,
+      enable_hlfb: bool = 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).
+      zero_centered_gamma (bool): Whether or not gamma has an offset.
+      enable_hlfb (bool): use HLFB in the op.
+    """
+    super().__init__()
+    self.enable_hlfb = enable_hlfb
+    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.
+    """
+    if self.zero_centered_gamma:
+      w = 1 + self.weight
+    else:
+      w = self.weight
+
+    if self.enable_hlfb:
+      return rms_norm_with_hlfb(
+          x,
+          w,
+          self.eps,
+      )
+    else:
+      output = self._norm(x.float()).type_as(x)
+      return output * w
+
+
+class GroupNorm(torch.nn.Module):
+
+  def __init__(
+      self,
+      group_num: int,
+      dim: int,
+      eps: float = 1e-5,
+      enable_hlfb: bool = False,
+  ):
+    """Initialize the GroupNorm layer.
+
+    Args:
+      group_num (int): Number of groups to separate the channels into.
+      dim (int): Dimension of the input tensor.
+      eps (float): A small float value to ensure numerical stability (default:
+        1e-5).
+      enable_hlfb (bool): Whether to convert this normalization into a single
+        op.
+    """
+    super().__init__()
+    self.enable_hlfb = enable_hlfb
+    self.group_num = group_num
+    self.eps = eps
+    self.weight = torch.nn.Parameter(torch.empty(dim))
+    self.bias = torch.nn.Parameter(torch.empty(dim))
+
+  def forward(self, x):
+    """Running the forward pass of GroupNorm layer.
+
+    Args:
+      x (torch.Tensor): input tensor.
+
+    Returns:
+      torch.Tensor: output tensor after applying GroupNorm.
+    """
+    return F.group_norm(x, self.group_num, self.weight, self.bias, self.eps)
+
+
+class LayerNorm(torch.nn.Module):
+
+  def __init__(
+      self,
+      dim: int,
+      eps: float = 1e-5,
+      enable_hlfb: bool = False,
+  ):
+    """Initialize the LayerNorm layer.
+
+    Args:
+      dim (int): dimension of the input tensor.
+      eps (float): A small float value to ensure numerical stability (default:
+        1e-5).
+      enable_hlfb (bool): Whether to convert this normalization into a single
+        op.
+    """
+    super().__init__()
+    self.enable_hlfb = enable_hlfb
+    self.normalized_shape = (dim,)
+    self.eps = eps
+    self.weight = torch.nn.Parameter(torch.empty(dim))
+    self.bias = torch.nn.Parameter(torch.empty(dim))
+
+  def forward(self, x):
+    """Running the forward pass of LayerNorm layer.
+
+    Args:
+      x (torch.Tensor): input tensor.
+
+    Returns:
+      torch.Tensor: output tensor after applying LayerNorm.
+    """
+    if self.enable_hlfb:
+      return layer_norm_with_hlfb(
+          x, self.normalized_shape, self.weight, self.bias, self.eps
+      )
+    return F.layer_norm(
+        x, self.normalized_shape, self.weight, self.bias, self.eps
+    )
+
+
+def rms_norm_with_hlfb(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    eps: float,
+):
+  """RMS Normalization with high-level function boundary enabled.
+
+  Args:
+    x (torch.Tensor): Input tensor for RMS Normalization, with BCHW shape.
+    w (torch.Tensor): The learned parameter tensor for normalization.
+    eps (float): A small float value to ensure numerical stability.
+
+  Returns:
+    The output tensor of RMS Normalization.
+  """
+  builder = StableHLOCompositeBuilder(
+      name="odml.rms_norm", attr={"epsilon": eps}
+  )
+
+  x, w = builder.mark_inputs(x, w)
+
+  def _norm(x):
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+  output = _norm(x.float()).type_as(x)
+  out = output * w
+
+  out = builder.mark_outputs(out)
+  return out
+
+
+def layer_norm_with_hlfb(
+    x: torch.Tensor,
+    normalized_shape: list[int],
+    w: torch.Tensor,
+    b: torch.Tensor,
+    eps: float,
+):
+  """Layer Normalization with high-level function boundary enabled.
+
+  Args:
+    x (torch.Tensor): Input tensor for Layer Normalization, with BCHW shape.
+    normalized_shape (list[int]): Input shape from an expected input of size,
+      same as https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html.
+    w (torch.Tensor): The weight tensor for the normalization.
+    b (torch.Tensor): The bias tensor for the normalization.
+    eps (float): A small float value to ensure numerical stability.
+
+  Returns:
+    The output tensor of Layer Normalization.
+  """
+  builder = StableHLOCompositeBuilder(
+      name="odml.group_norm",
+      attr={"num_groups": 1, "epsilon": eps, "channel_axis": 1},
+  )
+  x, w, b = builder.mark_inputs(x, w, b)
+  y = F.layer_norm(x, normalized_shape, w, b, eps=eps)
+  y = builder.mark_outputs(y)
+  return y

ai_edge_torch/generative/layers/rotary_position_embedding.py

@@ -0,0 +1,94 @@
+# 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.
+# ==============================================================================
+# Implementation for Rotary Position embedding. https://arxiv.org/pdf/2104.09864.pdf
+
+from typing import Tuple
+import torch
+
+
+def apply_rope(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+  """Computes rotary positional embedding.
+
+  Args:
+    x: the input tensor.
+    cos: cosine value for the rope.
+    sin: sin value for the rope.
+
+  Returns:
+    output tensor of RoPE.
+  """
+  x = x.transpose(1, 2)
+  rope_size = cos.size(-1)
+  x_splited = torch.split(x, rope_size, dim=-1)
+  left = x_splited[0] * cos - x_splited[1] * sin
+  right = x_splited[1] * cos + x_splited[0] * sin
+  roped = torch.cat((left, right) + x_splited[2:], dim=-1)
+  return roped.transpose(1, 2).type_as(x)
+
+
+def build_rope(
+    input_pos: torch.Tensor,
+    n_elem: int,
+    base: int = 10_000,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Computes rotary positional embedding cosine and sine tensors.
+
+  Args:
+    input_pos: the sequence indices for the query and key
+    n_elem: number of elements of the head dimension for RoPE computation
+    base: the base of the exponentiated value for RoPE.
+
+  Returns:
+    cos, sin tensors
+  """
+
+  if n_elem <= 0:
+    return None, None
+
+  freq_exponents = (2.0 / n_elem) * torch.arange(
+      n_elem // 2, dtype=torch.float32
+  )
+  timescale = float(base) ** freq_exponents
+  radians = input_pos.clone().unsqueeze(0).unsqueeze(-1) / timescale.unsqueeze(
+      0
+  ).unsqueeze(0)
+  cos = torch.cos(radians)
+  sin = torch.sin(radians)
+  return cos, sin
+
+
+def apply_rope_inline(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+  """Computes rotary positional embedding inline for a query and key.
+
+  Args:
+    q: the query tensor.
+    k: the key tensor.
+    cos: the cosine tensor.
+    sin: the sine tensor.
+
+  Returns:
+    output the RoPE'd query and key.
+  """
+
+  q_roped = apply_rope(q, cos, sin)
+  k_roped = apply_rope(k, cos, sin)
+  return q_roped, k_roped

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -0,0 +1,144 @@
+# 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.
+# ==============================================================================
+# Implements scaled dot product attention.
+
+import math
+from typing import Optional
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+import torch
+import torch.nn.functional as F
+
+
+def scaled_dot_product_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    softcap: Optional[float] = None,
+):
+  """Scaled dot product attention.
+
+  Args:
+    q (torch.Tensor): Query tensor, with shape [B, T, N, H].
+    k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
+    v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
+    head_size (int): head dimension.
+    mask (torch.Tensor): the optional mask tensor.
+
+  Returns:
+    The output tensor of scaled_dot_product_attention.
+  """
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  q = q.transpose(1, 2)
+  k = k.transpose(1, 2)
+  v = v.transpose(1, 2)
+  if q.size() != k.size():
+    # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
+    k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
+    v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
+  if softcap is None:
+    y = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=mask,
+        dropout_p=0.0,
+        is_causal=mask is None,
+        scale=scale,
+    )
+  else:
+    q.mul_(scale)
+    scores = q @ k.transpose(-1, -2)
+    scores = scores / softcap
+    scores = torch.tanh(scores)
+    scores = scores * softcap
+    scores = scores + mask
+    out = F.softmax(scores.float(), dim=-1).type_as(q)
+    y = torch.matmul(out, v)
+
+  return y.transpose(1, 2)
+
+
+def scaled_dot_product_attention_with_hlfb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    softcap: Optional[float] = None,
+):
+  """Scaled dot product attention with high-level function boundary enabled.
+
+  Args:
+    q (torch.Tensor): Query tensor, with shape [B, T, N, H].
+    k (torch.Tensor): Key tensor, with shape [B, T, KV_LEN, H].
+    v (torch.Tensor): Value tensor, with shape [B, T, KV_LEN, H].
+    head_size (int): head dimension.
+    mask (torch.Tensor): the optional mask tensor.
+
+  Returns:
+    The output tensor of scaled_dot_product_attention.
+  """
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  attrs = {"scale": scale}
+
+  if softcap is not None:
+    attrs["logit_cap"] = softcap
+
+  builder = StableHLOCompositeBuilder(
+      name="odml.scaled_dot_product_attention", attr=attrs
+  )
+  q, k, v, mask = builder.mark_inputs(q, k, v, mask)
+
+  q = q.transpose(1, 2)
+  k = k.transpose(1, 2)
+  v = v.transpose(1, 2)
+  if q.size() != k.size():
+    # Handle the GQA case, where q.shape[1] % k.shape[1] == 0.
+    k = k.repeat_interleave(q.shape[1] // k.shape[1], dim=1)
+    v = v.repeat_interleave(q.shape[1] // v.shape[1], dim=1)
+  if softcap is None:
+    y = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=mask,
+        dropout_p=0.0,
+        is_causal=mask is None,
+        scale=scale,
+    )
+  else:
+    q.mul_(scale)
+    scores = q @ k.transpose(-1, -2)
+    scores = scores / softcap
+    scores = torch.tanh(scores)
+    scores = scores * softcap
+    scores = scores + mask
+    out = F.softmax(scores.float(), dim=-1).type_as(q)
+    y = torch.matmul(out, v)
+
+  result = y.transpose(1, 2)
+  result = builder.mark_outputs(result)
+  return result

ai_edge_torch/generative/layers/unet/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+# ==============================================================================