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

Sign up to get free protection for your applications and to get access to all the features.
Files changed (169) hide show
  1. ai_edge_torch/__init__.py +5 -4
  2. ai_edge_torch/_convert/conversion.py +112 -0
  3. ai_edge_torch/_convert/conversion_utils.py +64 -0
  4. ai_edge_torch/{convert → _convert}/converter.py +94 -48
  5. ai_edge_torch/_convert/fx_passes/__init__.py +22 -0
  6. ai_edge_torch/{convert → _convert}/fx_passes/build_aten_composite_pass.py +107 -44
  7. ai_edge_torch/{convert → _convert}/fx_passes/build_interpolate_composite_pass.py +23 -20
  8. ai_edge_torch/{convert → _convert}/fx_passes/inject_mlir_debuginfo_pass.py +5 -6
  9. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/__init__.py +1 -1
  10. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_check.py +39 -9
  11. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_mark.py +2 -0
  12. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py +1 -0
  13. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py +17 -8
  14. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py +9 -8
  15. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py +31 -18
  16. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/op_func_registry.py +2 -2
  17. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/pass_body.py +34 -24
  18. ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/utils.py +2 -0
  19. ai_edge_torch/_convert/signature.py +66 -0
  20. ai_edge_torch/_convert/test/test_convert.py +495 -0
  21. ai_edge_torch/_convert/test/test_convert_composites.py +234 -0
  22. ai_edge_torch/_convert/test/test_convert_multisig.py +189 -0
  23. ai_edge_torch/{convert → _convert}/test/test_to_channel_last_io.py +5 -5
  24. ai_edge_torch/{convert → _convert}/to_channel_last_io.py +10 -3
  25. ai_edge_torch/config.py +27 -0
  26. ai_edge_torch/conftest.py +20 -0
  27. ai_edge_torch/debug/culprit.py +72 -40
  28. ai_edge_torch/debug/test/test_culprit.py +7 -5
  29. ai_edge_torch/debug/test/test_search_model.py +8 -7
  30. ai_edge_torch/debug/utils.py +14 -3
  31. ai_edge_torch/fx_pass_base.py +101 -0
  32. ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py +68 -0
  33. ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py +68 -0
  34. ai_edge_torch/generative/examples/gemma/{gemma.py → gemma1.py} +69 -55
  35. ai_edge_torch/generative/examples/gemma/gemma2.py +267 -0
  36. ai_edge_torch/generative/examples/gemma/verify_gemma1.py +56 -0
  37. ai_edge_torch/generative/examples/gemma/verify_gemma2.py +57 -0
  38. ai_edge_torch/generative/examples/gemma/verify_util.py +143 -0
  39. ai_edge_torch/generative/examples/openelm/convert_to_tflite.py +68 -0
  40. ai_edge_torch/generative/examples/openelm/openelm.py +206 -0
  41. ai_edge_torch/generative/examples/openelm/verify.py +64 -0
  42. ai_edge_torch/generative/examples/phi/__init__.py +14 -0
  43. ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py +68 -0
  44. ai_edge_torch/generative/examples/phi/convert_to_tflite.py +68 -0
  45. ai_edge_torch/generative/examples/{phi2 → phi}/phi2.py +70 -51
  46. ai_edge_torch/generative/examples/phi/phi3.py +286 -0
  47. ai_edge_torch/generative/examples/phi/verify.py +65 -0
  48. ai_edge_torch/generative/examples/phi/verify_phi3.py +70 -0
  49. ai_edge_torch/generative/examples/smollm/__init__.py +14 -0
  50. ai_edge_torch/generative/examples/smollm/convert_to_tflite.py +68 -0
  51. ai_edge_torch/generative/examples/smollm/smollm.py +101 -0
  52. ai_edge_torch/generative/examples/smollm/verify.py +62 -0
  53. ai_edge_torch/generative/examples/stable_diffusion/attention.py +3 -1
  54. ai_edge_torch/generative/examples/stable_diffusion/clip.py +83 -13
  55. ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py +27 -14
  56. ai_edge_torch/generative/examples/stable_diffusion/decoder.py +74 -9
  57. ai_edge_torch/generative/examples/stable_diffusion/diffusion.py +179 -37
  58. ai_edge_torch/generative/examples/stable_diffusion/encoder.py +4 -3
  59. ai_edge_torch/generative/examples/stable_diffusion/pipeline.py +83 -58
  60. ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler.py +4 -3
  61. ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler_ancestral.py +4 -3
  62. ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py +4 -3
  63. ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py +1 -0
  64. ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py +4 -1
  65. ai_edge_torch/generative/examples/stable_diffusion/util.py +9 -3
  66. ai_edge_torch/generative/examples/t5/convert_to_tflite.py +28 -25
  67. ai_edge_torch/generative/examples/t5/t5.py +208 -159
  68. ai_edge_torch/generative/examples/t5/t5_attention.py +45 -30
  69. ai_edge_torch/generative/examples/test_models/convert_toy_model.py +105 -0
  70. ai_edge_torch/generative/examples/test_models/toy_model.py +69 -41
  71. ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py +50 -64
  72. ai_edge_torch/generative/examples/tiny_llama/__init__.py +14 -0
  73. ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py +41 -39
  74. ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py +67 -54
  75. ai_edge_torch/generative/examples/tiny_llama/verify.py +64 -0
  76. ai_edge_torch/generative/fx_passes/__init__.py +4 -5
  77. ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py +10 -7
  78. ai_edge_torch/generative/layers/attention.py +141 -102
  79. ai_edge_torch/generative/layers/attention_utils.py +53 -12
  80. ai_edge_torch/generative/layers/builder.py +37 -7
  81. ai_edge_torch/generative/layers/feed_forward.py +39 -14
  82. ai_edge_torch/generative/layers/kv_cache.py +162 -50
  83. ai_edge_torch/generative/layers/model_config.py +84 -30
  84. ai_edge_torch/generative/layers/normalization.py +185 -7
  85. ai_edge_torch/generative/layers/rotary_position_embedding.py +6 -4
  86. ai_edge_torch/generative/layers/scaled_dot_product_attention.py +48 -21
  87. ai_edge_torch/generative/layers/unet/blocks_2d.py +136 -77
  88. ai_edge_torch/generative/layers/unet/builder.py +7 -4
  89. ai_edge_torch/generative/layers/unet/model_config.py +17 -15
  90. ai_edge_torch/generative/quantize/example.py +7 -8
  91. ai_edge_torch/generative/quantize/quant_recipe.py +10 -7
  92. ai_edge_torch/generative/quantize/quant_recipe_utils.py +12 -1
  93. ai_edge_torch/generative/quantize/quant_recipes.py +8 -0
  94. ai_edge_torch/generative/test/test_kv_cache.py +120 -0
  95. ai_edge_torch/generative/test/{loader_test.py → test_loader.py} +9 -7
  96. ai_edge_torch/generative/test/test_model_conversion.py +124 -188
  97. ai_edge_torch/generative/test/test_model_conversion_large.py +251 -0
  98. ai_edge_torch/generative/test/test_quantize.py +76 -60
  99. ai_edge_torch/generative/test/utils.py +54 -0
  100. ai_edge_torch/generative/utilities/converter.py +82 -0
  101. ai_edge_torch/generative/utilities/loader.py +120 -57
  102. ai_edge_torch/generative/utilities/stable_diffusion_loader.py +165 -57
  103. ai_edge_torch/generative/utilities/t5_loader.py +110 -81
  104. ai_edge_torch/generative/utilities/verifier.py +247 -0
  105. ai_edge_torch/hlfb/__init__.py +1 -1
  106. ai_edge_torch/hlfb/mark_pattern/__init__.py +9 -7
  107. ai_edge_torch/hlfb/mark_pattern/passes.py +23 -3
  108. ai_edge_torch/hlfb/mark_pattern/pattern.py +39 -30
  109. ai_edge_torch/hlfb/test/test_mark_pattern.py +46 -20
  110. ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py +24 -11
  111. ai_edge_torch/lowertools/__init__.py +18 -0
  112. ai_edge_torch/lowertools/_shim.py +80 -0
  113. ai_edge_torch/lowertools/common_utils.py +142 -0
  114. ai_edge_torch/lowertools/odml_torch_utils.py +255 -0
  115. ai_edge_torch/lowertools/test_utils.py +60 -0
  116. ai_edge_torch/lowertools/torch_xla_utils.py +284 -0
  117. ai_edge_torch/{generative/quantize/ai_edge_quantizer_glue → lowertools}/translate_recipe.py +29 -14
  118. ai_edge_torch/model.py +53 -18
  119. ai_edge_torch/odml_torch/__init__.py +20 -0
  120. ai_edge_torch/odml_torch/_torch_future.py +61 -0
  121. ai_edge_torch/odml_torch/_torch_library.py +19 -0
  122. ai_edge_torch/odml_torch/composite/__init__.py +16 -0
  123. ai_edge_torch/odml_torch/composite/mark_tensor.py +120 -0
  124. ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py +106 -0
  125. ai_edge_torch/odml_torch/debuginfo/__init__.py +16 -0
  126. ai_edge_torch/odml_torch/debuginfo/_build.py +43 -0
  127. ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py +55 -0
  128. ai_edge_torch/odml_torch/export.py +357 -0
  129. ai_edge_torch/odml_torch/export_utils.py +168 -0
  130. ai_edge_torch/odml_torch/jax_bridge/__init__.py +15 -0
  131. ai_edge_torch/odml_torch/jax_bridge/_wrap.py +150 -0
  132. ai_edge_torch/odml_torch/jax_bridge/utils.py +75 -0
  133. ai_edge_torch/odml_torch/lowerings/__init__.py +25 -0
  134. ai_edge_torch/odml_torch/lowerings/_basic.py +258 -0
  135. ai_edge_torch/odml_torch/lowerings/_batch_norm.py +65 -0
  136. ai_edge_torch/odml_torch/lowerings/_convolution.py +241 -0
  137. ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py +252 -0
  138. ai_edge_torch/odml_torch/lowerings/_layer_norm.py +78 -0
  139. ai_edge_torch/odml_torch/lowerings/context.py +42 -0
  140. ai_edge_torch/odml_torch/lowerings/registry.py +96 -0
  141. ai_edge_torch/odml_torch/lowerings/utils.py +185 -0
  142. ai_edge_torch/odml_torch/passes/__init__.py +38 -0
  143. ai_edge_torch/odml_torch/tf_integration.py +194 -0
  144. ai_edge_torch/quantize/pt2e_quantizer.py +52 -24
  145. ai_edge_torch/quantize/pt2e_quantizer_utils.py +43 -23
  146. ai_edge_torch/quantize/quant_config.py +13 -9
  147. ai_edge_torch/testing/model_coverage/model_coverage.py +29 -16
  148. ai_edge_torch/version.py +16 -0
  149. {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/METADATA +7 -3
  150. ai_edge_torch_nightly-0.3.0.dev20240926.dist-info/RECORD +177 -0
  151. {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/WHEEL +1 -1
  152. ai_edge_torch/convert/conversion.py +0 -117
  153. ai_edge_torch/convert/conversion_utils.py +0 -400
  154. ai_edge_torch/convert/fx_passes/__init__.py +0 -59
  155. ai_edge_torch/convert/fx_passes/_pass_base.py +0 -49
  156. ai_edge_torch/convert/fx_passes/canonicalize_pass.py +0 -37
  157. ai_edge_torch/convert/test/test_convert.py +0 -311
  158. ai_edge_torch/convert/test/test_convert_composites.py +0 -192
  159. ai_edge_torch/convert/test/test_convert_multisig.py +0 -139
  160. ai_edge_torch/generative/examples/gemma/convert_to_tflite.py +0 -66
  161. ai_edge_torch/generative/examples/phi2/convert_to_tflite.py +0 -64
  162. ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py +0 -161
  163. ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py +0 -0
  164. ai_edge_torch_nightly-0.2.0.dev20240714.dist-info/RECORD +0 -121
  165. /ai_edge_torch/{convert → _convert}/__init__.py +0 -0
  166. /ai_edge_torch/{convert → _convert}/test/__init__.py +0 -0
  167. /ai_edge_torch/generative/examples/{phi2 → openelm}/__init__.py +0 -0
  168. {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/LICENSE +0 -0
  169. {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/top_level.txt +0 -0
@@ -12,21 +12,17 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
- # Example of building phi-2 model from the Edge Generative API layers.
16
15
 
16
+ """Example of building a Phi-2 model."""
17
17
 
18
- import os
19
- from pathlib import Path
20
-
21
- import numpy as np
22
- import torch
23
- import torch.nn as nn
24
-
25
- from ai_edge_torch.generative.layers.attention import TransformerBlock
18
+ from ai_edge_torch.generative.layers import attention
19
+ from ai_edge_torch.generative.layers import builder
20
+ from ai_edge_torch.generative.layers import kv_cache as kv_utils
26
21
  import ai_edge_torch.generative.layers.attention_utils as attn_utils
27
- import ai_edge_torch.generative.layers.builder as builder
28
22
  import ai_edge_torch.generative.layers.model_config as cfg
29
23
  import ai_edge_torch.generative.utilities.loader as loading_utils
24
+ import torch
25
+ from torch import nn
30
26
 
31
27
  TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
32
28
  ff_up_proj="model.layers.{}.mlp.fc1",
@@ -43,11 +39,11 @@ TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
43
39
 
44
40
 
45
41
  class Phi2(nn.Module):
42
+ """A Phi-2 model built from the Edge Generative API layers."""
46
43
 
47
44
  def __init__(self, config: cfg.ModelConfig):
48
45
  super().__init__()
49
46
 
50
- self.config = config
51
47
  # Construct model layers.
52
48
  self.lm_head = nn.Linear(
53
49
  config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
@@ -55,35 +51,48 @@ class Phi2(nn.Module):
55
51
  self.tok_embedding = nn.Embedding(
56
52
  config.vocab_size, config.embedding_dim, padding_idx=0
57
53
  )
54
+ # Phi-2 has only one block config.
55
+ block_config = config.block_config(0)
58
56
  self.transformer_blocks = nn.ModuleList(
59
- TransformerBlock(config) for _ in range(config.num_layers)
57
+ attention.TransformerBlock(block_config, config)
58
+ for _ in range(config.num_layers)
60
59
  )
61
60
  self.final_norm = builder.build_norm(
62
61
  config.embedding_dim,
63
62
  config.final_norm_config,
64
63
  )
64
+ attn_config = block_config.attn_config
65
65
  self.rope_cache = attn_utils.build_rope_cache(
66
66
  size=config.kv_cache_max,
67
- dim=int(config.attn_config.rotary_percentage * config.head_dim),
67
+ dim=int(attn_config.rotary_percentage * attn_config.head_dim),
68
68
  base=10_000,
69
69
  condense_ratio=1,
70
70
  dtype=torch.float32,
71
71
  device=torch.device("cpu"),
72
72
  )
73
73
  self.mask_cache = attn_utils.build_causal_mask_cache(
74
- size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
74
+ size=config.kv_cache_max,
75
+ dtype=torch.float32,
76
+ device=torch.device("cpu"),
75
77
  )
76
78
  self.config = config
77
79
 
78
- # The model's forward function takes in additional k/v cache tensors
79
- # and returns the updated k/v cache tensors to the caller.
80
- # This can be eliminated if we handle k/v cache updates inside the model itself.
81
80
  @torch.inference_mode
82
- def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
83
- B, T = idx.size()
84
- assert (
85
- self.config.max_seq_len >= T
86
- ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
81
+ def forward(
82
+ self,
83
+ tokens: torch.Tensor,
84
+ input_pos: torch.Tensor,
85
+ kv_cache: kv_utils.KVCache,
86
+ ) -> dict[torch.Tensor, kv_utils.KVCache]:
87
+ _, seq_len = tokens.size()
88
+ assert self.config.max_seq_len >= seq_len, (
89
+ f"Cannot forward sequence of length {seq_len}, max seq length is only"
90
+ f" {self.config.max_seq_len}"
91
+ )
92
+ assert len(self.transformer_blocks) == len(kv_cache.caches), (
93
+ "The number of transformer blocks and the number of KV cache entries"
94
+ " must be the same."
95
+ )
87
96
 
88
97
  cos, sin = self.rope_cache
89
98
  cos = cos.index_select(0, input_pos)
@@ -91,20 +100,34 @@ class Phi2(nn.Module):
91
100
  mask = self.mask_cache.index_select(2, input_pos)
92
101
  mask = mask[:, :, :, : self.config.kv_cache_max]
93
102
 
94
- # forward the model itself
95
- x = self.tok_embedding(idx) # token embeddings of shape (b, t, n_embd)
103
+ x = self.tok_embedding(tokens)
96
104
 
105
+ updated_kv_entires = []
97
106
  for i, block in enumerate(self.transformer_blocks):
98
- x = block(x, (cos, sin), mask, input_pos)
107
+ kv_entry = kv_cache.caches[i] if kv_cache else None
108
+ x, kv_entry = block(x, (cos, sin), mask, input_pos, kv_entry)
109
+ if kv_entry:
110
+ updated_kv_entires.append(kv_entry)
111
+ updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
99
112
 
100
113
  x = self.final_norm(x)
101
- res = self.lm_head(x) # (b, t, vocab_size)
102
- return res
114
+ logits = self.lm_head(x) # (b, t, vocab_size)
115
+ return {"logits": logits, "kv_cache": updated_kv_cache}
103
116
 
104
117
 
105
118
  def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
119
+ """Returns the model config for a Phi-2 model.
120
+
121
+ Args:
122
+ kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
123
+ is 1024.
124
+
125
+ Returns:
126
+ The model config for a Phi-2 model.
127
+ """
106
128
  attn_config = cfg.AttentionConfig(
107
129
  num_heads=32,
130
+ head_dim=80,
108
131
  num_query_groups=32,
109
132
  rotary_percentage=0.4,
110
133
  qkv_use_bias=True,
@@ -116,49 +139,45 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
116
139
  intermediate_size=10240,
117
140
  use_bias=True,
118
141
  )
119
- norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.LAYER_NORM)
142
+ norm_config = cfg.NormalizationConfig(
143
+ type=cfg.NormalizationType.LAYER_NORM,
144
+ use_input_shape=False, # Phi-2 does layer-norm with the weight shape.
145
+ )
146
+ block_config = cfg.TransformerBlockConfig(
147
+ attn_config=attn_config,
148
+ ff_config=ff_config,
149
+ pre_attention_norm_config=norm_config,
150
+ parallel_residual=True,
151
+ )
120
152
  config = cfg.ModelConfig(
121
153
  vocab_size=51200,
122
154
  num_layers=32,
123
155
  max_seq_len=2048,
124
156
  kv_cache_max_len=kv_cache_max_len,
125
157
  embedding_dim=2560,
126
- attn_config=attn_config,
127
- ff_config=ff_config,
128
- pre_attention_norm_config=norm_config,
158
+ block_configs=block_config,
129
159
  final_norm_config=norm_config,
130
- parallel_residual=True,
131
160
  lm_head_use_bias=True,
132
161
  enable_hlfb=True,
133
162
  )
134
163
  return config
135
164
 
136
165
 
137
- def get_fake_model_config_for_test() -> cfg.ModelConfig:
138
- config = get_model_config()
166
+ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
167
+ config = get_model_config(kv_cache_max_len)
168
+ config.vocab_size = 128
139
169
  config.num_layers = 2
170
+ config.max_seq_len = 2 * kv_cache_max_len
171
+ # Phi-2 has only one block config.
172
+ config.block_config(0).ff_config.intermediate_size = 128
140
173
  return config
141
174
 
142
175
 
143
- def build_model(checkpoint_path, **kwargs) -> nn.Module:
176
+ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
177
+ """Instantiates the model instance and load checkpoint if provided."""
144
178
  config = get_model_config(**kwargs)
145
179
  model = Phi2(config)
146
180
  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
147
181
  loader.load(model)
182
+ model.eval()
148
183
  return model
149
-
150
-
151
- def define_and_run() -> None:
152
- kv_cache_max_len = 1024
153
- checkpoint_path = os.path.join(Path.home(), "Downloads/llm_data/phi2")
154
- model = build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
155
- idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
156
- tokens = torch.full((1, kv_cache_max_len), 0, dtype=torch.long, device="cpu")
157
- tokens[0, :4] = idx
158
- input_pos = torch.arange(0, kv_cache_max_len)
159
- print("running an inference")
160
- print(model.forward(tokens, input_pos))
161
-
162
-
163
- if __name__ == "__main__":
164
- define_and_run()
@@ -0,0 +1,286 @@
1
+ # Copyright 2024 The AI Edge Torch Authors.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+ """Example of building a Phi-3.5 model up to 4K tokens, not to 128K tokens."""
17
+
18
+ import math
19
+ from typing import Tuple
20
+
21
+ from ai_edge_torch.generative.layers import attention
22
+ from ai_edge_torch.generative.layers import builder
23
+ from ai_edge_torch.generative.layers import kv_cache as kv_utils
24
+ import ai_edge_torch.generative.layers.attention_utils as attn_utils
25
+ import ai_edge_torch.generative.layers.model_config as cfg
26
+ import ai_edge_torch.generative.utilities.loader as loading_utils
27
+ import torch
28
+ from torch import nn
29
+
30
+ TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
31
+ ff_up_proj="model.layers.{}.mlp.gate_up_proj",
32
+ ff_down_proj="model.layers.{}.mlp.down_proj",
33
+ attn_fused_qkv_proj="model.layers.{}.self_attn.qkv_proj",
34
+ attn_output_proj="model.layers.{}.self_attn.o_proj",
35
+ pre_attn_norm="model.layers.{}.input_layernorm",
36
+ post_attn_norm="model.layers.{}.post_attention_layernorm",
37
+ embedding="model.embed_tokens",
38
+ final_norm="model.norm",
39
+ lm_head="lm_head",
40
+ )
41
+
42
+ # max_position_embeddings / original_max_position_embeddings in Phi-3.5 config.
43
+ ROPE_SCALE_FACTOR = 32
44
+
45
+ # ROPE short factor in Phi-3.5 config. According to LOPE paper and its code in
46
+ # https://github.com/microsoft/LongRoPE, these values had been searched with
47
+ # min=1.0, step-0.01 to optimize the errors of sample dataset.
48
+ ROPE_SHORT_FACTOR = [
49
+ 1.0,
50
+ 1.0199999809265137,
51
+ 1.0299999713897705,
52
+ 1.0299999713897705,
53
+ 1.0499999523162842,
54
+ 1.0499999523162842,
55
+ 1.0499999523162842,
56
+ 1.0499999523162842,
57
+ 1.0499999523162842,
58
+ 1.0699999332427979,
59
+ 1.0999999046325684,
60
+ 1.1099998950958252,
61
+ 1.1599998474121094,
62
+ 1.1599998474121094,
63
+ 1.1699998378753662,
64
+ 1.2899998426437378,
65
+ 1.339999794960022,
66
+ 1.679999828338623,
67
+ 1.7899998426437378,
68
+ 1.8199998140335083,
69
+ 1.8499997854232788,
70
+ 1.8799997568130493,
71
+ 1.9099997282028198,
72
+ 1.9399996995925903,
73
+ 1.9899996519088745,
74
+ 2.0199997425079346,
75
+ 2.0199997425079346,
76
+ 2.0199997425079346,
77
+ 2.0199997425079346,
78
+ 2.0199997425079346,
79
+ 2.0199997425079346,
80
+ 2.0299997329711914,
81
+ 2.0299997329711914,
82
+ 2.0299997329711914,
83
+ 2.0299997329711914,
84
+ 2.0299997329711914,
85
+ 2.0299997329711914,
86
+ 2.0299997329711914,
87
+ 2.0299997329711914,
88
+ 2.0299997329711914,
89
+ 2.0799996852874756,
90
+ 2.0899996757507324,
91
+ 2.189999580383301,
92
+ 2.2199995517730713,
93
+ 2.5899994373321533,
94
+ 2.729999542236328,
95
+ 2.749999523162842,
96
+ 2.8399994373321533,
97
+ ]
98
+
99
+
100
+ def build_rope_cache(
101
+ size: int,
102
+ dim: int,
103
+ base: int = 10000,
104
+ condense_ratio: int = 1,
105
+ dtype: torch.dtype = torch.float32,
106
+ device: torch.device = None,
107
+ theta_factors: torch.Tensor = None,
108
+ scale: float = 1.0,
109
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
110
+ """Precomputes Rotary Positional Embeddings for Phi-3.5 model.
111
+
112
+ It's a modified version of attn_utils.build_rope_cache with additional
113
+ arguments for Phi-3.5 model. It precompute Rotary Positional Embedding Sin and
114
+ Cos values with scaling factors for quick lookup during the inference.
115
+
116
+ Args:
117
+ size (int): The size of the built cache.
118
+ dim (int): Each sequence's dimmension.
119
+ base (int, optional): Rope base value. Defaults to 10000.
120
+ condense_ratio (int, optional): The ratio by which sequence indicies are
121
+ condensed. Defaults to 1.
122
+ dtype (torch.dtype, optional): Output tensor's data type. Defaults to
123
+ torch.float32.
124
+ device (torch.device, optional): Output tensor's data type. Defaults to
125
+ None in which case "cpu" is used.
126
+ theta_factors (torch.Tensor, optional): A tensor of shape (dim,) used to
127
+ scale the theta values. Defaults to None.
128
+ scale (float, optional): A float used to scale the rope values. Defaults
129
+ to 1.0.
130
+
131
+ Returns:
132
+ Tuple[torch.Tensor, torch.Tensor]: Rope's Cosine and Sine waves.
133
+ """
134
+ if device is None:
135
+ device = torch.device('cpu')
136
+ theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
137
+ if theta_factors is not None:
138
+ theta = theta / theta_factors
139
+ seq_idx = torch.arange(size) / condense_ratio
140
+ idx_theta = torch.outer(seq_idx, theta)
141
+ cos = torch.cos(idx_theta).to(dtype=dtype, device=device) * scale
142
+ sin = torch.sin(idx_theta).to(dtype=dtype, device=device) * scale
143
+ return cos, sin
144
+
145
+
146
+ class Phi3_5Mini(nn.Module):
147
+ """A Phi-3.5 model built from the Edge Generative API layers."""
148
+
149
+ def __init__(self, config: cfg.ModelConfig):
150
+ super().__init__()
151
+
152
+ # Construct model layers.
153
+ self.lm_head = nn.Linear(
154
+ config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
155
+ )
156
+ self.tok_embedding = nn.Embedding(
157
+ config.vocab_size, config.embedding_dim, padding_idx=0
158
+ )
159
+ # Phi-3.5 has only one block config.
160
+ block_config = config.block_config(0)
161
+ self.transformer_blocks = nn.ModuleList(
162
+ attention.TransformerBlock(block_config, config)
163
+ for _ in range(config.num_layers)
164
+ )
165
+ self.final_norm = builder.build_norm(
166
+ config.embedding_dim,
167
+ config.final_norm_config,
168
+ )
169
+ attn_config = block_config.attn_config
170
+ self.rope_cache = build_rope_cache(
171
+ size=config.kv_cache_max,
172
+ dim=int(attn_config.rotary_percentage * attn_config.head_dim),
173
+ base=10_000,
174
+ condense_ratio=1,
175
+ dtype=torch.float32,
176
+ device=torch.device("cpu"),
177
+ theta_factors=torch.tensor(ROPE_SHORT_FACTOR),
178
+ scale=math.sqrt(
179
+ 1 + math.log(ROPE_SCALE_FACTOR) / math.log(config.max_seq_len)
180
+ ),
181
+ )
182
+ self.mask_cache = attn_utils.build_causal_mask_cache(
183
+ size=config.kv_cache_max,
184
+ dtype=torch.float32,
185
+ device=torch.device("cpu"),
186
+ )
187
+ self.config = config
188
+
189
+ @torch.inference_mode
190
+ def forward(
191
+ self,
192
+ tokens: torch.Tensor,
193
+ input_pos: torch.Tensor,
194
+ kv_cache: kv_utils.KVCache,
195
+ ) -> dict[torch.Tensor, kv_utils.KVCache]:
196
+ _, seq_len = tokens.size()
197
+ assert self.config.max_seq_len >= seq_len, (
198
+ f"Cannot forward sequence of length {seq_len}, max seq length is only"
199
+ f" {self.config.max_seq_len}"
200
+ )
201
+ assert len(self.transformer_blocks) == len(kv_cache.caches), (
202
+ "The number of transformer blocks and the number of KV cache entries"
203
+ " must be the same."
204
+ )
205
+
206
+ cos, sin = self.rope_cache
207
+ cos = cos.index_select(0, input_pos)
208
+ sin = sin.index_select(0, input_pos)
209
+ mask = self.mask_cache.index_select(2, input_pos)
210
+ mask = mask[:, :, :, : self.config.kv_cache_max]
211
+
212
+ x = self.tok_embedding(tokens)
213
+
214
+ updated_kv_entires = []
215
+ for i, block in enumerate(self.transformer_blocks):
216
+ kv_entry = kv_cache.caches[i] if kv_cache else None
217
+ x, kv_entry = block(x, (cos, sin), mask, input_pos, kv_entry)
218
+ if kv_entry:
219
+ updated_kv_entires.append(kv_entry)
220
+ updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entires))
221
+
222
+ x = self.final_norm(x)
223
+ logits = self.lm_head(x) # (b, t, vocab_size)
224
+ return {"logits": logits, "kv_cache": updated_kv_cache}
225
+
226
+
227
+ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
228
+ """Returns the model config for a Phi-3.5 model.
229
+
230
+ Args:
231
+ kv_cache_max_len (int): The maximum sequence length of the KV cache. Default
232
+ is 1024.
233
+
234
+ Returns:
235
+ The model config for a Phi-2 model.
236
+ """
237
+ attn_config = cfg.AttentionConfig(
238
+ num_heads=32,
239
+ head_dim=96,
240
+ num_query_groups=32,
241
+ rotary_percentage=1.0,
242
+ qkv_transpose_before_split=True,
243
+ )
244
+ ff_config = cfg.FeedForwardConfig(
245
+ type=cfg.FeedForwardType.SEQUENTIAL,
246
+ activation=cfg.ActivationConfig(cfg.ActivationType.SILU_GLU),
247
+ intermediate_size=8192,
248
+ )
249
+ norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
250
+ block_config = cfg.TransformerBlockConfig(
251
+ attn_config=attn_config,
252
+ ff_config=ff_config,
253
+ pre_attention_norm_config=norm_config,
254
+ post_attention_norm_config=norm_config,
255
+ )
256
+ config = cfg.ModelConfig(
257
+ vocab_size=32064,
258
+ num_layers=32,
259
+ max_seq_len=4096,
260
+ kv_cache_max_len=kv_cache_max_len,
261
+ embedding_dim=3072,
262
+ block_configs=block_config,
263
+ final_norm_config=norm_config,
264
+ enable_hlfb=True,
265
+ )
266
+ return config
267
+
268
+
269
+ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
270
+ config = get_model_config(kv_cache_max_len)
271
+ config.vocab_size = 128
272
+ config.num_layers = 2
273
+ config.max_seq_len = 2 * kv_cache_max_len
274
+ # Phi-3.5 has only one block config.
275
+ config.block_config(0).ff_config.intermediate_size = 128
276
+ return config
277
+
278
+
279
+ def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
280
+ """Instantiates the model instance and load checkpoint if provided."""
281
+ config = get_model_config(**kwargs)
282
+ model = Phi3_5Mini(config)
283
+ loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
284
+ loader.load(model)
285
+ model.eval()
286
+ return model
@@ -0,0 +1,65 @@
1
+ # Copyright 2024 The AI Edge Torch Authors.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+ """Verifies the reauthored Phi-2 model."""
17
+ import logging
18
+
19
+ from absl import app
20
+ from absl import flags
21
+ from ai_edge_torch.generative.examples.phi import phi2
22
+ from ai_edge_torch.generative.utilities import verifier
23
+ import kagglehub
24
+ import transformers
25
+
26
+
27
+ _PROMPTS = flags.DEFINE_multi_string(
28
+ "prompts",
29
+ "Instruct: Write an email about the weather Output:",
30
+ "The input prompts to generate answers.",
31
+ )
32
+ _MAX_NEW_TOKENS = flags.DEFINE_integer(
33
+ "max_new_tokens",
34
+ 30,
35
+ "The maximum size of the generated tokens.",
36
+ )
37
+
38
+
39
+ def main(_):
40
+ checkpoint = kagglehub.model_download("Microsoft/phi/transformers/2")
41
+ logging.info("Loading the original model from: %s", checkpoint)
42
+ generation_config = transformers.GenerationConfig.from_pretrained(checkpoint)
43
+ generation_config.max_new_tokens = _MAX_NEW_TOKENS.value
44
+ wrapper_model = verifier.ModelWrapper(
45
+ model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
46
+ hf_generation_config=generation_config,
47
+ )
48
+
49
+ logging.info("Building the reauthored model from: %s", checkpoint)
50
+ reauthored_model = phi2.build_model(checkpoint)
51
+
52
+ logging.info("Loading the tokenizer from: %s", checkpoint)
53
+ tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
54
+
55
+ verifier.verify_reauthored_model(
56
+ original_model=wrapper_model,
57
+ reauthored_model=reauthored_model,
58
+ tokenizer=tokenizer,
59
+ generate_prompts=_PROMPTS.value,
60
+ atol=1e-03,
61
+ )
62
+
63
+
64
+ if __name__ == "__main__":
65
+ app.run(main)
@@ -0,0 +1,70 @@
1
+ # Copyright 2024 The AI Edge Torch Authors.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+ """Verifies the reauthored Phi-3.5 model."""
17
+
18
+ import logging
19
+ import pathlib
20
+
21
+ from absl import app
22
+ from absl import flags
23
+ from ai_edge_torch.generative.examples.phi import phi3
24
+ from ai_edge_torch.generative.utilities import verifier
25
+ import transformers
26
+
27
+
28
+ _PROMPTS = flags.DEFINE_multi_string(
29
+ "prompts",
30
+ "Instruct: Write an email about the weather Output:",
31
+ "The input prompts to generate answers.",
32
+ )
33
+ _MAX_NEW_TOKENS = flags.DEFINE_integer(
34
+ "max_new_tokens",
35
+ 30,
36
+ "The maximum size of the generated tokens.",
37
+ )
38
+
39
+
40
+ def main(_):
41
+ checkpoint = "microsoft/Phi-3.5-mini-instruct"
42
+ logging.info("Loading the original model from: %s", checkpoint)
43
+ generation_config = transformers.GenerationConfig.from_pretrained(checkpoint)
44
+ generation_config.max_new_tokens = _MAX_NEW_TOKENS.value
45
+ wrapper_model = verifier.ModelWrapper(
46
+ model=transformers.AutoModelForCausalLM.from_pretrained(checkpoint),
47
+ hf_generation_config=generation_config,
48
+ )
49
+
50
+ # Locate the cached dir.
51
+ cached_config_file = transformers.utils.cached_file(
52
+ checkpoint, transformers.utils.CONFIG_NAME
53
+ )
54
+ reauthored_checkpoint = pathlib.Path(cached_config_file).parent
55
+ logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
56
+ reauthored_model = phi3.build_model(reauthored_checkpoint)
57
+
58
+ logging.info("Loading the tokenizer from: %s", checkpoint)
59
+ tokenizer = transformers.AutoTokenizer.from_pretrained(checkpoint)
60
+
61
+ verifier.verify_reauthored_model(
62
+ original_model=wrapper_model,
63
+ reauthored_model=reauthored_model,
64
+ tokenizer=tokenizer,
65
+ generate_prompts=_PROMPTS.value,
66
+ )
67
+
68
+
69
+ if __name__ == "__main__":
70
+ app.run(main)
@@ -0,0 +1,14 @@
1
+ # Copyright 2024 The AI Edge Torch Authors.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================