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

ai_edge_torch/generative/examples/t5/t5_attention.py

@@ -0,0 +1,255 @@
+# 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.
+# ==============================================================================
+# Attention modules for the T5 encoder-decoder model family.
+
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from ai_edge_torch.generative.layers.attention import scaled_dot_product_attention  # NOQA
+from ai_edge_torch.generative.layers.attention import scaled_dot_product_attention_with_hlfb  # NOQA
+import ai_edge_torch.generative.layers.builder as builder
+from ai_edge_torch.generative.layers.kv_cache import KVCache
+import ai_edge_torch.generative.layers.model_config as cfg
+
+
+class EncoderDecoderBlock(nn.Module):
+
+  def __init__(
+      self, config: cfg.ModelConfig, has_relative_attention_bias: bool = False
+  ) -> None:
+    """Initialize an instance of the EncoderDecoderBlock.
+
+    Args:
+      config (cfg.ModelConfig): the configuration object
+        for this transformer block.
+      has_relative_attention_bias (bool): whether the
+        self attention block has relative bias.
+    """
+
+    super().__init__()
+    self.atten_func = T5Attention(
+        config.embedding_dim,
+        config.attn_config,
+        config.pre_attention_norm_config,
+        config.kv_cache_max,
+        config.enable_hlfb,
+        has_relative_attention_bias=has_relative_attention_bias,
+    )
+    # For a decoder, we add a cross attention.
+    if config.is_decoder:
+      self.cross_atten_func = T5Attention(
+          config.embedding_dim,
+          config.attn_config,
+          config.pre_attention_norm_config,
+          config.kv_cache_max,
+          config.enable_hlfb,
+          # Cross Attention does not have relative attention bias.
+          has_relative_attention_bias=False,
+      )
+    else:
+      self.cross_atten_func = None
+
+    self.pre_ff_norm = builder.build_norm(
+        config.embedding_dim, config.pre_ff_norm_config
+    )
+    self.ff = builder.build_ff(config.embedding_dim, config.ff_config)
+    self.config = config
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      input_pos: Optional[torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      relative_position: Optional[torch.Tensor] = None,
+      position_bias: Optional[torch.Tensor] = None,
+      encoder_hidden_states: Optional[torch.Tensor] = None,
+      encoder_attention_mask: Optional[torch.Tensor] = None,
+      encoder_decoder_position_bias: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the EncoderDecoderBlock.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this transformer block.
+    """
+
+    hidden_states, position_bias = self.atten_func(
+        x,
+        input_pos=input_pos,
+        mask=mask,
+        relative_position=relative_position,
+        position_bias=position_bias,
+    )
+
+    attn_out = hidden_states + x
+
+    if self.cross_atten_func:
+      hidden_states, encoder_decoder_position_bias = self.cross_atten_func(
+          attn_out,
+          input_pos=input_pos,
+          key_value_states=encoder_hidden_states,
+          mask=encoder_attention_mask,
+          relative_position=relative_position,
+          position_bias=encoder_decoder_position_bias,
+      )
+      attn_out = hidden_states + attn_out
+
+    forwarded = self.pre_ff_norm(attn_out)
+    forwarded = self.ff(forwarded)
+    hidden_states = attn_out + forwarded
+
+    # encoder_deocder_position_bias is from CrossAttention
+    return hidden_states, position_bias, encoder_decoder_position_bias
+
+
+class T5Attention(nn.Module):
+
+  def __init__(
+      self,
+      dim: int,
+      config: cfg.AttentionConfig,
+      norm_config: cfg.NormalizationConfig,
+      kv_cache_max: int,
+      enable_hlfb: bool,
+      has_relative_attention_bias=False,
+  ) -> None:
+    """Initialize an instance of T5Attention.
+
+    Args:
+      dim (int): causal attention's input/output dimmension.
+      config (cfg.AttentionConfig): attention specific configurations.
+      kv_cache_max (int): determines the size of the KV Cache buffer, if enabled.
+      enable_hlfb (bool): whether hlfb is enabled or not.
+      has_relative_attention_bias (bool): whether we compute relative bias.
+    """
+    super().__init__()
+    self.pre_atten_norm = builder.build_norm(dim, norm_config)
+
+    self.has_relative_attention_bias = has_relative_attention_bias
+    self.relative_attention_num_buckets = config.relative_attention_num_buckets
+    self.d_model = dim
+    self.head_dim = dim // config.num_heads
+    self.n_heads = config.num_heads
+    self.inner_dim = self.n_heads * self.head_dim
+
+    self.q = nn.Linear(self.d_model, self.inner_dim, bias=config.qkv_use_bias)
+    self.k = nn.Linear(self.d_model, self.inner_dim, bias=config.qkv_use_bias)
+    self.v = nn.Linear(self.d_model, self.inner_dim, bias=config.qkv_use_bias)
+    # output projection
+    self.proj = nn.Linear(
+        self.inner_dim, self.d_model, bias=config.output_proj_use_bias
+    )
+
+    if self.has_relative_attention_bias:
+      self.relative_attention_bias = nn.Embedding(
+          self.relative_attention_num_buckets, self.n_heads
+      )
+
+    self.config = config
+    self.kv_cache = None
+    # Build a k/v cache with size (batch_size, kv_cache_max, n_heads, head_dim).
+    # Now only supports a max batch_size of 1.
+    if config.enable_kv_cache:
+      self.kv_cache = KVCache(
+          1,
+          kv_cache_max,
+          config.num_query_groups,
+          self.head_dim,
+          enable_hlfb,
+      )
+
+    if enable_hlfb:
+      self.sdpa_func = scaled_dot_product_attention_with_hlfb
+    else:
+      self.sdpa_func = scaled_dot_product_attention
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      input_pos: Optional[torch.Tensor] = None,
+      key_value_states: Optional[torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      relative_position: Optional[torch.Tensor] = None,
+      position_bias: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the T5Attention layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this self attention layer.
+    """
+
+    x = self.pre_atten_norm(x)
+    B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+    query_states = self.q(x)
+    query_states = query_states.reshape(B, T, -1, self.head_dim)  # (B, T, nh_q, hs)
+
+    if key_value_states is not None:
+      (
+          kvB,
+          kvT,
+          kvC,
+      ) = (
+          key_value_states.size()
+      )  # batch size, sequence length, embedding dimensionality (n_embd)
+      key_states = self.k(key_value_states)
+      value_states = self.v(key_value_states)
+      key_states = key_states.reshape(kvB, kvT, -1, self.head_dim)
+      value_states = value_states.reshape(kvB, kvT, -1, self.head_dim)
+    else:
+      key_states = self.k(x)
+      value_states = self.v(x)
+      key_states = key_states.reshape(B, T, -1, self.head_dim)
+      value_states = value_states.reshape(B, T, -1, self.head_dim)
+
+    if key_value_states is None and self.kv_cache is not None:
+      key_states, value_states = self.kv_cache.update_cache(
+          input_pos, key_states, value_states
+      )
+
+    if position_bias is None:
+      # handle the encoder case first
+      if self.has_relative_attention_bias:
+        position_bias = self.relative_attention_bias(
+            relative_position
+        )  # shape (query_length, key_length, num_heads)
+        position_bias = position_bias.permute([0, 1, 4, 2, 3]).squeeze(
+            0
+        )  # shape (1, num_heads, query_length, key_length)
+      else:
+        # position_bias = torch.zeros(B, self.n_heads, T, self.head_dim, dtype=torch.float32)
+        position_bias = torch.zeros_like(mask, dtype=torch.float32)
+
+    mask = mask + position_bias
+    y = self.sdpa_func(
+        query_states, key_states, value_states, self.head_dim, mask=mask, scale=1.0
+    )
+    y = y.reshape(B, T, C)  # re-assemble all head outputs side by side
+    # output projection
+    y = self.proj(y)
+    return y, position_bias

ai_edge_torch/generative/examples/test_models/__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/examples/test_models/toy_model.py

@@ -0,0 +1,119 @@
+# 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.
+# ==============================================================================
+# A toy example which has a single-layer transformer block.
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ai_edge_torch
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+KV_CACHE_MAX_LEN = 100
+
+
+class ToySingleLayerModel(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_block = TransformerBlock(config)
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        base=10_000,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device('cpu'),
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+
+    x = self.transformer_block(x, (cos, sin), mask, input_pos)
+    x = self.final_norm(x)
+    return self.lm_head(x)
+
+
+def define_and_run() -> None:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32, num_query_groups=4, rotary_percentage=1.0, enable_kv_cache=False
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationType.SILU,
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=400,
+      num_layers=1,
+      max_seq_len=KV_CACHE_MAX_LEN,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+  )
+
+  model = ToySingleLayerModel(config)
+  idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN), 0)
+  input_pos = torch.arange(0, KV_CACHE_MAX_LEN)
+  print('running an inference')
+  print(
+      model.forward(
+          idx,
+          input_pos,
+      )
+  )
+
+  # Convert model to tflite.
+  print('converting model to tflite')
+  edge_model = ai_edge_torch.convert(
+      model,
+      (
+          idx,
+          input_pos,
+      ),
+  )
+  edge_model.export('/tmp/toy_model.tflite')
+
+
+if __name__ == '__main__':
+  define_and_run()

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -0,0 +1,143 @@
+# 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.
+# ==============================================================================
+# A toy example which has basic transformer block (w/ KV-Cache).
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_xla
+
+import ai_edge_torch
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class ToyModelWithKV(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_blocks = nn.ModuleList(
+        TransformerBlock(config) for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        base=10_000,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device('cpu'),
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+    for i, block in enumerate(self.transformer_blocks):
+      x = block(x, (cos, sin), mask, input_pos)
+    x = self.final_norm(x)
+    return self.lm_head(x)
+
+
+def _export_stablehlo_mlir(model, args):
+  ep = torch.export.export(model, args)
+  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+  return stablehlo_gm.get_stablehlo_text()
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32, num_query_groups=4, rotary_percentage=1.0
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationType.SILU,
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=150,
+      num_layers=2,
+      max_seq_len=500,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.unsqueeze(torch.arange(0, 100), 0)
+  input_pos = torch.arange(0, 100)
+  return idx, input_pos
+
+
+def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.tensor([[1]], dtype=torch.long)
+  input_pos = torch.tensor([10], dtype=torch.int64)
+  return idx, input_pos
+
+
+def define_and_run() -> None:
+  dump_mlir = False
+
+  config = get_model_config()
+  model = ToyModelWithKV(config)
+  print('running an inference')
+  idx, input_pos = get_sample_prefill_inputs()
+  decode_idx, decode_input_pos = get_sample_decode_inputs()
+  print(model.forward(idx, input_pos))
+
+  if dump_mlir:
+    mlir_text = _export_stablehlo_mlir(model, (idx, input_pos))
+    with open('/tmp/toy_model_with_kv.stablehlo.mlir', 'w') as f:
+      f.write(mlir_text)
+
+  # Convert model to tflite with 2 signatures (prefill + decode).
+  print('converting toy model to tflite with 2 signatures (prefill + decode)')
+  edge_model = (
+      ai_edge_torch.signature('prefill', model, (idx, input_pos))
+      .signature('decode', model, (decode_idx, decode_input_pos))
+      .convert()
+  )
+  edge_model.export('/tmp/toy_kv_cache.tflite')
+
+
+if __name__ == '__main__':
+  define_and_run()

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+
+import os
+from pathlib import Path
+
+import torch
+
+import ai_edge_torch
+from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
+from ai_edge_torch.generative.quantize import quant_recipes
+
+
+def convert_tiny_llama_to_tflite(
+    checkpoint_path: str,
+    prefill_seq_len: int = 512,
+    kv_cache_max_len: int = 1024,
+    quantize: bool = True,
+):
+  """An example method for converting TinyLlama model to multi-signature
+  tflite model.
+
+  Args:
+      checkpoint_path (str): The filepath to the model checkpoint, or directory holding the checkpoint.
+      prefill_seq_len (int, optional): The maximum size of prefill input tensor.
+        Defaults to 512.
+      kv_cache_max_len (int, optional): The maximum size of KV cache buffer,
+        including both prefill and decode. Defaults to 1024.
+      quantize (bool, optional): Whether the model should be quanized.
+        Defaults to True.
+  """
+  pytorch_model = tiny_llama.build_model(
+      checkpoint_path, kv_cache_max_len=kv_cache_max_len
+  )
+  # Tensors used to trace the model graph during conversion.
+  prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
+  prefill_input_pos = torch.arange(0, prefill_seq_len)
+  decode_token = torch.tensor([[0]], dtype=torch.long)
+  decode_input_pos = torch.tensor([0], dtype=torch.int64)
+
+  quant_config = quant_recipes.full_linear_int8_dynamic_recipe() if quantize else None
+  edge_model = (
+      ai_edge_torch.signature(
+          'prefill', pytorch_model, (prefill_tokens, prefill_input_pos)
+      )
+      .signature('decode', pytorch_model, (decode_token, decode_input_pos))
+      .convert(quant_config=quant_config)
+  )
+  edge_model.export(f'/tmp/tiny_llama_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite')
+
+
+if __name__ == '__main__':
+  checkpoint_path = os.path.join(Path.home(), 'Downloads/llm_data/tiny_llama')
+  convert_tiny_llama_to_tflite(checkpoint_path)