Add support for GPTBigCodeForCausalLM

mergekit/architecture.py

@@ -37,35 +37,53 @@ def layer_weight_formats(self) -> List[str]:
         ...
 
     @abstractmethod
-    def embed_weights(self) -> List[str]:
+    def input_embed(self) -> Optional[str]:
+        ...
+
+    @abstractmethod
+    def output_embed(self) -> Optional[str]:
         ...
 
+    def embed_weights(self) -> List[str]:
+        return [self.input_embed(), self.output_embed()]
+
     def num_layers(self, config: PretrainedConfig) -> int:
         return config.num_hidden_layers
 
     def num_layers_config_key(self) -> str:
         """Key in config that represents number of layers"""
         return "num_hidden_layers"
 
+    def possible_name_prefix(self) -> Optional[str]:
+        """
+        Return a prefix that is allowed but not required for tensor names.
+        """
+        return None
+
 
 class StaticTensorNames(ArchitectureInfo, BaseModel, frozen=True):
-    name: str
+    names: List[str]
 
     pre_weight_names: List[str]  # weights applied before first layer
     post_weight_names: List[str]  # weights applied after last layer
-    embed_weight_names: List[str]  # weights for embed/lm_head
+    input_embed_name: str
+    output_embed_name: Optional[str]
     layer_prefix_format: str
     layer_weight_suffixes: List[str]
     num_layers_key: Optional[str] = None
+    allowable_prefix: Optional[str] = None
 
     def pre_weights(self) -> List[str]:
         return self.pre_weight_names
 
     def post_weights(self) -> List[str]:
         return self.post_weight_names
 
-    def embed_weights(self) -> List[str]:
-        return self.embed_weight_names
+    def input_embed(self) -> Optional[str]:
+        return self.input_embed_name
+
+    def output_embed(self) -> Optional[str]:
+        return self.output_embed_name
 
     def layer_weight_formats(self) -> List[str]:
         res = []
@@ -81,12 +99,16 @@ def num_layers_config_key(self) -> str:
     def num_layers(self, config: PretrainedConfig) -> int:
         return getattr(config, self.num_layers_config_key())
 
+    def possible_name_prefix(self) -> Optional[str]:
+        return self.allowable_prefix
+
 
 LLAMA_INFO = StaticTensorNames(
-    name="LlamaForCausalLM",
+    names=["LlamaForCausalLM", "LLaMAForCausalLM"],
     pre_weight_names=["model.embed_tokens.weight"],
     post_weight_names=["model.norm.weight", "lm_head.weight"],
-    embed_weight_names=["model.embed_tokens.weight", "lm_head.weight"],
+    input_embed_name="model.embed_tokens.weight",
+    output_embed_name="lm_head.weight",
     layer_prefix_format="model.layers.{idx}",
     layer_weight_suffixes=[
         "input_layernorm.weight",
@@ -102,34 +124,36 @@ def num_layers(self, config: PretrainedConfig) -> int:
 )
 
 MISTRAL_INFO = StaticTensorNames(
-    name="MistralForCausalLM",
+    names=["MistralForCausalLM"],
     # lol
-    **LLAMA_INFO.model_dump(exclude=["name"]),
+    **LLAMA_INFO.model_dump(exclude=["names"]),
 )
 
 
 STABLELM_INFO = StaticTensorNames(
-    name="StableLMEpochForCausalLM",
+    names=["StableLMEpochForCausalLM"],
     post_weight_names=LLAMA_INFO.post_weight_names + ["model.norm.bias"],
     layer_weight_suffixes=LLAMA_INFO.layer_weight_suffixes
     + [
         "input_layernorm.bias",
         "post_attention_layernorm.bias",
     ],
     **LLAMA_INFO.model_dump(
-        exclude=["name", "layer_weight_suffixes", "post_weight_names"]
+        exclude=["names", "layer_weight_suffixes", "post_weight_names"]
     ),
 )
 
+
 GPT_NEOX_INFO = StaticTensorNames(
-    name="GPTNeoXForCausalLM",
+    names=["GPTNeoXForCausalLM"],
     pre_weight_names=["gpt_neox.embed_in.weight"],
     post_weight_names=[
         "gpt_neox.final_layer_norm.bias",
         "gpt_neox.final_layer_norm.weight",
         "embed_out.weight",
     ],
-    embed_weight_names=["gpt_neox.embed_in.weight", "embed_out.weight"],
+    input_embed_name="gpt_neox.embed_in.weight",
+    output_embed_name="embed_out.weight",
     layer_prefix_format="gpt_neox.layers.{idx}",
     layer_weight_suffixes=sum(
         (
@@ -148,11 +172,13 @@ def num_layers(self, config: PretrainedConfig) -> int:
     + ["attention.bias", "attention.masked_bias", "attention.rotary_emb.inv_freq"],
 )
 
+
 GPT2_INFO = StaticTensorNames(
-    name="GPT2LMHeadModel",
+    names=["GPT2Model", "GPT2LMHeadModel"],
     pre_weight_names=["wte.weight", "wpe.weight"],
-    post_weight_names=["ln_f.weight", "ln_f.bias"],
-    embed_weight_names=["wte.weight"],
+    input_embed_name="wte.weight",
+    output_embed_name="lm_head.weight",
+    post_weight_names=["ln_f.weight", "ln_f.bias", "lm_head.weight"],
     layer_prefix_format="h.{idx}",
     layer_weight_suffixes=[
         "attn.c_attn.weight",
@@ -167,32 +193,46 @@ def num_layers(self, config: PretrainedConfig) -> int:
         "mlp.c_proj.bias",
         "mlp.c_fc.weight",
         "mlp.c_fc.bias",
-        "mlp.c_proj.weight",
-        "mlp.c_proj.bias",
     ],
     num_layers_key="n_layer",
+    allowable_prefix="transformer.",
 )
 
 GPT2_SEQCLASS_INFO = StaticTensorNames(
-    name="GPT2ForSequenceClassification",
+    names=["GPT2ForSequenceClassification"],
     pre_weight_names=["transformer.wte.weight", "transformer.wpe.weight"],
+    input_embed_name="transformer.wte.weight",
+    output_embed_name=None,
     post_weight_names=[
         "transformer.ln_f.weight",
         "transformer.ln_f.bias",
         "score.weight",
     ],
     layer_prefix_format="transformer.h.{idx}",
-    embed_weight_names=GPT2_INFO.embed_weight_names,
     layer_weight_suffixes=GPT2_INFO.layer_weight_suffixes,
     num_layers_key=GPT2_INFO.num_layers_key,
 )
 
+GPT_BIGCODE_INFO = StaticTensorNames(
+    names=["GPTBigCodeForCausalLM"],
+    output_embed_name="lm_head.weight",
+    post_weight_names=[
+        "transformer.ln_f.weight",
+        "transformer.ln_f.bias",
+        "lm_head.weight",
+    ],
+    **GPT2_SEQCLASS_INFO.model_dump(
+        exclude=["names", "output_embed_name", "post_weight_names"]
+    ),
+)
+
 
 QWEN_INFO = StaticTensorNames(
-    name="QWenLMHeadModel",
+    names=["QWenLMHeadModel"],
     pre_weight_names=["transformer.wte.weight"],
+    input_embed_name="transformer.wte.weight",
+    output_embed_name="lm_head.weight",
     post_weight_names=["transformer.ln_f.weight", "lm_head.weight"],
-    embed_weight_names=["transformer.wte.weight", "lm_head.weight"],
     layer_prefix_format="transformer.h.{idx}",
     layer_weight_suffixes=[
         "attn.c_attn.bias",
@@ -206,8 +246,9 @@ def num_layers(self, config: PretrainedConfig) -> int:
     ],
 )
 
+
 CHATGLM_INFO = StaticTensorNames(
-    name="ChatGLMModel",
+    names=["ChatGLMModel"],
     pre_weight_names=[
         "transformer.embedding.word_embeddings.weight",
         "transformer.rotary_pos_emb.inv_freq",
@@ -216,10 +257,8 @@ def num_layers(self, config: PretrainedConfig) -> int:
         "transformer.encoder.final_layernorm.weight",
         "transformer.output_layer.weight",
     ],
-    embed_weight_names=[
-        "transformer.embedding.word_embeddings.weight",
-        "transformer.output_layer.weight",
-    ],
+    input_embed_name="transformer.embedding.word_embeddings.weight",
+    output_embed_name="transformer.output_layer.weight",
     layer_prefix_format="transformer.encoder.layers.{idx}",
     layer_weight_suffixes=[
         "input_layernorm.weight",
@@ -262,6 +301,12 @@ def embed_weights(self) -> List[str]:
             f"layers.{fake_layer_idx}.linear.bias",
         ]
 
+    def input_embed(self) -> Optional[str]:
+        return self.embed_weights()[0]
+
+    def output_embed(self) -> Optional[str]:
+        return self.embed_weights()[1]
+
     def layer_weight_formats(self) -> List[str]:
         return [
             ("layers.{idx}." + suffix)
@@ -288,15 +333,16 @@ def num_layers_config_key(self) -> str:
 
 
 PHI2_INFO = StaticTensorNames(
-    name="PhiForCausalLM",
+    names=["PhiForCausalLM"],
     pre_weight_names=["transformer.embd.wte.weight"],
     post_weight_names=[
         "lm_head.linear.bias",
         "lm_head.linear.weight",
         "lm_head.ln.bias",
         "lm_head.ln.weight",
     ],
-    embed_weight_names=["lm_head.linear.weight", "transformer.embd.wte.weight"],
+    input_embed_name="transformer.embd.wte.weight",
+    output_embed_name="lm_head.linear.weight",
     layer_prefix_format="transformer.h.{idx}",
     layer_weight_suffixes=[
         "ln.bias",
@@ -331,9 +377,10 @@ def get_architecture_info(config: PretrainedConfig) -> StaticTensorNames:
         CHATGLM_INFO,
         STABLELM_INFO,
         PHI2_INFO,
+        GPT_BIGCODE_INFO,
     ]
     for arch in supported:
-        if arch.name == arch_name:
+        if arch_name in arch.names:
             return arch
 
     raise RuntimeError(f"Unsupported architecture {arch_name}")

mergekit/graph.py

@@ -36,6 +36,7 @@
 from pydantic import BaseModel
 from typing_extensions import Protocol
 
+from mergekit.architecture import ArchitectureInfo, get_architecture_info
 from mergekit.common import ModelReference
 from mergekit.io import LazyTensorLoader, TensorWriter
 
@@ -91,11 +92,32 @@ def generate_rule(self, component: TensorReference) -> Optional[Operation]:
         ...
 
 
+def _remap_tensor_name(
+    name: str, arch_info: ArchitectureInfo, tensor_paths: Dict[str, str]
+) -> Optional[str]:
+    if name in tensor_paths:
+        return name
+    pp = arch_info.possible_name_prefix()
+    if pp and (pp + name) in tensor_paths:
+        return pp + name
+
+    e_out = arch_info.output_embed()
+    e_in = arch_info.input_embed()
+    if e_out and name == e_out:
+        # if output embedding tensor is missing, assume weights are tied
+        if e_in in tensor_paths:
+            return e_in
+        if pp and (pp + e_in) in tensor_paths:
+            return pp + e_in
+    return None
+
+
 class LoadTensorRule(ProceduralRule):
     """Rule for loading tensors from input models."""
 
     model: ModelReference
     tensor_paths: Dict[str, str]
+    dtype: torch.dtype
 
     def __init__(
         self,
@@ -117,6 +139,7 @@ def can_generate_rule(self, component: TensorReference) -> bool:
     def generate_rule(self, component: TensorReference) -> Operation:
         if not self.can_generate_rule(component):
             return None
+
         return Operation(
             function="load_tensor",
             inputs=[],
@@ -192,6 +215,8 @@ class Executor:
     loaders: Dict[ModelReference, LazyTensorLoader]
     targets: List[TensorReference]
     operations: Dict[str, OperationProtocol]
+    arch_info: ArchitectureInfo
+    real_tensor_names: Dict[TensorReference, str]
     low_cpu_memory: bool = False
 
     def __init__(
@@ -211,6 +236,10 @@ def __init__(
         if lora_cache_dir is None and transformers_cache_dir is not None:
             lora_cache_dir = transformers_cache_dir
 
+        self.arch_info = get_architecture_info(
+            models[0].config(trust_remote_code=trust_remote_code)
+        )
+
         self.targets = targets
         self.loaders = {
             ref: LazyTensorLoader(
@@ -225,7 +254,11 @@ def __init__(
         }
         for model, loader in self.loaders.items():
             rules.procedural.append(
-                LoadTensorRule(model, loader.index.tensor_paths, dtype=dtype)
+                LoadTensorRule(
+                    model,
+                    loader.index.tensor_paths,
+                    dtype=dtype,
+                )
             )
 
         if operations is None:
@@ -234,6 +267,19 @@ def __init__(
         self.rules = rules
         self.cuda = cuda
         self.low_cpu_memory = low_cpu_memory
+        self.real_tensor_names = {}
+
+    def get_real_name(self, tensor_ref: TensorReference) -> str:
+        if not tensor_ref.model:
+            return tensor_ref.key
+
+        if tensor_ref not in self.real_tensor_names:
+            self.real_tensor_names[tensor_ref] = _remap_tensor_name(
+                tensor_ref.key,
+                self.arch_info,
+                self.loaders[tensor_ref.model].index.tensor_paths,
+            )
+        return self.real_tensor_names[tensor_ref]
 
     def run(self, out_path: str, max_shard_size: int, clone_tensors: bool = False):
         """
@@ -327,9 +373,10 @@ def _load_tensor(self, operation: Operation):
         """Load a tensor from an input model."""
         assert operation.function == "load_tensor"
 
-        res = self.loaders[operation.kwargs["model"]].get_tensor(
-            operation.kwargs["key"]
+        tensor_ref = TensorReference(
+            model=operation.kwargs["model"], key=operation.kwargs["key"]
         )
+        res = self.loaders[tensor_ref.model].get_tensor(self.get_real_name(tensor_ref))
         if operation.kwargs["dtype"]:
             res = res.to(dtype=operation.kwargs["dtype"])
 
@@ -346,8 +393,9 @@ def _compare_key(self, ref: TensorReference):
         at any given time.
         """
         if ref.model:
+            tensor_name = self.get_real_name(ref)
             shard_key = _normalized_shard_name(
-                self.loaders[ref.model].index.tensor_paths[ref.key]
+                self.loaders[ref.model].index.tensor_paths[tensor_name]
             )
         else:
             shard_key = ""
@@ -389,7 +437,12 @@ def _visit(node: TensorReference):
             if node in ops:
                 return
 
-            operation = self.rules.get(node)
+            real_name = self.get_real_name(node)
+            if not real_name:
+                raise RuntimeError(
+                    f"Model {node.model} does not contain tensor {node.key}"
+                )
+            operation = self.rules.get(TensorReference(model=node.model, key=real_name))
             if not operation:
                 raise RuntimeError(f"No rule to produce {node}")
             ops[node] = operation