Explicit head_dim configuration (google-ai-edge#120)

talumbau · mbrenon · web-flow · commit 7300dd153530 · 2024-08-06T16:42:21.000Z
* Allow explictly setting head_dim instead of deducing it.

head_dim is also moved down to the attention config: this parameter is
related to attention so it has no reason to live under the main model
config.

This is also a requirement for the upcoming OpenELM models.

* Add head_dim for T5 Attention config.

* Stable Diffusion attention configs

- Use head_dim from AttentionConfig

* Config head_dim in generative/experimental

* Fix attention configs for SD loader utilities

---------

Co-authored-by: mbrenon &lt;mbrenon@google.com&gt;
diff --git a/ai_edge_torch/generative/examples/experimental/gemma/gemma.py b/ai_edge_torch/generative/examples/experimental/gemma/gemma.py
@@ -73,7 +73,9 @@ def __init__(self, config: cfg.ModelConfig):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -125,6 +127,7 @@ def forward(
 def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=8,
+      head_dim=256,
       num_query_groups=1,
       rotary_percentage=1.0,
   )
diff --git a/ai_edge_torch/generative/examples/experimental/phi/phi2.py b/ai_edge_torch/generative/examples/experimental/phi/phi2.py
@@ -68,7 +68,9 @@ def __init__(self, config: cfg.ModelConfig):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -118,6 +120,7 @@ def forward(
 def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=32,
+      head_dim=80,
       num_query_groups=32,
       rotary_percentage=0.4,
       qkv_use_bias=True,
diff --git a/ai_edge_torch/generative/examples/experimental/tiny_llama/tiny_llama.py b/ai_edge_torch/generative/examples/experimental/tiny_llama/tiny_llama.py
@@ -70,7 +70,9 @@ def __init__(self, config: cfg.ModelConfig):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -121,6 +123,7 @@ def forward(
 def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=32,
+      head_dim=64,
       num_query_groups=4,
       rotary_percentage=1.0,
   )
diff --git a/ai_edge_torch/generative/examples/gemma/gemma.py b/ai_edge_torch/generative/examples/gemma/gemma.py
@@ -68,7 +68,9 @@ def __init__(self, config: cfg.ModelConfig):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -113,6 +115,7 @@ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
 def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=8,
+      head_dim=256,
       num_query_groups=1,
       rotary_percentage=1.0,
   )
diff --git a/ai_edge_torch/generative/examples/phi2/phi2.py b/ai_edge_torch/generative/examples/phi2/phi2.py
@@ -63,7 +63,9 @@ def __init__(self, config: cfg.ModelConfig):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -107,6 +109,7 @@ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
 def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=32,
+      head_dim=80,
       num_query_groups=32,
       rotary_percentage=0.4,
       qkv_use_bias=True,
diff --git a/ai_edge_torch/generative/examples/stable_diffusion/clip.py b/ai_edge_torch/generative/examples/stable_diffusion/clip.py
@@ -92,6 +92,7 @@ def get_model_config() -> cfg.ModelConfig:
 
   attn_config = cfg.AttentionConfig(
       num_heads=num_heads,
+      head_dim=embedding_dim // num_heads,
       num_query_groups=num_query_groups,
       rotary_percentage=0.0,
       qkv_use_bias=True,
diff --git a/ai_edge_torch/generative/examples/stable_diffusion/decoder.py b/ai_edge_torch/generative/examples/stable_diffusion/decoder.py
@@ -288,6 +288,7 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
       normalization_config=norm_config,
       attention_config=layers_cfg.AttentionConfig(
           num_heads=1,
+          head_dim=block_out_channels[-1],
           num_query_groups=1,
           qkv_use_bias=True,
           output_proj_use_bias=True,
diff --git a/ai_edge_torch/generative/examples/stable_diffusion/diffusion.py b/ai_edge_torch/generative/examples/stable_diffusion/diffusion.py
@@ -195,6 +195,31 @@
 )
 
 
+def build_attention_config(
+    num_heads,
+    dim,
+    num_query_groups,
+    rotary_percentage=0.0,
+    qkv_transpose_before_split=True,
+    qkv_use_bias=False,
+    output_proj_use_bias=True,
+    enable_kv_cache=False,
+    qkv_fused_interleaved=False,
+):
+
+  return layers_cfg.AttentionConfig(
+      num_heads=num_heads,
+      head_dim=dim // num_heads,
+      num_query_groups=num_query_groups,
+      rotary_percentage=rotary_percentage,
+      qkv_transpose_before_split=qkv_transpose_before_split,
+      qkv_use_bias=qkv_use_bias,
+      output_proj_use_bias=output_proj_use_bias,
+      enable_kv_cache=enable_kv_cache,
+      qkv_fused_interleaved=qkv_fused_interleaved,
+  )
+
+
 class TimeEmbedding(nn.Module):
 
   def __init__(self, in_dim, out_dim):
@@ -267,17 +292,6 @@ def __init__(self, config: unet_cfg.DiffusionModelConfig):
         config.in_channels, block_out_channels[0], kernel_size=3, padding=1
     )
 
-    attention_config = layers_cfg.AttentionConfig(
-        num_heads=config.transformer_num_attention_heads,
-        num_query_groups=config.transformer_num_attention_heads,
-        rotary_percentage=0.0,
-        qkv_transpose_before_split=True,
-        qkv_use_bias=False,
-        output_proj_use_bias=True,
-        enable_kv_cache=False,
-        qkv_fused_interleaved=False,
-    )
-
     # Down encoders.
     down_encoders = []
     output_channel = block_out_channels[0]
@@ -312,15 +326,23 @@ def __init__(self, config: unet_cfg.DiffusionModelConfig):
                             dim=output_channel,
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
-                            attention_config=attention_config,
+                            attention_config=build_attention_config(
+                                num_heads=config.transformer_num_attention_heads,
+                                dim=output_channel,
+                                num_query_groups=config.transformer_num_attention_heads,
+                            ),
                             enable_hlfb=False,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
                             cross_dim=config.transformer_cross_attention_dim,
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
-                            attention_config=attention_config,
+                            attention_config=build_attention_config(
+                                num_heads=config.transformer_num_attention_heads,
+                                dim=output_channel,
+                                num_query_groups=config.transformer_num_attention_heads,
+                            ),
                             enable_hlfb=False,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
@@ -374,15 +396,23 @@ def __init__(self, config: unet_cfg.DiffusionModelConfig):
                     dim=mid_block_channels,
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
-                    attention_config=attention_config,
+                    attention_config=build_attention_config(
+                        num_heads=config.transformer_num_attention_heads,
+                        dim=mid_block_channels,
+                        num_query_groups=config.transformer_num_attention_heads,
+                    ),
                     enable_hlfb=False,
                 ),
                 cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                     query_dim=mid_block_channels,
                     cross_dim=config.transformer_cross_attention_dim,
                     attention_batch_size=config.transformer_batch_size,
                     normalization_config=config.transformer_norm_config,
-                    attention_config=attention_config,
+                    attention_config=build_attention_config(
+                        num_heads=config.transformer_num_attention_heads,
+                        dim=mid_block_channels,
+                        num_query_groups=config.transformer_num_attention_heads,
+                    ),
                     enable_hlfb=False,
                 ),
                 pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
@@ -437,15 +467,23 @@ def __init__(self, config: unet_cfg.DiffusionModelConfig):
                             dim=output_channel,
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
-                            attention_config=attention_config,
+                            attention_config=build_attention_config(
+                                num_heads=config.transformer_num_attention_heads,
+                                dim=output_channel,
+                                num_query_groups=config.transformer_num_attention_heads,
+                            ),
                             enable_hlfb=False,
                         ),
                         cross_attention_block_config=unet_cfg.CrossAttentionBlock2DConfig(
                             query_dim=output_channel,
                             cross_dim=config.transformer_cross_attention_dim,
                             attention_batch_size=config.transformer_batch_size,
                             normalization_config=config.transformer_norm_config,
-                            attention_config=attention_config,
+                            attention_config=build_attention_config(
+                                num_heads=config.transformer_num_attention_heads,
+                                dim=output_channel,
+                                num_query_groups=config.transformer_num_attention_heads,
+                            ),
                             enable_hlfb=False,
                         ),
                         pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
diff --git a/ai_edge_torch/generative/examples/t5/t5.py b/ai_edge_torch/generative/examples/t5/t5.py
@@ -371,6 +371,7 @@ def forward(
 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,
diff --git a/ai_edge_torch/generative/examples/t5/t5_attention.py b/ai_edge_torch/generative/examples/t5/t5_attention.py
@@ -185,7 +185,7 @@ def forward(
     )  # batch size, sequence length, embedding dimensionality (n_embd)
     query_states = self.q_projection(x)
     query_states = query_states.reshape(
-        B, T, -1, self.head_dim
+        B, T, -1, self.config.head_dim
     )  # (B, T, nh_q, hs)
 
     if key_value_states is not None:
@@ -198,13 +198,13 @@ def forward(
       )  # batch size, sequence length, embedding dimensionality (n_embd)
       key_states = self.k_projection(key_value_states)
       value_states = self.v_projection(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)
+      key_states = key_states.reshape(kvB, kvT, -1, self.config.head_dim)
+      value_states = value_states.reshape(kvB, kvT, -1, self.config.head_dim)
     else:
       key_states = self.k_projection(x)
       value_states = self.v_projection(x)
-      key_states = key_states.reshape(B, T, -1, self.head_dim)
-      value_states = value_states.reshape(B, T, -1, self.head_dim)
+      key_states = key_states.reshape(B, T, -1, self.config.head_dim)
+      value_states = value_states.reshape(B, T, -1, self.config.head_dim)
 
     if key_value_states is None and self.kv_cache is not None:
       key_states, value_states = self.kv_cache.update_cache(
@@ -221,15 +221,15 @@ def forward(
             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(B, self.n_heads, T, self.config.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,
+        self.config.head_dim,
         mask=mask,
         scale=1.0,
     )
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model.py b/ai_edge_torch/generative/examples/test_models/toy_model.py
@@ -43,7 +43,9 @@ def __init__(self, config: cfg.ModelConfig) -> None:
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.max_seq_len,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -72,6 +74,7 @@ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
 def get_model_config() -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=32,
+      head_dim=4,
       num_query_groups=4,
       rotary_percentage=1.0,
       enable_kv_cache=False,
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py b/ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py
@@ -46,7 +46,9 @@ def __init__(self, config: cfg.ModelConfig) -> None:
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.max_seq_len,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -90,7 +92,7 @@ def _export_stablehlo_mlir(model, args):
 
 def get_model_config() -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
-      num_heads=32, num_query_groups=4, rotary_percentage=1.0
+      num_heads=32, head_dim=4, num_query_groups=4, rotary_percentage=1.0
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
diff --git a/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py b/ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py
@@ -45,7 +45,9 @@ def __init__(self, config: cfg.ModelConfig) -> None:
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.max_seq_len,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -78,7 +80,7 @@ def _export_stablehlo_mlir(model, args):
 
 def get_model_config() -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
-      num_heads=32, num_query_groups=4, rotary_percentage=1.0
+      num_heads=32, head_dim=4, num_query_groups=4, rotary_percentage=1.0
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,
diff --git a/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py b/ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py
@@ -64,7 +64,9 @@ def __init__(self, config: cfg.ModelConfig):
     )
     self.rope_cache = attn_utils.build_rope_cache(
         size=config.kv_cache_max,
-        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        dim=int(
+            config.attn_config.rotary_percentage * config.attn_config.head_dim
+        ),
         base=10_000,
         condense_ratio=1,
         dtype=torch.float32,
@@ -109,6 +111,7 @@ def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
 def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   attn_config = cfg.AttentionConfig(
       num_heads=32,
+      head_dim=64,
       num_query_groups=4,
       rotary_percentage=1.0,
   )
diff --git a/ai_edge_torch/generative/fx_passes/test/test_remove_sdpa_zero_mask_pass.py b/ai_edge_torch/generative/fx_passes/test/test_remove_sdpa_zero_mask_pass.py
@@ -99,6 +99,7 @@ def get_model_config() -> unet_cfg.AttentionBlock2DConfig:
           normalization_config=norm_config,
           attention_config=layers_cfg.AttentionConfig(
               num_heads=1,
+              head_dim=block_out_channels[-1],
               num_query_groups=1,
               qkv_use_bias=True,
               output_proj_use_bias=True,
diff --git a/ai_edge_torch/generative/layers/attention.py b/ai_edge_torch/generative/layers/attention.py
diff --git a/ai_edge_torch/generative/layers/experimental/ekv_cache.py b/ai_edge_torch/generative/layers/experimental/ekv_cache.py
diff --git a/ai_edge_torch/generative/layers/model_config.py b/ai_edge_torch/generative/layers/model_config.py
diff --git a/ai_edge_torch/generative/test/test_experimental_ekv.py b/ai_edge_torch/generative/test/test_experimental_ekv.py
diff --git a/ai_edge_torch/generative/utilities/loader.py b/ai_edge_torch/generative/utilities/loader.py
diff --git a/ai_edge_torch/generative/utilities/stable_diffusion_loader.py b/ai_edge_torch/generative/utilities/stable_diffusion_loader.py
diff --git a/ai_edge_torch/generative/utilities/t5_loader.py b/ai_edge_torch/generative/utilities/t5_loader.py
