ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/utilities/loader.py

@@ -0,0 +1,398 @@
+# 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.
+# ==============================================================================
+# Common utility functions for data loading etc.
+from dataclasses import dataclass
+import glob
+import os
+from typing import Callable, Dict, List, Tuple
+
+from ai_edge_torch.generative.layers import model_config
+from safetensors import safe_open
+import torch
+
+
+def load_safetensors(full_path: str):
+  """Loads safetensors into a single state dictionary.
+
+  Args:
+    full_path (string): the directory that contains the safetensor files.
+
+  Returns:
+    A state dictionary contating loaded tensors.
+
+  Raises:
+    ValueError: If no tensors are loaded from the provided directory or file.
+  """
+  pattern = (
+      os.path.join(full_path, "*.safetensors")
+      if os.path.isdir(full_path)
+      else full_path
+  )
+  files = []
+  for file in glob.glob(pattern):
+    files.append(file)
+
+  tensors = {}
+  for file in files:
+    with safe_open(file, framework="pt") as fp:
+      for k in fp.keys():
+        assert k not in tensors
+        tensors[k] = fp.get_tensor(k)
+
+  if not tensors:
+    raise ValueError("Failed to load SafeTensors.")
+  return tensors
+
+
+def load_pytorch_statedict(full_path: str):
+  """Loads state dictionary binaries into a single state dictionary.
+
+  Args:
+    full_path (string): the directory that contains the bin files.
+
+  Returns:
+    A state dictionary contating loaded tensors.
+
+  Raises:
+    ValueError: If no tensors are loaded from the provided directory or file.
+  """
+  files = []
+  patterns = []
+  if os.path.isdir(full_path):
+    patterns.append(os.path.join(full_path, "*.bin"))
+    patterns.append(os.path.join(full_path, "*pt"))
+  else:
+    patterns.append(full_path)
+  for pattern in patterns:
+    for file in glob.glob(pattern):
+      files.append(file)
+
+  tensors = {}
+  for file in files:
+    this_file_tensors = torch.load(file)
+    for k in this_file_tensors:
+      assert k not in tensors
+    tensors.update(this_file_tensors)
+
+  if not tensors:
+    raise ValueError("Failed to load torch bin files.")
+  return tensors
+
+
+class ModelLoader:
+  """Utlity for loading model checkpoints to the Edge Generative API layer."""
+
+  @dataclass
+  class TensorNames:
+    attn_query_proj: str = None
+    attn_key_proj: str = None
+    attn_value_proj: str = None
+    attn_fused_qkv_proj: str = None
+    attn_output_proj: str = None
+    attn_query_norm: str = None
+    attn_key_norm: str = None
+
+    ff_up_proj: str = None
+    ff_down_proj: str = None
+    ff_gate_proj: str = None
+
+    pre_attn_norm: str = None
+    post_attn_norm: str = None
+    pre_ff_norm: str = None
+    post_ff_norm: str = None
+    embedding: str = None
+    embedding_position: str = None
+    final_norm: str = None
+    lm_head: str = None
+
+  def __init__(self, file_name: str, names: TensorNames) -> None:
+    """ModelLoader constructor.
+
+    Can be used to load multiple models of the same type.
+
+    Args:
+        file_name (str): Path to the checkpoint. Can be a directory or an exact
+          file.
+        names (TensorNames): An instance of `TensorNames` to determine mappings.
+    """
+    self._file_name = file_name
+    self._names = names
+    self._loader = self._get_loader()
+
+  def get_state(self) -> Dict[str, torch.Tensor]:
+    return self._loader(self._file_name)
+
+  def load(
+      self, model: torch.nn.Module, strict: bool = True
+  ) -> Tuple[List[str], List[str]]:
+    """Load the model from the checkpoint.
+
+    Args:
+        model (torch.nn.Module): The pytorch model that needs to be loaded.
+        strict (bool, optional): Whether the converted keys are strictly
+          matched. Defaults to True.
+
+    Returns:
+        missing_keys (List[str]): a list of str containing the missing keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected
+        keys.
+
+    Raises:
+        ValueError: If conversion results in unmapped tensors and strict mode is
+          enabled.
+    """
+    state = self.get_state()
+    state = state["model_state_dict"] if "model_state_dict" in state else state
+    converted_state = dict()
+    if self._names.embedding is not None:
+      converted_state["tok_embedding.weight"] = state.pop(
+          f"{self._names.embedding}.weight"
+      )
+      if model.config.embedding_use_bias:
+        converted_state["tok_embedding.bias"] = state.pop(
+            f"{self._names.embedding}.bias"
+        )
+      if self._names.embedding_position is not None:
+        converted_state["tok_embedding_position"] = state.pop(
+            f"{self._names.embedding_position}"
+        )
+    if self._names.lm_head is not None:
+      converted_state["lm_head.weight"] = state.pop(
+          f"{self._names.lm_head}.weight"
+      )
+      if model.config.lm_head_use_bias:
+        converted_state["lm_head.bias"] = state.pop(
+            f"{self._names.lm_head}.bias"
+        )
+    if self._names.final_norm is not None:
+      final_norm_name = self._names.final_norm
+      converted_state["final_norm.weight"] = state.pop(
+          f"{final_norm_name}.weight"
+      )
+      if f"{final_norm_name}.bias" in state:
+        converted_state["final_norm.bias"] = state.pop(
+            f"{final_norm_name}.bias"
+        )
+
+    for i in range(model.config.num_layers):
+      self._map_norm(i, model.config, state, converted_state)
+      self._map_feedforward(i, model.config, state, converted_state)
+      self._map_attention(i, model.config, state, converted_state)
+
+    if strict and state:
+      raise ValueError(
+          f"Failed to map all tensor. Remaing tensor are: {list(state.keys())}"
+      )
+    return model.load_state_dict(converted_state, strict=strict)
+
+  def _get_loader(self) -> Callable[[str], Dict[str, torch.Tensor]]:
+    """A best effort method for finding appropriate state loader.
+
+    Raises:
+        ValueError: If it fails to find an appropriate loader.
+
+    Returns:
+        Callable[[str], Dict[str, torch.Tensor]]: State loader to be used.
+    """
+    if os.path.isdir(self._file_name):
+      if glob.glob(os.path.join(self._file_name, "*.safetensors")):
+        return load_safetensors
+      if glob.glob(os.path.join(self._file_name, "*.bin")) or glob.glob(
+          os.path.join(self._file_name, "*pt")
+      ):
+        return load_pytorch_statedict
+
+    if self._file_name.endswith(".safetensors"):
+      return load_safetensors
+
+    if self._file_name.endswith(".bin") or self._file_name.endswith("pt"):
+      return load_pytorch_statedict
+
+    raise ValueError("File format not supported.")
+
+  def _map_feedforward(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+  ):
+    prefix = f"transformer_blocks.{idx}"
+    ff_config = config.block_config(idx).ff_config
+    if ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
+      ff_up_proj_name = self._names.ff_up_proj.format(idx)
+      ff_down_proj_name = self._names.ff_down_proj.format(idx)
+      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+          f"{ff_up_proj_name}.weight"
+      )
+      converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+          f"{ff_down_proj_name}.weight"
+      )
+      if ff_config.use_bias:
+        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+            f"{ff_up_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+            f"{ff_down_proj_name}.bias"
+        )
+    else:
+      ff_up_proj_name = self._names.ff_up_proj.format(idx)
+      ff_down_proj_name = self._names.ff_down_proj.format(idx)
+      ff_gate_proj_name = self._names.ff_gate_proj.format(idx)
+      converted_state[f"{prefix}.ff.w3.weight"] = state.pop(
+          f"{ff_up_proj_name}.weight"
+      )
+      converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
+          f"{ff_down_proj_name}.weight"
+      )
+      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+          f"{ff_gate_proj_name}.weight"
+      )
+      if ff_config.use_bias:
+        converted_state[f"{prefix}.ff.w3.bias"] = state.pop(
+            f"{ff_up_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+            f"{ff_down_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+            f"{ff_gate_proj_name}.bias"
+        )
+
+    if self._names.pre_ff_norm is not None:
+      pre_ff_norm_name = self._names.pre_ff_norm.format(idx)
+      converted_state[f"{prefix}.ff.pre_ff_norm.weight"] = state.pop(
+          f"{pre_ff_norm_name}.weight"
+      )
+      if f"{pre_ff_norm_name}.bias" in state:
+        converted_state[f"{prefix}.ff.pre_ff_norm.bias"] = state.pop(
+            f"{pre_ff_norm_name}.bias"
+        )
+
+    if self._names.post_ff_norm is not None:
+      post_ff_norm_name = self._names.post_ff_norm.format(idx)
+      converted_state[f"{prefix}.ff.post_ff_norm.weight"] = state.pop(
+          f"{post_ff_norm_name}.weight"
+      )
+      if f"{post_ff_norm_name}.bias" in state:
+        converted_state[f"{prefix}.ff.post_ff_norm.bias"] = state.pop(
+            f"{post_ff_norm_name}.bias"
+        )
+
+  def _map_attention(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+  ):
+    prefix = f"transformer_blocks.{idx}"
+    attn_config = config.block_config(idx).attn_config
+    if self._names.attn_fused_qkv_proj:
+      fused_qkv_name = self._names.attn_fused_qkv_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = state.pop(
+          f"{fused_qkv_name}.weight"
+      )
+    else:
+      q_name = self._names.attn_query_proj.format(idx)
+      k_name = self._names.attn_key_proj.format(idx)
+      v_name = self._names.attn_value_proj.format(idx)
+      converted_state[f"{prefix}.atten_func.qkv_projection.weight"] = (
+          self._fuse_qkv(
+              attn_config,
+              state.pop(f"{q_name}.weight"),
+              state.pop(f"{k_name}.weight"),
+              state.pop(f"{v_name}.weight"),
+          )
+      )
+    if attn_config.qkv_use_bias:
+      if self._names.attn_fused_qkv_proj:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = state.pop(
+            f"{fused_qkv_name}.bias"
+        )
+      else:
+        converted_state[f"{prefix}.atten_func.qkv_projection.bias"] = (
+            self._fuse_qkv(
+                attn_config,
+                state.pop(f"{q_name}.bias"),
+                state.pop(f"{k_name}.bias"),
+                state.pop(f"{v_name}.bias"),
+            )
+        )
+
+    if self._names.attn_query_norm is not None:
+      attn_query_norm_name = self._names.attn_query_norm.format(idx)
+      converted_state[f"{prefix}.atten_func.query_norm.weight"] = state.pop(
+          f"{attn_query_norm_name}.weight"
+      )
+    if self._names.attn_key_norm is not None:
+      attn_key_norm_name = self._names.attn_key_norm.format(idx)
+      converted_state[f"{prefix}.atten_func.key_norm.weight"] = state.pop(
+          f"{attn_key_norm_name}.weight"
+      )
+
+    o_name = self._names.attn_output_proj.format(idx)
+    converted_state[f"{prefix}.atten_func.output_projection.weight"] = (
+        state.pop(f"{o_name}.weight")
+    )
+    if attn_config.output_proj_use_bias:
+      converted_state[f"{prefix}.atten_func.output_projection.bias"] = (
+          state.pop(f"{o_name}.bias")
+      )
+
+  def _map_norm(
+      self,
+      idx: int,
+      config: model_config.ModelConfig,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+  ):
+    prefix = f"transformer_blocks.{idx}"
+    if self._names.pre_attn_norm is not None:
+      pre_attn_norm_name = self._names.pre_attn_norm.format(idx)
+      converted_state[f"{prefix}.pre_atten_norm.weight"] = state.pop(
+          f"{pre_attn_norm_name}.weight"
+      )
+      if f"{pre_attn_norm_name}.bias" in state:
+        converted_state[f"{prefix}.pre_atten_norm.bias"] = state.pop(
+            f"{pre_attn_norm_name}.bias"
+        )
+
+    if self._names.post_attn_norm is not None:
+      post_attn_norm_name = self._names.post_attn_norm.format(idx)
+      converted_state[f"{prefix}.post_atten_norm.weight"] = state.pop(
+          f"{post_attn_norm_name}.weight"
+      )
+      if f"{post_attn_norm_name}.bias" in state:
+        converted_state[f"{prefix}.post_atten_norm.bias"] = state.pop(
+            f"{post_attn_norm_name}.bias"
+        )
+
+  def _fuse_qkv(
+      self,
+      attn_config: model_config.AttentionConfig,
+      q: torch.Tensor,
+      k: torch.Tensor,
+      v: torch.Tensor,
+  ) -> torch.Tensor:
+    if attn_config.qkv_fused_interleaved:
+      q_per_kv = attn_config.num_heads // attn_config.num_query_groups
+      qs = torch.split(q, attn_config.head_dim * q_per_kv)
+      ks = torch.split(k, attn_config.head_dim)
+      vs = torch.split(v, attn_config.head_dim)
+      cycled = [t for group in zip(qs, ks, vs) for t in group]
+      return torch.cat(cycled)
+    else:
+      return torch.cat([q, k, v], dim=0)

ai_edge_torch/generative/utilities/model_builder.py

@@ -0,0 +1,180 @@
+# 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."""
+
+import copy
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import builder
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.generative.layers import lora as lora_utils
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
+
+
+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"
+
+
+@dataclass
+class ExportConfig:
+  """Model generating configuration settings."""
+
+  # On prefill signatures, should the model produce logit output?
+  # When False, only decode signatures will produce output.
+  output_logits_on_prefill: bool = False
+
+
+class DecoderOnlyModel(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
+    )
+    if config.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)
+        for idx in range(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
+  def forward(
+      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[ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    _, seq_len = tokens.size()
+    assert self.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)
+
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.kv_cache_max]
+
+    return self._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[ExportConfig] = None,
+  ) -> dict[torch.Tensor, kv_utils.KVCache]:
+    """Forwards the model with input embeddings."""
+    assert len(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
+    if self.config.embedding_scale is not None:
+      x = x * self.config.embedding_scale
+
+    updated_kv_entries = []
+    for i, block in enumerate(self.transformer_blocks):
+      kv_entry = kv_cache.caches[i] if kv_cache else None
+      lora_adapter = lora.adapters[i] if lora else None
+      x, kv_entry = block(x, rope, mask, input_pos, kv_entry, lora_adapter)
+      if kv_entry:
+        updated_kv_entries.append(kv_entry)
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+
+    if export_config is not None:
+      if (
+          torch.numel(input_pos) > 1
+          and not export_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}
+
+
+def build_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=not config.lm_head_share_weight_with_embedding
+  )
+  transformer.eval()
+  return transformer