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

ai_edge_torch/generative/layers/builder.py

@@ -0,0 +1,160 @@
+# 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 typing import Callable
+
+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
+
+
+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)
+
+
+class SwiGLU(nn.Module):
+  """SwiGLU is an activation function which is a variant of GLU.
+
+  SwiGLU is same as SiLU_GLU, because The SiLU function is also known as the
+  swish function.
+
+  SwiGLU(x) = Swish(xW+b) * (xV+c)
+  See: https://paperswithcode.com/method/swiglu
+  """
+
+  def forward(self, x: torch.Tensor):
+    x, y = x.chunk(2, dim=-1)
+    return F.silu(x) * y
+
+
+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,
+        enable_hlfb=config.enable_hlfb,
+    )
+  elif config.type == cfg.NormalizationType.LAYER_NORM:
+    return normalization.LayerNorm(dim, config.epsilon, config.enable_hlfb)
+  elif config.type == cfg.NormalizationType.GROUP_NORM:
+    return normalization.GroupNorm(
+        config.group_num, dim, config.epsilon, config.enable_hlfb
+    )
+  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 (`FeedForwardConfig` 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)
+
+  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
+      ),
+      pre_ff_norm=pre_ff_norm,
+      post_ff_norm=post_ff_norm,
+  )
+
+
+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
+  elif config.type == cfg.ActivationType.SILU_GLU:
+    return SwiGLU()
+  else:
+    raise ValueError("Unsupported activation type.")

ai_edge_torch/generative/layers/feed_forward.py

@@ -0,0 +1,120 @@
+# 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, Optional
+
+import torch
+from torch import nn
+
+
+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,
+      use_glu=False,
+      pre_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+      post_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+  ):
+    """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.
+    """
+    super().__init__()
+    self.act = activation
+    if use_glu:
+      self.w1 = nn.Linear(dim, hidden_dim * 2, bias=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
+
+  def forward(self, x):
+    """Forward pass for Feedforward layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+
+    Returns:
+      torch.Tensor: output tensor after feedforward.
+    """
+    x_norm = self.pre_ff_norm(x)
+    out = self.w2(self.act(self.w1(x_norm)))
+    return self.post_ff_norm(out)
+
+
+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,
+      use_glu=False,
+      pre_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+      post_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+  ):
+    """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.
+    """
+    super().__init__()
+    self.act = activation
+    if use_glu:
+      self.w1 = nn.Linear(dim, hidden_dim * 2, bias=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
+
+  def forward(self, x):
+    """Forward pass for Feedforward layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+
+    Returns:
+      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))
+    return self.post_ff_norm(out)

ai_edge_torch/generative/layers/kv_cache.py

@@ -0,0 +1,204 @@
+# 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.
+# ==============================================================================
+
+"""Utility functions for externalized KV Cache."""
+
+import dataclasses
+from typing import List, Tuple
+
+from ai_edge_torch import hlfb
+from ai_edge_torch.generative.layers import model_config
+from ai_edge_torch.generative.utilities.dynamic_update_slice import dynamic_update_slice
+import torch
+import torch.utils._pytree as pytree
+
+BATCH_SIZE = 1
+
+
+@dataclasses.dataclass
+class KVCacheEntry:
+  """A single cache entry that includes K and V caches.
+
+  The chaches are built based on the provided config with the shape of
+  (batch_size=1, kv_cache_max, num_query_groups, head_dim).
+  """
+
+  k_cache: torch.Tensor
+  v_cache: torch.Tensor
+
+  @classmethod
+  def from_model_config(
+      cls,
+      kv_cache_max: int,
+      config: model_config.AttentionConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCacheEntry":
+    """Build an instance of the class based on model config."""
+    shape = (BATCH_SIZE, kv_cache_max, config.num_query_groups, config.head_dim)
+    k = torch.zeros(shape, dtype=dtype, device=device)
+    v = torch.zeros(shape, dtype=dtype, device=device)
+    obj = cls(k_cache=k, v_cache=v)
+    return obj
+
+
+@dataclasses.dataclass
+class KVCache:
+  """A utility class for holding KV cache entries per layer."""
+
+  caches: Tuple[KVCacheEntry, ...]
+
+  @classmethod
+  def from_model_config(
+      cls,
+      config: model_config.ModelConfig,
+      dtype: torch.dtype = torch.float32,
+      device: torch.device = None,
+  ) -> "KVCache":
+    """Build an instance of the class based on model config.
+
+    Args:
+        config (ModelConfig): Model config used for building the cache.
+        dtype (torch.dtype, optional): The data type of the cache tensor.
+          Defaults to torch.float32.
+        device (torch.device, optional): The device placement of the cache
+          tensors. Defaults to None.
+
+    Returns:
+        KVCache: The created cache object.
+    """
+    caches = [
+        KVCacheEntry.from_model_config(
+            config.kv_cache_max,
+            config.block_config(idx).attn_config,
+            dtype,
+            device,
+        )
+        for idx in range(config.num_layers)
+    ]
+    obj = cls(caches=tuple(caches))
+    return obj
+
+  def flatten(self) -> List[torch.Tensor]:
+    """Flatten the cache entries into a list of tensors with order k_i, v_i."""
+    flattened, _ = _flatten_kvc(self)
+    return flattened
+
+
+def _flatten_kvc(kvc: KVCache) -> Tuple[List[str], List[str]]:
+  flattened = []
+  flat_names = []
+  none_names = []
+  for i, kv_entry in enumerate(kvc.caches):
+    flattened.append(kv_entry.k_cache)
+    flat_names.append(f"k_{i}")
+    flattened.append(kv_entry.v_cache)
+    flat_names.append(f"v_{i}")
+  return flattened, [flat_names, none_names]
+
+
+def _flatten_kvc_with_keys(kvc: KVCache) -> Tuple[List, List]:
+  flattened, (flat_names, none_names) = _flatten_kvc(kvc)
+  return [
+      (pytree.MappingKey(k), v) for k, v in zip(flat_names, flattened)
+  ], flat_names
+
+
+def _unflatten_kvc(
+    values: List[torch.Tensor], context: Tuple[List, List]
+) -> KVCache:
+  assert len(values) % 2 == 0, "Found odd number of K and V entries."
+  num_layers = len(values) // 2
+  flat_names = context[0]
+  kv_entries = []
+  for i in range(num_layers):
+    k_cache_idx = flat_names.index(f"k_{i}")
+    v_cache_idx = flat_names.index(f"v_{i}")
+    kv_entries.append(
+        KVCacheEntry(k_cache=values[k_cache_idx], v_cache=values[v_cache_idx])
+    )
+  obj = KVCache(tuple(kv_entries))
+  return obj
+
+
+pytree.register_pytree_node(
+    KVCache,
+    _flatten_kvc,
+    _unflatten_kvc,
+    flatten_with_keys_fn=_flatten_kvc_with_keys,
+    serialized_type_name="",
+)
+
+
+def update(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+    use_dus: bool = True,
+) -> KVCacheEntry:
+  """Out of place update of Cache buffer.
+
+  Args:
+      cache (KVCacheEntry): The original cache buffer.
+      input_pos (torch.Tensor): The update slice positions.
+      k_slice (torch.Tensor): The K slice to be updated in the new cache.
+      v_slice (torch.Tensor): The V slice to be updated in the new cache.
+
+  Returns:
+      KVCacheEntry: The updated KVCache entry based on the passed inputs.
+  """
+  update_kv_cache = _update_kv_impl if use_dus else _update_kv_base_impl
+  return update_kv_cache(cache, input_pos, k_slice, v_slice)
+
+
+def _update_kv_base_impl(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntry:
+  """Update the cache buffer without High Level Function Boundary annotation."""
+  k = cache.k_cache.index_copy(1, input_pos.to(torch.long), k_slice)
+  v = cache.v_cache.index_copy(1, input_pos.to(torch.long), v_slice)
+  updated_cache = KVCacheEntry(k, v)
+  return updated_cache
+
+
+def _get_slice_indices(positions: torch.Tensor) -> torch.Tensor:
+  """Dynamic Update Slice updates are a variadic sequence of 0-rank tensors."""
+
+  zero = torch.zeros([]).int()
+  positions = positions.int()[0].reshape([])
+  return [zero, positions, zero, zero]
+
+
+def _update_kv_impl(
+    cache: KVCacheEntry,
+    input_pos: torch.Tensor,
+    k_slice: torch.Tensor,
+    v_slice: torch.Tensor,
+) -> KVCacheEntry:
+  """Update the cache buffer for K and V caches."""
+  # NB: Here assume that input_pos == range(input_pos[0], len(input_pos))
+
+  k_slice_indices = _get_slice_indices(input_pos)
+  v_slice_indices = _get_slice_indices(input_pos)
+
+  k = dynamic_update_slice(cache.k_cache, k_slice, k_slice_indices)
+  v = dynamic_update_slice(cache.v_cache, v_slice, v_slice_indices)
+
+  updated_cache = KVCacheEntry(k, v)
+  return updated_cache