ai-edge-torch-nightly 0.1.dev202405131930__py3-none-any.whl

ai_edge_torch/generative/layers/builder.py

@@ -0,0 +1,103 @@
+# 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.
+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
+
+
+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)
+  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(type_: cfg.ActivationType):
+  """Get pytorch callable activation from the name.
+
+  Args:
+    name (string): activation's name.
+
+  Returns:
+    Activation function.
+
+  Raises:
+    ValueError: If activation name is not supported.
+  """
+  if type_ == cfg.ActivationType.SILU:
+    return F.silu
+  elif type_ == cfg.ActivationType.GELU:
+    return F.gelu
+  elif type_ == cfg.ActivationType.GELU_TANH:
+    return lambda x: F.gelu(x, approximate="tanh")
+  elif 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,135 @@
+# 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()
+  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()
+
+
+@enum.unique
+class FeedForwardType(enum.Enum):
+  """Different variations of the Feed Forward module."""
+
+  # `output = linear(act(linear(x)))`.
+  SEQUENTIAL = enum.auto()
+  # `output = linear(act(linear(x)) * lienar(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 use bias with Query, Key, and Value projection.
+  qkv_use_bias: bool = False
+  # 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 FeedForwardConfig:
+  """FeedForward module's parameters."""
+
+  type: FeedForwardType
+  activation: ActivationType
+  intermediate_size: int
+  use_bias: bool = False
+
+
+@dataclass
+class NormalizationConfig:
+  """Normalizater parameters."""
+
+  type: NormalizationType = NormalizationType.NONE
+  epsilon: float = 1e-5
+  zero_centered: bool = False
+
+
+@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
+
+  @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

ai_edge_torch/generative/layers/rotary_position_embedding.py

@@ -0,0 +1,36 @@
+# 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
+import torch
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+  """Computes rotary positional embedding.
+
+  Args:
+    x(torch.Tensor): the input tensor.
+    cos(torch.Tensor): cosine value for the rope.
+    sin(torch.Tensor): sin value for the rope.
+
+  Returns:
+    output tensor of RoPE.
+  """
+  x = x.transpose(1, 2)
+  head_size = x.size(-1)
+  x1 = x[..., : head_size // 2]  # (B, nh, T, hs/2)
+  x2 = x[..., head_size // 2 :]  # (B, nh, T, hs/2)
+  rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
+  roped = (x * cos) + (rotated * sin)
+  return roped.transpose(1, 2).type_as(x)

ai_edge_torch/generative/quantize/__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.
+# ==============================================================================

ai_edge_torch/generative/quantize/example.py

@@ -0,0 +1,45 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+import numpy as np
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.quantize import quant_recipes
+
+
+def main():
+  # Build a PyTorch model as usual
+  config = gemma.get_fake_model_config_2b_for_test()
+  model = gemma.Gemma(config)
+  idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
+  tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+  tokens[0, :4] = idx
+  input_pos = torch.arange(0, 10)
+
+  # Create a quantization recipe to be applied to the model
+  quant_config = quant_recipes.full_linear_int8_dynamic_recipe()
+  print(quant_config)
+
+  # Convert with quantization
+  edge_model = ai_edge_torch.convert(
+      model, (tokens, input_pos), quant_config=quant_config
+  )
+  edge_model.export("/tmp/gemma_2b_quantized.tflite")
+
+
+if __name__ == "__main__":
+  main()

ai_edge_torch/generative/quantize/quant_attrs.py

@@ -0,0 +1,66 @@
+# Copyright 2024 The AI Edge Torch Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+import enum
+
+
+@enum.unique
+class Dtype(enum.Enum):
+  """Data types and precision of tensors."""
+
+  FP32 = enum.auto()
+  FP16 = enum.auto()
+  INT8 = enum.auto()
+
+
+@enum.unique
+class Algorithm(enum.Enum):
+  """Algorithm used to calculate quantization parameters.
+
+  Attributes:
+    MIN_MAX: Maps the min/max of floating point space to the min/max of
+      quantized space and quantize uniformly.
+  """
+
+  MIN_MAX = enum.auto()
+
+
+@enum.unique
+class Mode(enum.Enum):
+  """Mode of quantization.
+
+  Attributes:
+    DYNAMIC_RANGE: Quantize activations during runtime and weights statically to
+      perform computation in integers.
+    WEIGHT_ONLY: Quantize weights statically and dequantize during runtime to
+      perform computation in floating points.
+  """
+
+  DYNAMIC_RANGE = enum.auto()
+  WEIGHT_ONLY = enum.auto()
+
+
+@enum.unique
+class Granularity(enum.Enum):
+  """Granularity of quantization parameters.
+
+  Attributes:
+    NONE: Granularity not applicable to this quantization scheme.
+    CHANNELWISE: Or per-channel quantization. Each channel of relevant tensors
+      is quantized independently of one another.
+  """
+
+  NONE = enum.auto()
+  CHANNELWISE = enum.auto()