ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/generative/examples/t5/t5.py

@@ -17,17 +17,15 @@
 import copy
 import os
 from pathlib import Path
-from typing import Optional, Tuple
-
-import numpy as np
-import torch
-import torch.nn as nn
+from typing import Optional
 
 from ai_edge_torch.generative.examples.t5.t5_attention import EncoderDecoderBlock  # NOQA
 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
 import ai_edge_torch.generative.utilities.t5_loader as loading_utils
+import torch
+import torch.nn as nn
 
 ENCDEC_TENSOR_NAMES = {
     "ff_up_proj": "{prefix}.block.{}.layer.{num}.DenseReluDense.wi",
@@ -36,9 +34,11 @@ ENCDEC_TENSOR_NAMES = {
     "attn_key_proj": "{prefix}.block.{}.layer.0.SelfAttention.k",
     "attn_value_proj": "{prefix}.block.{}.layer.0.SelfAttention.v",
     "attn_output_proj": "{prefix}.block.{}.layer.0.SelfAttention.o",
-    "relative_attn_bias": "{prefix}.block.0.layer.0.SelfAttention.relative_attention_bias",
+    "relative_attn_bias": (
+        "{prefix}.block.0.layer.0.SelfAttention.relative_attention_bias"
+    ),
     "pre_attn_norm": "{prefix}.block.{}.layer.0.layer_norm",
-    "pre_ff_norm": "{prefix}.block.{}.layer.1.layer_norm",
+    "post_attn_norm": "{prefix}.block.{}.layer.1.layer_norm",
     "final_norm": "{prefix}.final_layer_norm",
 }
 
@@ -52,13 +52,19 @@ class T5Stack(nn.Module):
     self.config = config
     self.embed_tokens = embed_tokens
     self.is_decoder = config.is_decoder
-    self.transformer_blocks = nn.ModuleList(
-        [
-            EncoderDecoderBlock(config, has_relative_attention_bias=bool(i == 0))
-            for i in range(config.num_layers)
-        ]
+    # T5 has only one block config.
+    block_config = config.block_config(0)
+    self.transformer_blocks = nn.ModuleList([
+        EncoderDecoderBlock(
+            block_config,
+            config,
+            has_relative_attention_bias=bool(idx == 0),
+        )
+        for idx in range(config.num_layers)
+    ])
+    self.final_norm = builder.build_norm(
+        config.embedding_dim, config.final_norm_config
     )
-    self.final_norm = builder.build_norm(config.embedding_dim, config.final_norm_config)
 
   def forward(
       self,
@@ -73,23 +79,23 @@ class T5Stack(nn.Module):
           torch.Tensor
       ] = None,  # should be for decoder case
   ):
-    input_shape = input_ids.size()
     inputs_embeds = self.embed_tokens(input_ids)
-    batch_size, seq_length = input_shape
     hidden_states = inputs_embeds
     position_bias = None
     encoder_decoder_position_bias = None
-    for i, layer_module in enumerate(self.transformer_blocks):
+    for _, layer_module in enumerate(self.transformer_blocks):
       # EncoderDecoderBlock.forward
-      hidden_states, position_bias, encoder_decoder_position_bias = layer_module(
-          hidden_states,
-          input_pos,
-          mask=attention_mask,
-          relative_position=relative_position,
-          position_bias=position_bias,
-          encoder_hidden_states=encoder_hidden_states,
-          encoder_attention_mask=encoder_attention_mask,
-          encoder_decoder_position_bias=encoder_decoder_position_bias,
+      hidden_states, position_bias, encoder_decoder_position_bias = (
+          layer_module(
+              hidden_states,
+              input_pos,
+              mask=attention_mask,
+              relative_position=relative_position,
+              position_bias=position_bias,
+              encoder_hidden_states=encoder_hidden_states,
+              encoder_attention_mask=encoder_attention_mask,
+              encoder_decoder_position_bias=encoder_decoder_position_bias,
+          )
       )
 
     hidden_states = self.final_norm(hidden_states)
@@ -109,7 +115,8 @@ class T5(nn.Module):
 
     encoder_config = copy.deepcopy(config)
     encoder_config.is_decoder = False
-    encoder_config.attn_config.enable_kv_cache = False
+    # T5 has only one block config.
+    encoder_config.block_config(0).attn_config.enable_kv_cache = False
     self.encoder = T5Stack(encoder_config, self.tok_embedding)
 
     decoder_config = copy.deepcopy(config)
@@ -130,23 +137,27 @@ class T5(nn.Module):
     )
 
     self.dec_attn_mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
 
+    # T5 has only one block config.
+    attn_config = config.block_config(0).attn_config
     self.enc_rel_pos_mask = attn_utils.build_relative_position_buckets(
         bidirectional=True,
         query_length=config.kv_cache_max,
         key_length=config.kv_cache_max,
-        num_buckets=config.attn_config.relative_attention_num_buckets,
-        max_distance=config.attn_config.relative_attention_max_distance,
+        num_buckets=attn_config.relative_attention_num_buckets,
+        max_distance=attn_config.relative_attention_max_distance,
     )
 
     self.dec_rel_pos_mask = attn_utils.build_relative_position_buckets(
         bidirectional=False,
         query_length=config.kv_cache_max,
         key_length=config.kv_cache_max,
-        num_buckets=config.attn_config.relative_attention_num_buckets,
-        max_distance=config.attn_config.relative_attention_max_distance,
+        num_buckets=attn_config.relative_attention_num_buckets,
+        max_distance=attn_config.relative_attention_max_distance,
     )
 
   @torch.inference_mode
@@ -159,9 +170,10 @@ class T5(nn.Module):
       pad_mask: torch.Tensor,
   ) -> torch.Tensor:
     B, T = input_ids.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     enc_mask = self.enc_attn_mask_cache.index_select(2, input_pos)
     enc_mask = enc_mask[:, :, :, : self.config.kv_cache_max]
@@ -170,10 +182,18 @@ class T5(nn.Module):
     dec_mask = self.dec_attn_mask_cache.index_select(2, decoder_input_pos)
     dec_mask = dec_mask[:, :, :, : self.config.kv_cache_max]
     enc_relative_position = self.enc_rel_pos_mask.index_select(2, input_pos)
-    enc_relative_position = enc_relative_position[:, :, :, : self.config.kv_cache_max]
-    dec_relative_position = self.enc_rel_pos_mask.index_select(2, decoder_input_pos)
-    dec_relative_position = dec_relative_position[:, :, :, : self.config.kv_cache_max]
-    enc_attention_mask = self.enc_attn_mask_cache.index_select(2, decoder_input_pos)
+    enc_relative_position = enc_relative_position[
+        :, :, :, : self.config.kv_cache_max
+    ]
+    dec_relative_position = self.enc_rel_pos_mask.index_select(
+        2, decoder_input_pos
+    )
+    dec_relative_position = dec_relative_position[
+        :, :, :, : self.config.kv_cache_max
+    ]
+    enc_attention_mask = self.enc_attn_mask_cache.index_select(
+        2, decoder_input_pos
+    )
     # Mask off any "pad" tokens that shouldn't contribute to cross attention
     enc_attention_mask[:, :, :, :] += pad_mask
 
@@ -210,12 +230,15 @@ class T5Encoder(nn.Module):
 
     self.config = config
     # Construct model layers.
-    assert embedding_layer != None, "Passed in embedding layer should not be None!"
+    assert (
+        embedding_layer != None
+    ), "Passed in embedding layer should not be None!"
     self.tok_embedding = embedding_layer
 
     encoder_config = copy.deepcopy(config)
     encoder_config.is_decoder = False
-    encoder_config.attn_config.enable_kv_cache = False
+    # T5 has only one block config.
+    encoder_config.block_config(0).attn_config.enable_kv_cache = False
     self.encoder = T5Stack(encoder_config, self.tok_embedding)
 
     self.enc_attn_mask_cache = (
@@ -228,12 +251,14 @@ class T5Encoder(nn.Module):
         .unsqueeze(0)
     )
 
+    # T5 has only one block config.
+    attn_config = config.block_config(0).attn_config
     self.enc_rel_pos_mask = attn_utils.build_relative_position_buckets(
         bidirectional=True,
         query_length=config.kv_cache_max,
         key_length=config.kv_cache_max,
-        num_buckets=config.attn_config.relative_attention_num_buckets,
-        max_distance=config.attn_config.relative_attention_max_distance,
+        num_buckets=attn_config.relative_attention_num_buckets,
+        max_distance=attn_config.relative_attention_max_distance,
     )
 
   @torch.inference_mode
@@ -244,16 +269,19 @@ class T5Encoder(nn.Module):
       pad_mask: torch.Tensor,
   ) -> torch.Tensor:
     B, T = input_ids.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     enc_mask = self.enc_attn_mask_cache.index_select(2, input_pos)
     enc_mask = enc_mask[:, :, :, : self.config.kv_cache_max]
     # Mask off any "pad" tokens that shouldn't contribute to self-attention
     enc_mask[:, :, :, :] += pad_mask
     enc_relative_position = self.enc_rel_pos_mask.index_select(2, input_pos)
-    enc_relative_position = enc_relative_position[:, :, :, : self.config.kv_cache_max]
+    enc_relative_position = enc_relative_position[
+        :, :, :, : self.config.kv_cache_max
+    ]
 
     # Convert encoder inputs in embeddings if needed
     encoder_hidden_states = self.encoder(
@@ -273,7 +301,9 @@ class T5Decoder(nn.Module):
 
     self.config = config
     # Construct model layers.
-    assert embedding_layer != None, "Passed in embedding layer should not be None!"
+    assert (
+        embedding_layer != None
+    ), "Passed in embedding layer should not be None!"
     self.tok_embedding = embedding_layer
 
     decoder_config = copy.deepcopy(config)
@@ -293,16 +323,20 @@ class T5Decoder(nn.Module):
         .unsqueeze(0)
     )
 
+    # T5 has only one block config.
+    attn_config = config.block_config(0).attn_config
     self.enc_rel_pos_mask = attn_utils.build_relative_position_buckets(
         bidirectional=True,
         query_length=config.kv_cache_max,
         key_length=config.kv_cache_max,
-        num_buckets=config.attn_config.relative_attention_num_buckets,
-        max_distance=config.attn_config.relative_attention_max_distance,
+        num_buckets=attn_config.relative_attention_num_buckets,
+        max_distance=attn_config.relative_attention_max_distance,
     )
 
     self.dec_attn_mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
 
   @torch.inference_mode
@@ -315,9 +349,15 @@ class T5Decoder(nn.Module):
   ) -> torch.Tensor:
     dec_mask = self.dec_attn_mask_cache.index_select(2, decoder_input_pos)
     dec_mask = dec_mask[:, :, :, : self.config.kv_cache_max]
-    dec_relative_position = self.enc_rel_pos_mask.index_select(2, decoder_input_pos)
-    dec_relative_position = dec_relative_position[:, :, :, : self.config.kv_cache_max]
-    enc_attention_mask = self.enc_attn_mask_cache.index_select(2, decoder_input_pos)
+    dec_relative_position = self.enc_rel_pos_mask.index_select(
+        2, decoder_input_pos
+    )
+    dec_relative_position = dec_relative_position[
+        :, :, :, : self.config.kv_cache_max
+    ]
+    enc_attention_mask = self.enc_attn_mask_cache.index_select(
+        2, decoder_input_pos
+    )
     # Mask off any "pad" tokens that shouldn't contribute to cross attention
     enc_attention_mask[:, :, :, :] += pad_mask
 
@@ -342,6 +382,7 @@ class T5Decoder(nn.Module):
 def get_model_config_t5() -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=12,
+      head_dim=64,
       num_query_groups=12,
       qkv_use_bias=False,
       relative_attention_num_buckets=32,
@@ -357,19 +398,20 @@ def get_model_config_t5() -> cfg.ModelConfig:
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-6,
   )
-
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      relative_attention=True,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
   config = cfg.ModelConfig(
       vocab_size=32128,
       num_layers=12,
       max_seq_len=512,
       embedding_dim=768,
-      attn_config=attn_config,
-      relative_attention=True,
-      ff_config=ff_config,
-      pre_attention_norm_config=norm_config,
-      pre_ff_norm_config=norm_config,
+      block_configs=block_config,
       final_norm_config=norm_config,
-      parallel_residual=False,
       lm_head_use_bias=False,
       enable_hlfb=True,
   )
@@ -390,7 +432,7 @@ def build_t5_model(checkpoint_path: str) -> nn.Module:
       "cross_attn_value_proj": "{prefix}.block.{}.layer.1.EncDecAttention.v",
       "cross_attn_output_proj": "{prefix}.block.{}.layer.1.EncDecAttention.o",
       # In the decoder, the FF is layer 2 in the Transformer block
-      "pre_ff_norm": "{prefix}.block.{}.layer.2.layer_norm",
+      "post_attn_norm": "{prefix}.block.{}.layer.2.layer_norm",
       # In the decoder, the cross attention is layer 1 in the Transformer block
       "pre_cross_attn_norm": "{prefix}.block.{}.layer.1.layer_norm",
   }
@@ -446,7 +488,7 @@ def build_t5_decoder_model(
       "cross_attn_value_proj": "{prefix}.block.{}.layer.1.EncDecAttention.v",
       "cross_attn_output_proj": "{prefix}.block.{}.layer.1.EncDecAttention.o",
       # In the decoder, the FF is layer 2 in the Transformer block
-      "pre_ff_norm": "{prefix}.block.{}.layer.2.layer_norm",
+      "post_attn_norm": "{prefix}.block.{}.layer.2.layer_norm",
       # In the decoder, the cross attention is layer 1 in the Transformer block
       "pre_cross_attn_norm": "{prefix}.block.{}.layer.1.layer_norm",
   }
@@ -470,104 +512,101 @@ def build_t5_decoder_model(
 
 
 def get_sample_encoder_input_ids() -> torch.Tensor:
-  idx = torch.tensor(
-      [
-          [
-              3856,
-              27111,
-              10,
-              4425,
-              51,
-              4008,
-              31,
-              7,
-              2306,
-              16576,
-              47,
-              4381,
-              16,
-              8,
-              3414,
-              13,
-              1410,
-              16,
-              932,
-              11,
-              1515,
-              2766,
-              6,
-              11,
-              4838,
-              16,
-              23964,
-              16,
-              1797,
-              13,
-              24,
-              215,
-              5,
-              94,
-              47,
-              2017,
-              168,
-              1204,
-              57,
-              6800,
-              7,
-              11,
-              9443,
-              38,
-              3673,
-              8,
-              4016,
-              13,
-              66,
-              70,
-              14234,
-              5,
-              2449,
-              1215,
-              83,
-              17,
-              16,
-              8782,
-              70,
-              723,
-              30,
-              8,
-              6162,
-              13,
-              1410,
-              12,
-              48,
-              833,
-              250,
-              13,
-              149,
-              231,
-              79,
-              1858,
-              16576,
-              5,
-              1,
-          ]
-      ]
-  )
+  idx = torch.tensor([[
+      3856,
+      27111,
+      10,
+      4425,
+      51,
+      4008,
+      31,
+      7,
+      2306,
+      16576,
+      47,
+      4381,
+      16,
+      8,
+      3414,
+      13,
+      1410,
+      16,
+      932,
+      11,
+      1515,
+      2766,
+      6,
+      11,
+      4838,
+      16,
+      23964,
+      16,
+      1797,
+      13,
+      24,
+      215,
+      5,
+      94,
+      47,
+      2017,
+      168,
+      1204,
+      57,
+      6800,
+      7,
+      11,
+      9443,
+      38,
+      3673,
+      8,
+      4016,
+      13,
+      66,
+      70,
+      14234,
+      5,
+      2449,
+      1215,
+      83,
+      17,
+      16,
+      8782,
+      70,
+      723,
+      30,
+      8,
+      6162,
+      13,
+      1410,
+      12,
+      48,
+      833,
+      250,
+      13,
+      149,
+      231,
+      79,
+      1858,
+      16576,
+      5,
+      1,
+  ]])
   return idx
 
 
 def define_and_run_t5(checkpoint_path: str) -> None:
-  t5_goldens = torch.load("t5_lm_logits.pt")
+  current_dir = Path(__file__).parent.resolve()
+  t5_goldens = torch.load(current_dir / "t5_lm_logits.pt")
 
   model = build_t5_model(checkpoint_path)
 
   idx = get_sample_encoder_input_ids()
-  tokens = torch.full((1, 512), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, 512), 0, dtype=torch.int, device="cpu")
   tokens[0, :77] = idx
-  input_pos = torch.arange(0, 512)
+  input_pos = torch.arange(0, 512, dtype=torch.int)
 
-  decode_d_token = torch.tensor([[0]], dtype=torch.int64)
-  decode_d_input_pos = torch.tensor([0], dtype=torch.int64)
+  decode_d_token = torch.tensor([[0]], dtype=torch.int)
+  decode_d_input_pos = torch.tensor([0], dtype=torch.int)
   pad_mask = torch.zeros([model.config.kv_cache_max], dtype=torch.float32)
   pad_mask[77:] = float("-inf")
   lm_logits = model.forward(
@@ -579,20 +618,30 @@ def define_and_run_t5(checkpoint_path: str) -> None:
 
 # TODO(haoliang): Move those tests.
 def define_and_run_t5_split(checkpoint_path: str) -> None:
-  t5_goldens = torch.load("t5_lm_logits.pt")
+  current_dir = Path(__file__).parent.resolve()
+  t5_goldens = torch.load(current_dir / "t5_lm_logits.pt")
+
   config = get_model_config_t5()
-  embedding_layer = nn.Embedding(config.vocab_size, config.embedding_dim, padding_idx=0)
-  t5_encoder_model = build_t5_encoder_model(config, embedding_layer, checkpoint_path)
-  t5_decoder_model = build_t5_decoder_model(config, embedding_layer, checkpoint_path)
+  embedding_layer = nn.Embedding(
+      config.vocab_size, config.embedding_dim, padding_idx=0
+  )
+  t5_encoder_model = build_t5_encoder_model(
+      config, embedding_layer, checkpoint_path
+  )
+  t5_decoder_model = build_t5_decoder_model(
+      config, embedding_layer, checkpoint_path
+  )
   idx = get_sample_encoder_input_ids()
 
-  tokens = torch.full((1, 512), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, 512), 0, dtype=torch.int, device="cpu")
   tokens[0, :77] = idx
-  input_pos = torch.arange(0, 512)
+  input_pos = torch.arange(0, 512, dtype=torch.int)
 
-  decode_d_token = torch.tensor([[0]], dtype=torch.int64)
-  decode_d_input_pos = torch.tensor([0], dtype=torch.int64)
-  pad_mask = torch.zeros([t5_encoder_model.config.kv_cache_max], dtype=torch.float32)
+  decode_d_token = torch.tensor([[0]], dtype=torch.int)
+  decode_d_input_pos = torch.tensor([0], dtype=torch.int)
+  pad_mask = torch.zeros(
+      [t5_encoder_model.config.kv_cache_max], dtype=torch.float32
+  )
   pad_mask[77:] = float("-inf")
   hidden_states = t5_encoder_model.forward(tokens, input_pos, pad_mask)
   lm_logits = t5_decoder_model.forward(