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.
- ai_edge_torch/__init__.py +5 -4
- ai_edge_torch/_convert/conversion.py +112 -0
- ai_edge_torch/_convert/conversion_utils.py +64 -0
- ai_edge_torch/{convert → _convert}/converter.py +94 -48
- ai_edge_torch/_convert/fx_passes/__init__.py +22 -0
- ai_edge_torch/{convert → _convert}/fx_passes/build_aten_composite_pass.py +107 -44
- ai_edge_torch/{convert → _convert}/fx_passes/build_interpolate_composite_pass.py +23 -20
- ai_edge_torch/{convert → _convert}/fx_passes/inject_mlir_debuginfo_pass.py +5 -6
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/__init__.py +1 -1
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_check.py +39 -9
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_mark.py +2 -0
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py +1 -0
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py +17 -8
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py +9 -8
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py +31 -18
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/op_func_registry.py +2 -2
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/pass_body.py +34 -24
- ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/utils.py +2 -0
- ai_edge_torch/_convert/signature.py +66 -0
- ai_edge_torch/_convert/test/test_convert.py +495 -0
- ai_edge_torch/_convert/test/test_convert_composites.py +234 -0
- ai_edge_torch/_convert/test/test_convert_multisig.py +189 -0
- ai_edge_torch/{convert → _convert}/test/test_to_channel_last_io.py +5 -5
- ai_edge_torch/{convert → _convert}/to_channel_last_io.py +10 -3
- ai_edge_torch/config.py +27 -0
- ai_edge_torch/conftest.py +20 -0
- ai_edge_torch/debug/culprit.py +72 -40
- ai_edge_torch/debug/test/test_culprit.py +7 -5
- ai_edge_torch/debug/test/test_search_model.py +8 -7
- ai_edge_torch/debug/utils.py +14 -3
- ai_edge_torch/fx_pass_base.py +101 -0
- ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py +68 -0
- ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py +68 -0
- ai_edge_torch/generative/examples/gemma/{gemma.py → gemma1.py} +69 -55
- ai_edge_torch/generative/examples/gemma/gemma2.py +267 -0
- ai_edge_torch/generative/examples/gemma/verify_gemma1.py +56 -0
- ai_edge_torch/generative/examples/gemma/verify_gemma2.py +57 -0
- ai_edge_torch/generative/examples/gemma/verify_util.py +143 -0
- ai_edge_torch/generative/examples/openelm/convert_to_tflite.py +68 -0
- ai_edge_torch/generative/examples/openelm/openelm.py +206 -0
- ai_edge_torch/generative/examples/openelm/verify.py +64 -0
- ai_edge_torch/generative/examples/phi/__init__.py +14 -0
- ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py +68 -0
- ai_edge_torch/generative/examples/phi/convert_to_tflite.py +68 -0
- ai_edge_torch/generative/examples/{phi2 → phi}/phi2.py +70 -51
- ai_edge_torch/generative/examples/phi/phi3.py +286 -0
- ai_edge_torch/generative/examples/phi/verify.py +65 -0
- ai_edge_torch/generative/examples/phi/verify_phi3.py +70 -0
- ai_edge_torch/generative/examples/smollm/__init__.py +14 -0
- ai_edge_torch/generative/examples/smollm/convert_to_tflite.py +68 -0
- ai_edge_torch/generative/examples/smollm/smollm.py +101 -0
- ai_edge_torch/generative/examples/smollm/verify.py +62 -0
- ai_edge_torch/generative/examples/stable_diffusion/attention.py +3 -1
- ai_edge_torch/generative/examples/stable_diffusion/clip.py +83 -13
- ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py +27 -14
- ai_edge_torch/generative/examples/stable_diffusion/decoder.py +74 -9
- ai_edge_torch/generative/examples/stable_diffusion/diffusion.py +179 -37
- ai_edge_torch/generative/examples/stable_diffusion/encoder.py +4 -3
- ai_edge_torch/generative/examples/stable_diffusion/pipeline.py +83 -58
- ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler.py +4 -3
- ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler_ancestral.py +4 -3
- ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py +4 -3
- ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py +1 -0
- ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py +4 -1
- ai_edge_torch/generative/examples/stable_diffusion/util.py +9 -3
- ai_edge_torch/generative/examples/t5/convert_to_tflite.py +28 -25
- ai_edge_torch/generative/examples/t5/t5.py +208 -159
- ai_edge_torch/generative/examples/t5/t5_attention.py +45 -30
- ai_edge_torch/generative/examples/test_models/convert_toy_model.py +105 -0
- ai_edge_torch/generative/examples/test_models/toy_model.py +69 -41
- ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py +50 -64
- ai_edge_torch/generative/examples/tiny_llama/__init__.py +14 -0
- ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py +41 -39
- ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py +67 -54
- ai_edge_torch/generative/examples/tiny_llama/verify.py +64 -0
- ai_edge_torch/generative/fx_passes/__init__.py +4 -5
- ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py +10 -7
- ai_edge_torch/generative/layers/attention.py +141 -102
- ai_edge_torch/generative/layers/attention_utils.py +53 -12
- ai_edge_torch/generative/layers/builder.py +37 -7
- ai_edge_torch/generative/layers/feed_forward.py +39 -14
- ai_edge_torch/generative/layers/kv_cache.py +162 -50
- ai_edge_torch/generative/layers/model_config.py +84 -30
- ai_edge_torch/generative/layers/normalization.py +185 -7
- ai_edge_torch/generative/layers/rotary_position_embedding.py +6 -4
- ai_edge_torch/generative/layers/scaled_dot_product_attention.py +48 -21
- ai_edge_torch/generative/layers/unet/blocks_2d.py +136 -77
- ai_edge_torch/generative/layers/unet/builder.py +7 -4
- ai_edge_torch/generative/layers/unet/model_config.py +17 -15
- ai_edge_torch/generative/quantize/example.py +7 -8
- ai_edge_torch/generative/quantize/quant_recipe.py +10 -7
- ai_edge_torch/generative/quantize/quant_recipe_utils.py +12 -1
- ai_edge_torch/generative/quantize/quant_recipes.py +8 -0
- ai_edge_torch/generative/test/test_kv_cache.py +120 -0
- ai_edge_torch/generative/test/{loader_test.py → test_loader.py} +9 -7
- ai_edge_torch/generative/test/test_model_conversion.py +124 -188
- ai_edge_torch/generative/test/test_model_conversion_large.py +251 -0
- ai_edge_torch/generative/test/test_quantize.py +76 -60
- ai_edge_torch/generative/test/utils.py +54 -0
- ai_edge_torch/generative/utilities/converter.py +82 -0
- ai_edge_torch/generative/utilities/loader.py +120 -57
- ai_edge_torch/generative/utilities/stable_diffusion_loader.py +165 -57
- ai_edge_torch/generative/utilities/t5_loader.py +110 -81
- ai_edge_torch/generative/utilities/verifier.py +247 -0
- ai_edge_torch/hlfb/__init__.py +1 -1
- ai_edge_torch/hlfb/mark_pattern/__init__.py +9 -7
- ai_edge_torch/hlfb/mark_pattern/passes.py +23 -3
- ai_edge_torch/hlfb/mark_pattern/pattern.py +39 -30
- ai_edge_torch/hlfb/test/test_mark_pattern.py +46 -20
- ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py +24 -11
- ai_edge_torch/lowertools/__init__.py +18 -0
- ai_edge_torch/lowertools/_shim.py +80 -0
- ai_edge_torch/lowertools/common_utils.py +142 -0
- ai_edge_torch/lowertools/odml_torch_utils.py +255 -0
- ai_edge_torch/lowertools/test_utils.py +60 -0
- ai_edge_torch/lowertools/torch_xla_utils.py +284 -0
- ai_edge_torch/{generative/quantize/ai_edge_quantizer_glue → lowertools}/translate_recipe.py +29 -14
- ai_edge_torch/model.py +53 -18
- ai_edge_torch/odml_torch/__init__.py +20 -0
- ai_edge_torch/odml_torch/_torch_future.py +61 -0
- ai_edge_torch/odml_torch/_torch_library.py +19 -0
- ai_edge_torch/odml_torch/composite/__init__.py +16 -0
- ai_edge_torch/odml_torch/composite/mark_tensor.py +120 -0
- ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py +106 -0
- ai_edge_torch/odml_torch/debuginfo/__init__.py +16 -0
- ai_edge_torch/odml_torch/debuginfo/_build.py +43 -0
- ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py +55 -0
- ai_edge_torch/odml_torch/export.py +357 -0
- ai_edge_torch/odml_torch/export_utils.py +168 -0
- ai_edge_torch/odml_torch/jax_bridge/__init__.py +15 -0
- ai_edge_torch/odml_torch/jax_bridge/_wrap.py +150 -0
- ai_edge_torch/odml_torch/jax_bridge/utils.py +75 -0
- ai_edge_torch/odml_torch/lowerings/__init__.py +25 -0
- ai_edge_torch/odml_torch/lowerings/_basic.py +258 -0
- ai_edge_torch/odml_torch/lowerings/_batch_norm.py +65 -0
- ai_edge_torch/odml_torch/lowerings/_convolution.py +241 -0
- ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py +252 -0
- ai_edge_torch/odml_torch/lowerings/_layer_norm.py +78 -0
- ai_edge_torch/odml_torch/lowerings/context.py +42 -0
- ai_edge_torch/odml_torch/lowerings/registry.py +96 -0
- ai_edge_torch/odml_torch/lowerings/utils.py +185 -0
- ai_edge_torch/odml_torch/passes/__init__.py +38 -0
- ai_edge_torch/odml_torch/tf_integration.py +194 -0
- ai_edge_torch/quantize/pt2e_quantizer.py +52 -24
- ai_edge_torch/quantize/pt2e_quantizer_utils.py +43 -23
- ai_edge_torch/quantize/quant_config.py +13 -9
- ai_edge_torch/testing/model_coverage/model_coverage.py +29 -16
- ai_edge_torch/version.py +16 -0
- {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/METADATA +7 -3
- ai_edge_torch_nightly-0.3.0.dev20240926.dist-info/RECORD +177 -0
- {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/WHEEL +1 -1
- ai_edge_torch/convert/conversion.py +0 -117
- ai_edge_torch/convert/conversion_utils.py +0 -400
- ai_edge_torch/convert/fx_passes/__init__.py +0 -59
- ai_edge_torch/convert/fx_passes/_pass_base.py +0 -49
- ai_edge_torch/convert/fx_passes/canonicalize_pass.py +0 -37
- ai_edge_torch/convert/test/test_convert.py +0 -311
- ai_edge_torch/convert/test/test_convert_composites.py +0 -192
- ai_edge_torch/convert/test/test_convert_multisig.py +0 -139
- ai_edge_torch/generative/examples/gemma/convert_to_tflite.py +0 -66
- ai_edge_torch/generative/examples/phi2/convert_to_tflite.py +0 -64
- ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py +0 -161
- ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py +0 -0
- ai_edge_torch_nightly-0.2.0.dev20240714.dist-info/RECORD +0 -121
- /ai_edge_torch/{convert → _convert}/__init__.py +0 -0
- /ai_edge_torch/{convert → _convert}/test/__init__.py +0 -0
- /ai_edge_torch/generative/examples/{phi2 → openelm}/__init__.py +0 -0
- {ai_edge_torch_nightly-0.2.0.dev20240714.dist-info → ai_edge_torch_nightly-0.3.0.dev20240926.dist-info}/LICENSE +0 -0
- {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
|
19
|
-
from
|
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(
|
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(
|
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,
|
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(
|
83
|
-
|
84
|
-
|
85
|
-
|
86
|
-
|
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
|
-
|
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
|
-
|
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
|
-
|
102
|
-
return
|
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(
|
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
|
-
|
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
|
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
|
+
# ==============================================================================
|