440 lines
12 KiB
Python
440 lines
12 KiB
Python
"""
|
|
Recreate the Core ML model from scratch using
|
|
coremltools' neural_network.NeuralNetworkBuilder
|
|
"""
|
|
import coremltools
|
|
import coremltools.models.datatypes as datatypes
|
|
from coremltools.models import neural_network as neural_network
|
|
from coremltools.models.utils import save_spec
|
|
import numpy as np
|
|
|
|
# get weights
|
|
from transformers import GPT2LMHeadModel, GPT2Tokenizer
|
|
|
|
model_name = "./distilgpt2-base-pretrained-he"
|
|
save_directory = "tmp/coreml/"
|
|
#!mkdir -p $save_directory
|
|
file_name = "model.mlmodel"
|
|
|
|
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
|
|
lm_head_model = GPT2LMHeadModel.from_pretrained(model_name).eval()
|
|
model = lm_head_model.transformer
|
|
|
|
wte = model.wte.weight.data.numpy().transpose() # shape (768, 50257) /!\ i hate this
|
|
wpe = model.wpe.weight.data.numpy().transpose() # shape (768, 1024)
|
|
|
|
sequence_length = 64
|
|
steps = 6
|
|
|
|
# build model
|
|
input_features = [
|
|
('input_ids', datatypes.Array(sequence_length)),
|
|
('position_ids', datatypes.Array(sequence_length)),
|
|
]
|
|
output_features = [('output_logits', None)]
|
|
|
|
builder = neural_network.NeuralNetworkBuilder(
|
|
input_features,
|
|
output_features,
|
|
mode=None,
|
|
disable_rank5_shape_mapping=True,
|
|
)
|
|
builder.add_expand_dims(
|
|
name='input_ids_expanded_to_rank5',
|
|
input_name='input_ids',
|
|
output_name='input_ids_expanded_to_rank5',
|
|
axes=(1, 2, 3, 4)
|
|
)
|
|
builder.add_expand_dims(
|
|
name='position_ids_expanded_to_rank5',
|
|
input_name='position_ids',
|
|
output_name='position_ids_expanded_to_rank5',
|
|
axes=(1, 2, 3, 4)
|
|
)
|
|
builder.add_embedding(
|
|
name='token_embeddings',
|
|
input_name='input_ids_expanded_to_rank5',
|
|
output_name='token_embeddings',
|
|
W=wte,
|
|
b=None,
|
|
input_dim=50257,
|
|
output_channels=768,
|
|
has_bias=False,
|
|
)
|
|
builder.add_embedding(
|
|
name='positional_embeddings',
|
|
input_name='position_ids_expanded_to_rank5',
|
|
output_name='positional_embeddings',
|
|
W=wpe,
|
|
b=None,
|
|
input_dim=1024,
|
|
output_channels=768,
|
|
has_bias=False,
|
|
)
|
|
|
|
# Input:, Output: (seq, 1, 768, 1, 1)
|
|
builder.add_add_broadcastable(
|
|
name='embeddings_addition',
|
|
input_names=['token_embeddings', 'positional_embeddings'],
|
|
output_name=f'{0}_previous_block'
|
|
)
|
|
|
|
for i in range(steps):
|
|
print(i)
|
|
ln_weight = model.h[i].ln_1.weight.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
ln_bias = model.h[i].ln_1.bias.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
ln_epsilon = model.h[i].ln_1.eps
|
|
|
|
builder.add_mvn(
|
|
name=f"{i}_block_ln_1",
|
|
input_name=f"{i}_previous_block",
|
|
# output_name=f"{i}_block_ln_1_output",
|
|
output_name=f"{i}_block_ln_1",
|
|
across_channels=True,
|
|
normalize_variance=True,
|
|
epsilon=ln_epsilon
|
|
)
|
|
|
|
builder.add_scale(
|
|
name=f"{i}_block_ln_1_scaled",
|
|
input_name=f"{i}_block_ln_1",
|
|
output_name=f"{i}_block_ln_1_scaled",
|
|
W=ln_weight,
|
|
b=ln_bias,
|
|
has_bias=True,
|
|
shape_scale=[768],
|
|
shape_bias=[768]
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_ln_1_reshape",
|
|
input_name=f"{i}_block_ln_1_scaled",
|
|
output_name=f"{i}_block_ln_1_scaled_transposed",
|
|
axes=(1, 0, 2, 3, 4)
|
|
)
|
|
|
|
|
|
conv_1D_bias = model.h[i].attn.c_attn.bias.data.numpy().reshape((1, 1, 2304, 1, 1))
|
|
conv_1D_weights = model.h[i].attn.c_attn.weight.data.numpy().transpose().reshape((1, 768, 2304, 1, 1))
|
|
|
|
builder.add_inner_product(
|
|
name=f"{i}_block_attn_conv",
|
|
input_name=f"{i}_block_ln_1_scaled_transposed",
|
|
output_name=f"{i}_block_attn_conv",
|
|
input_channels=768,
|
|
output_channels=2304,
|
|
W=conv_1D_weights,
|
|
b=conv_1D_bias,
|
|
has_bias=True
|
|
)
|
|
|
|
builder.add_split(
|
|
name=f"{i}_block_attn_qkv_split",
|
|
input_name=f"{i}_block_attn_conv",
|
|
output_names=[f"{i}_block_attn_q", f"{i}_block_attn_k", f"{i}_block_attn_v"]
|
|
)
|
|
|
|
builder.add_rank_preserving_reshape(
|
|
name=f"{i}_block_attn_q_reshape",
|
|
input_name=f"{i}_block_attn_q",
|
|
output_name=f"{i}_block_attn_q_reshape",
|
|
output_shape=(1, 1, sequence_length, 12, 64)
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_attn_q_reshape_permuted",
|
|
input_name=f"{i}_block_attn_q_reshape",
|
|
output_name=f"{i}_block_attn_q_reshape_permuted",
|
|
axes=(0, 1, 3, 2, 4)
|
|
)
|
|
|
|
builder.add_rank_preserving_reshape(
|
|
name=f"{i}_block_attn_k_reshape",
|
|
input_name=f"{i}_block_attn_k",
|
|
output_name=f"{i}_block_attn_k_reshape",
|
|
output_shape=(1, 1, sequence_length, 12, 64)
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_attn_k_reshape_permuted",
|
|
input_name=f"{i}_block_attn_k_reshape",
|
|
output_name=f"{i}_block_attn_k_reshape_permuted",
|
|
axes=(0, 1, 3, 4, 2)
|
|
)
|
|
|
|
builder.add_rank_preserving_reshape(
|
|
name=f"{i}_block_attn_v_reshape",
|
|
input_name=f"{i}_block_attn_v",
|
|
output_name=f"{i}_block_attn_v_reshape",
|
|
output_shape=(1, 1, sequence_length, 12, 64)
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_attn_v_reshape_permuted",
|
|
input_name=f"{i}_block_attn_v_reshape",
|
|
output_name=f"{i}_block_attn_v_reshape_permuted",
|
|
axes=(0, 1, 3, 2, 4)
|
|
)
|
|
|
|
builder.add_batched_mat_mul(
|
|
name=f"{i}_block_attn_qv_matmul",
|
|
input_names=[f"{i}_block_attn_q_reshape_permuted", f"{i}_block_attn_k_reshape_permuted"],
|
|
output_name=f"{i}_block_attn_qv_matmul"
|
|
)
|
|
|
|
builder.add_scale(
|
|
name=f"{i}_block_attn_qv_matmul_scaled",
|
|
input_name=f"{i}_block_attn_qv_matmul",
|
|
output_name=f"{i}_block_attn_qv_matmul_scaled",
|
|
W=np.array(1/8),
|
|
b=0,
|
|
has_bias=False
|
|
)
|
|
|
|
bias_0 = model.h[i].attn.bias
|
|
nd = ns = sequence_length
|
|
b = (model.h[i].attn.bias[:, :, ns-nd:ns, :ns]).unsqueeze(0)
|
|
|
|
builder.add_scale(
|
|
name=f"{i}_block_attn_bias",
|
|
input_name=f"{i}_block_attn_qv_matmul_scaled",
|
|
output_name=f"{i}_block_attn_bias",
|
|
W=b,
|
|
b=None,
|
|
has_bias=False,
|
|
shape_scale=[1, sequence_length, sequence_length]
|
|
)
|
|
|
|
bias_constant_0 = - 1e4 * (1 - b)
|
|
|
|
builder.add_bias(
|
|
name=f"{i}_block_attn_afterbias",
|
|
input_name=f"{i}_block_attn_bias",
|
|
output_name=f"{i}_block_attn_afterbias",
|
|
# output_name=f"output_logits",
|
|
b=bias_constant_0,
|
|
shape_bias=[1, sequence_length, sequence_length],
|
|
)
|
|
|
|
builder.add_squeeze(
|
|
name=f"{i}_squeezit",
|
|
input_name=f"{i}_block_attn_afterbias",
|
|
output_name=f"{i}_squeezit",
|
|
axes=[0, 1]
|
|
)
|
|
|
|
builder.add_softmax(
|
|
name=f"{i}_block_attn_softmax",
|
|
input_name=f"{i}_squeezit",
|
|
output_name=f"{i}_block_attn_softmax",
|
|
)
|
|
|
|
builder.add_expand_dims(
|
|
name=f"{i}_expandit",
|
|
input_name=f"{i}_block_attn_softmax",
|
|
output_name=f"{i}_expandit",
|
|
axes=[0, 1]
|
|
)
|
|
|
|
builder.add_batched_mat_mul(
|
|
name=f"{i}_block_full_attention",
|
|
input_names=[f"{i}_expandit", f"{i}_block_attn_v_reshape_permuted"],
|
|
output_name=f"{i}_block_full_attention"
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_full_attention_merged_t",
|
|
input_name=f"{i}_block_full_attention",
|
|
output_name=f"{i}_block_full_attention_merged_t",
|
|
axes=[0, 1, 3, 2, 4]
|
|
)
|
|
|
|
builder.add_rank_preserving_reshape(
|
|
name=f"{i}_block_full_attention_merged",
|
|
input_name=f"{i}_block_full_attention_merged_t",
|
|
output_name=f"{i}_block_full_attention_merged",
|
|
output_shape=[1, 1, 1, sequence_length, 768]
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_attn_conv_proj_t",
|
|
input_name=f"{i}_block_full_attention_merged",
|
|
output_name=f"{i}_block_attn_conv_proj_t",
|
|
axes=[0, 3, 4, 1, 2]
|
|
)
|
|
|
|
conv_1D_proj_bias = model.h[i].attn.c_proj.bias.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
conv_1D_proj_weights = model.h[i].attn.c_proj.weight.data.numpy().transpose().reshape((1, 768, 768, 1, 1))
|
|
|
|
# Input:, Output: (1, 3, 768, 1, 1)
|
|
builder.add_inner_product(
|
|
name=f"{i}_block_attn_conv_proj",
|
|
input_name=f"{i}_block_attn_conv_proj_t",
|
|
output_name=f"{i}_block_attn_conv_proj",
|
|
input_channels=768,
|
|
output_channels=768,
|
|
W=conv_1D_proj_weights,
|
|
b=conv_1D_proj_bias,
|
|
has_bias=True
|
|
)
|
|
|
|
# Input: (seq, 1, 768, 1, 1), Output: (1, seq, 768, 1, 1)
|
|
builder.add_transpose(
|
|
name=f"{i}_previous_block_t",
|
|
input_name=f'{i}_previous_block',
|
|
output_name=f"{i}_previous_block_t",
|
|
axes=[1, 0, 2, 3, 4]
|
|
)
|
|
|
|
# Input: [(1, seq, 768, 1, 1), (1, seq, 768, 1, 1)], Output: (1, seq, 768, 1, 1)
|
|
builder.add_add_broadcastable(
|
|
name=f"{i}_block_xa_sum",
|
|
input_names=[f"{i}_previous_block_t", f"{i}_block_attn_conv_proj"],
|
|
output_name=f"{i}_block_xa_sum",
|
|
# output_name=f"output_logits"
|
|
)
|
|
|
|
ln_2_weight = model.h[i].ln_2.weight.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
ln_2_bias = model.h[i].ln_2.bias.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
ln_2_epsilon = model.h[i].ln_2.eps
|
|
|
|
# Input: (1, seq, 768, 1, 1), Output:
|
|
builder.add_mvn(
|
|
name=f"{i}_block_ln_2",
|
|
input_name=f"{i}_block_xa_sum",
|
|
output_name=f"{i}_block_ln_2",
|
|
across_channels=True,
|
|
normalize_variance=True,
|
|
epsilon=ln_2_epsilon
|
|
)
|
|
|
|
builder.add_scale(
|
|
name=f"{i}_block_ln_2_scaled",
|
|
input_name=f"{i}_block_ln_2",
|
|
# output_name=f"output_logits",
|
|
output_name=f"{i}_block_ln_2_scaled",
|
|
W=ln_2_weight,
|
|
b=ln_2_bias,
|
|
has_bias=True,
|
|
shape_scale=[768],
|
|
shape_bias=[768]
|
|
)
|
|
|
|
mlp_conv_1D_fc_bias = model.h[i].mlp.c_fc.bias.data.numpy().reshape((1, 1, 3072, 1, 1))
|
|
mlp_conv_1D_fc_weights = model.h[i].mlp.c_fc.weight.data.numpy().transpose().reshape((1, 768, 3072, 1, 1))
|
|
|
|
# Input:, Output: (1, 3, 3072, 1, 1)
|
|
builder.add_inner_product(
|
|
name=f"{i}_block_mlp_conv_fc",
|
|
input_name=f"{i}_block_ln_2_scaled",
|
|
output_name=f"{i}_block_mlp_conv_fc",
|
|
# output_name=f"output_logits",
|
|
input_channels=768,
|
|
output_channels=3072,
|
|
W=mlp_conv_1D_fc_weights,
|
|
b=mlp_conv_1D_fc_bias,
|
|
has_bias=True
|
|
)
|
|
|
|
builder.add_gelu(
|
|
name=f"{i}_block_mlp_gelu",
|
|
input_name=f"{i}_block_mlp_conv_fc",
|
|
output_name=f"{i}_block_mlp_gelu",
|
|
# output_name=f"output_logits",
|
|
mode='TANH_APPROXIMATION'
|
|
)
|
|
|
|
mlp_conv_1D_proj_bias = model.h[i].mlp.c_proj.bias.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
mlp_conv_1D_proj_weights = model.h[i].mlp.c_proj.weight.data.numpy().transpose().reshape((1, 3072, 768, 1, 1))
|
|
|
|
# Input:, Output: (1, 3, 3072, 1, 1)
|
|
builder.add_inner_product(
|
|
name=f"{i}_block_mlp_conv_proj",
|
|
input_name=f"{i}_block_mlp_gelu",
|
|
output_name=f"{i}_block_mlp_conv_proj",
|
|
# output_name=f"output_logits",
|
|
input_channels=3072,
|
|
output_channels=768,
|
|
W=mlp_conv_1D_proj_weights,
|
|
b=mlp_conv_1D_proj_bias,
|
|
has_bias=True
|
|
)
|
|
|
|
builder.add_add_broadcastable(
|
|
name=f"{i}_block_xm_sum",
|
|
input_names=[f"{i}_block_xa_sum", f"{i}_block_mlp_conv_proj"],
|
|
# output_name=f"output_logits"
|
|
output_name=f"{i + 1}_previous_block_final"
|
|
)
|
|
|
|
builder.add_transpose(
|
|
name=f"{i}_block_xm_sum_t",
|
|
input_name=f"{i + 1}_previous_block_final",
|
|
output_name=f"{i + 1}_previous_block",
|
|
axes=[1, 0, 2, 3, 4]
|
|
)
|
|
|
|
|
|
ln_f_weight = model.ln_f.weight.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
ln_f_bias = model.ln_f.bias.data.numpy().reshape((1, 1, 768, 1, 1))
|
|
ln_f_epsilon = model.ln_f.eps
|
|
|
|
# Input: (1, seq, 768, 1, 1), Output:
|
|
builder.add_mvn(
|
|
name=f"ln_f",
|
|
input_name=f"{steps}_previous_block_final",
|
|
output_name=f"ln_f",
|
|
# output_name=f"output_logits",
|
|
across_channels=True,
|
|
normalize_variance=True,
|
|
epsilon=ln_f_epsilon
|
|
)
|
|
|
|
builder.add_scale(
|
|
name=f"ln_f_scaled",
|
|
input_name=f"ln_f",
|
|
output_name=f"ln_f_scaled",
|
|
# output_name=f"output_logits",
|
|
W=ln_f_weight,
|
|
b=ln_f_bias,
|
|
has_bias=True,
|
|
shape_scale=[768],
|
|
shape_bias=[768]
|
|
)
|
|
|
|
lm_head_weights = lm_head_model.lm_head.weight.data.numpy().reshape((1, 50257, 768, 1, 1))
|
|
|
|
builder.add_inner_product(
|
|
name="lm_head",
|
|
input_name="ln_f_scaled",
|
|
output_name="output_logits",
|
|
input_channels=768,
|
|
output_channels=50257,
|
|
W=lm_head_weights,
|
|
b=None,
|
|
has_bias=False
|
|
)
|
|
|
|
# compile spec to model
|
|
mlmodel = coremltools.models.MLModel(builder.spec)
|
|
|
|
#save_spec(builder.spec, f'./{model_name}-{sequence_length}-{steps}.mlmodel')
|
|
save_spec(builder.spec, f'./{save_directory}{file_name}')
|
|
# model = coremltools.models.MLModel('gpt2.mlmodel')
|
|
|
|
# input_ids = np.zeros(sequence_length)
|
|
# position_ids = np.arange(sequence_length).astype(np.float)
|
|
|
|
# input_data = {
|
|
# 'input_ids': input_ids,
|
|
# 'position_ids': position_ids,
|
|
# }
|
|
|
|
# predictions = mlmodel.predict(input_data)["output_logits"]
|
|
# equal = np.amax(predictions - mlp_conv_proj.detach().numpy())
|
|
|
|
# print(predictions)
|
|
|
|
|
|
# save_spec(builder.spec, 'gpt2.mlmodel')
|