model_builder.py

# Copyright 2024 The AI Edge Torch Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Utilities to be used for re-authoring transformer models."""

importcopy
fromtypingimportOptional, Tuple

fromai_edge_torch.generative.layersimportattention
fromai_edge_torch.generative.layersimportbuilder
fromai_edge_torch.generative.layersimportkv_cacheaskv_utils
fromai_edge_torch.generative.layersimportloraaslora_utils
importai_edge_torch.generative.layers.attention_utilsasattn_utils
importai_edge_torch.generative.layers.model_configascfg
fromai_edge_torch.generative.utilitiesimportexport_configasexport_cfg
importai_edge_torch.generative.utilities.loaderasloading_utils
importtorch
fromtorchimportnn


TENSOR_NAMES=loading_utils.ModelLoader.TensorNames(
ff_up_proj="model.layers.{}.mlp.up_proj",
ff_down_proj="model.layers.{}.mlp.down_proj",
ff_gate_proj="model.layers.{}.mlp.gate_proj",
attn_query_proj="model.layers.{}.self_attn.q_proj",
attn_key_proj="model.layers.{}.self_attn.k_proj",
attn_value_proj="model.layers.{}.self_attn.v_proj",
attn_output_proj="model.layers.{}.self_attn.o_proj",
pre_attn_norm="model.layers.{}.input_layernorm",
post_attn_norm="model.layers.{}.post_attention_layernorm",
embedding="model.embed_tokens",
final_norm="model.norm",
)

TENSOR_NAMES_WITH_SEPARATE_LM_HEAD=copy.copy(TENSOR_NAMES)
TENSOR_NAMES_WITH_SEPARATE_LM_HEAD.lm_head="lm_head"


classDecoderOnlyModel(nn.Module):
"""A simple decoder-only transformer model built from the Edge Generative API.

 This model is used for re-authoring. model_config is used to specify the
 details of model architecture and parameters.

 It assumes that the attention configs for ROPE, i.e. head_dim, rotary_base,
 and rotary_percentage are the same for all layers.
 """

def__init__(self, config: cfg.ModelConfig):
super().__init__()

# Construct model layers.
self.tok_embedding=nn.Embedding(
config.vocab_size, config.embedding_dim, padding_idx=0
 )
self.lm_head=nn.Linear(
config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
 )
ifconfig.lm_head_share_weight_with_embedding:
self.lm_head.weight.data=self.tok_embedding.weight.data
self.transformer_blocks=nn.ModuleList(
attention.TransformerBlock(config.block_config(idx), config)
foridxinrange(config.num_layers)
 )
self.final_norm=builder.build_norm(
config.embedding_dim,
config.final_norm_config,
 )
self.mask_cache=attn_utils.build_causal_mask_cache(
size=config.kv_cache_max,
 )
self.config=config

@torch.inference_mode
defforward(
self,
tokens: torch.Tensor,
input_pos: torch.Tensor,
kv_cache: kv_utils.KVCache,
mask: Optional[torch.Tensor] =None,
lora: Optional[lora_utils.LoRA] =None,
export_config: Optional[export_cfg.ExportConfig] =None,
 ) ->dict[torch.Tensor, kv_utils.KVCache]:
_, seq_len=tokens.size()
assertself.config.max_seq_len>=seq_len, (
f"Cannot forward sequence of length {seq_len}, max seq length is only"
f" {self.config.max_seq_len}"
 )

# token embeddings of shape (b, t, n_embd)
input_embeds=self.tok_embedding(tokens)

# ROPE parameters for all attn_configs are the same. Take the first one.
attn_config=self.config.block_config(0).attn_config
n_elem=int(attn_config.rotary_percentage*attn_config.head_dim)
rope=self.config.build_rope(input_pos, n_elem, attn_config.rotary_base)

ifmaskisNone:
mask=self.mask_cache.index_select(2, input_pos)
mask=mask[:, :, :, : self.config.kv_cache_max]

returnself._forward_with_embeds(
input_embeds, rope, mask, input_pos, kv_cache, lora, export_config
 )

def_forward_with_embeds(
self,
input_embeds: torch.Tensor,
rope: Tuple[torch.Tensor, torch.Tensor],
mask: torch.Tensor,
input_pos: torch.Tensor,
kv_cache: kv_utils.KVCache,
lora: Optional[lora_utils.LoRA] =None,
export_config: Optional[export_cfg.ExportConfig] =None,
 ) ->dict[torch.Tensor, kv_utils.KVCache]:
"""Forwards the model with input embeddings."""
assertlen(self.transformer_blocks) ==len(kv_cache.caches), (
"The number of transformer blocks and the number of KV cache entries"
" must be the same."
 )

x=input_embeds
ifself.config.embedding_scaleisnotNone:
x=x*self.config.embedding_scale

updated_kv_entries= []
fori, blockinenumerate(self.transformer_blocks):
kv_entry=kv_cache.caches[i] ifkv_cacheelseNone
lora_adapter=lora.adapters[i] ifloraelseNone
x, kv_entry=block(x, rope, mask, input_pos, kv_entry, lora_adapter)
ifkv_entry:
updated_kv_entries.append(kv_entry)
updated_kv_cache=kv_utils.KVCache(tuple(updated_kv_entries))

ifexport_configisnotNone:
if (
torch.numel(input_pos) >1
andnotexport_config.output_logits_on_prefill
 ):
return {"kv_cache": updated_kv_cache}

x=self.final_norm(x)
logits=self.lm_head(x) # (b, t, vocab_size)
return {"logits": logits, "kv_cache": updated_kv_cache}


defbuild_decoder_only_model(
checkpoint_path: str,
config: cfg.ModelConfig,
tensor_names: loading_utils.ModelLoader.TensorNames,
model_class: type[nn.Module] =DecoderOnlyModel,
) ->nn.Module:
transformer=model_class(config)
loader=loading_utils.ModelLoader(checkpoint_path, tensor_names)
loader.load(
transformer, strict=notconfig.lm_head_share_weight_with_embedding
 )
transformer.eval()
returntransformer