Weights Metrics

.gitignore

@@ -1,3 +1,15 @@
+*.pkl
+*.h5
+offload_folder/
+
+# Environment
+
+mergekit/bin/
+mergekit/share/
+mergekit/etc/
+merged/
+.vscode/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

examples/metrics-llama-1v2.yml

@@ -0,0 +1,8 @@
+models:
+  - model: huggyllama/llama-7b
+  - model: TheBloke/Llama-2-7B-fp16
+
+metric_method: all
+parameters:
+  intra_model_metrics: true
+dtype: float32

examples/metrics-llama-codev2.yml

@@ -0,0 +1,6 @@
+models:
+  - model: meta-llama/CodeLlama-7b-Python-hf
+  - model: meta-llama/Llama-2-7b-hf
+
+metric_method: all
+dtype: float32

examples/metrics-small.yml

@@ -0,0 +1,9 @@
+models:
+  - model: BEE-spoke-data/smol_llama-220M-GQA
+  - model: BEE-spoke-data/smol_llama-220M-openhermes
+
+metric_method: all
+parameters:
+  intra_model_metrics: true
+  inter_model_metrics: true
+dtype: float32

mergekit/_data/architectures/llama.json

@@ -21,21 +21,27 @@
             {
                 "name": "model.layers.${layer_index}.self_attn.q_proj.weight",
                 "input_space": "h_${layer_index}",
+                "num_attention_heads": "${num_attention_heads}",
                 "output_space": "attn_qk_${layer_index}"
             },
             {
                 "name": "model.layers.${layer_index}.self_attn.k_proj.weight",
                 "input_space": "h_${layer_index}",
-                "output_space": "attn_qk_${layer_index}"
+                "output_space": "attn_qk_${layer_index}",
+                "num_attention_heads": "${num_attention_heads}",
+                "num_key_value_heads": "${num_key_value_heads}"
             },
             {
                 "name": "model.layers.${layer_index}.self_attn.v_proj.weight",
                 "input_space": "h_${layer_index}",
-                "output_space": "attn_v_${layer_index}"
+                "output_space": "attn_v_${layer_index}",
+                "num_attention_heads": "${num_attention_heads}",
+                "num_key_value_heads": "${num_key_value_heads}"
             },
             {
                 "name": "model.layers.${layer_index}.self_attn.o_proj.weight",
                 "input_space": "attn_v_${layer_index}",
+                "num_attention_heads": "${num_attention_heads}",
                 "output_space": "post_attn_${layer_index}"
             },
             {

mergekit/_data/models_and_datasets.json

@@ -0,0 +1,12 @@
+{
+    "model_names": [
+        "Qwen/Qwen2-7B-Instruct",
+        "arcee-ai/qwen2-7b-math-tess",
+        "arcee-ai/qwen2-mmamo-2",
+        "arcee-ai/Qwen2-Merger"
+    ],
+    "dataset_names": [
+        "microsoft/orca-math-word-problems-200k",
+        "MuskumPillerum/General-Knowledge"
+    ]
+}

mergekit/_data/models_and_datasets.py

@@ -0,0 +1,45 @@
+import json
+from pathlib import Path
+
+def save_model_and_dataset(model_name, dataset_name):
+    models_and_datasets = presets().load()
+    if model_name not in models_and_datasets['model_names']:
+        models_and_datasets['model_names'].append(model_name)
+    if dataset_name not in models_and_datasets['dataset_names']:
+        models_and_datasets['dataset_names'].append(dataset_name)
+    presets().save(models_and_datasets['model_names'], models_and_datasets['dataset_names'])
+
+def model_and_dataset_to_index(model_name, dataset_name):
+    models_and_datasets = presets().load()
+    model_index = models_and_datasets['model_names'].index(model_name) if model_name in models_and_datasets['model_names'] else []
+    dataset_index = models_and_datasets['dataset_names'].index(dataset_name) if dataset_name in models_and_datasets['dataset_names'] else []
+
+    return model_index, dataset_index
+
+def index_to_model_and_dataset(model_index, dataset_index):
+    models_and_datasets = presets().load()
+    model_name = models_and_datasets['model_names'][model_index] if len(models_and_datasets['model_names']) > model_index else []
+    dataset_name = models_and_datasets['dataset_names'][dataset_index] if len(models_and_datasets['dataset_names']) > dataset_index else []
+    return model_name, dataset_name
+
+class presets():
+    def __init__(self):
+        self.FILE_PATH = Path(__file__).parent / 'models_and_datasets.json'
+
+    def load(self):
+        """Load the lists from the JSON file."""
+        if self.FILE_PATH.exists():
+            with open(self.FILE_PATH, 'r') as file:
+                data = json.load(file)
+                return data
+        print(f"File {self.FILE_PATH} does not exist or is empty.")
+        return {}
+
+    def save(self, model_names, dataset_names):
+        """Save the lists to the JSON file."""
+        data = {
+            'model_names': model_names,
+            'dataset_names': dataset_names
+        }
+        with open(self.FILE_PATH, 'w') as file:
+            json.dump(data, file, indent=4)

mergekit/architecture.py

@@ -43,6 +43,10 @@ class WeightInfo(BaseModel, frozen=True):
             List of alternative names for the weight, if applicable.
         force_dtype (Optional[str]):
             Mandatory dtype for the weight, if applicable.
+        num_key_value_heads (Optional[int]):
+            Number of key-value heads in the weight, relevant for GQA, if applicable.
+        num_attention_heads (Optional[int]):
+            Number of attention heads in the weight, if applicable.
     """
 
     name: str
@@ -53,6 +57,9 @@ class WeightInfo(BaseModel, frozen=True):
     aliases: Optional[Tuple[str, ...]] = None
     force_dtype: Optional[str] = None
 
+    num_key_value_heads: Union[int, str, None] = None
+    num_attention_heads: Union[int, str, None] = None
+
 
 class ProceduralSpaceInfo(BaseModel, frozen=True):
     """Defines a procedural space computed from one or more other spaces.
@@ -172,30 +179,29 @@ class JSONArchitectureDefinition(BaseModel, frozen=True):
 class TemplateWithArithmetic(string.Template):
     idpattern = r"(?a:[_a-z][_a-z0-9]*([+-]1)?)"
 
-
-def _template_substitution(
-    template: str, num_layers: int, layer_idx: Optional[int] = None
+def _generic_template_substitution(
+    template: str, match_str: str = "num_layers", replace_with: int = 0
 ) -> str:
-    if "{" not in template:
+    if f"{match_str}" not in template:
         return template
 
     substitutions = {
-        "num_layers": num_layers,
-        "num_layers+1": num_layers + 1,
-        "num_layers-1": num_layers - 1,
+        match_str: replace_with,
+        f"{match_str}+1": replace_with + 1,
+        f"{match_str}-1": replace_with - 1,
     }
 
-    if layer_idx is not None:
-        substitutions.update(
-            {
-                "layer_index": layer_idx,
-                "layer_index+1": layer_idx + 1,
-                "layer_index-1": layer_idx - 1,
-            }
-        )
-
     return TemplateWithArithmetic(template).substitute(substitutions)
 
+def _template_substitution(
+    template: str, substitute: Dict[str, Optional[int]]
+) -> str:
+    for key, value in substitute.items():
+        if value is not None:
+            template = _generic_template_substitution(template, key, value)
+
+    return template
+
 
 class JsonArchitectureInfo(ArchitectureInfo, BaseModel, frozen=True):
     definition: JSONArchitectureDefinition
@@ -206,18 +212,23 @@ def _substitute(
         config: PretrainedConfig,
         layer_idx: Optional[int] = None,
     ) -> Union[WeightInfo, ProceduralSpaceInfo]:
-        num_layers = self.num_layers(config)
-
-        obj_dict = item.model_dump(mode="json", exclude_unset=True)
+        substitute = {
+            "num_layers": self.num_layers(config),
+            "layer_index": layer_idx,
+            "num_attention_heads": getattr(config, "num_attention_heads") if getattr(config, "num_attention_heads") is not None else None,
+            "num_key_value_heads": getattr(config, "num_key_value_heads") if getattr(config, "num_key_value_heads") is not None else None,
+        }
+
+        obj_dict = item.model_dump(mode="json", exclude_unset=False)
         for key in obj_dict:
             if isinstance(obj_dict[key], str):
                 obj_dict[key] = _template_substitution(
-                    obj_dict[key], num_layers, layer_idx
+                    obj_dict[key], substitute
                 )
             elif isinstance(obj_dict[key], list):
                 obj_dict[key] = [
                     (
-                        _template_substitution(s, num_layers, layer_idx)
+                        _template_substitution(s, substitute)
                         if isinstance(s, str)
                         else s
                     )

mergekit/config.py

@@ -83,7 +83,8 @@ class OutputSliceDefinition(BaseModel):
 
 
 class MergeConfiguration(BaseModel):
-    merge_method: str
+    merge_method: Optional[str] = None
+    metric_method: Optional[str] = None
     slices: Optional[List[OutputSliceDefinition]] = None
     models: Optional[List[InputModelDefinition]] = None
     parameters: Optional[Dict[str, ParameterSetting]] = None
@@ -114,6 +115,14 @@ def validate_inputs(self):
         if ((not self.slices) and (not self.models)) or (self.slices and self.models):
             raise RuntimeError("Must specify either output slices or models to merge")
         return self
+
+    @model_validator(mode="after")
+    def validate_methods(self):
+        if not self.merge_method and not self.metric_method:
+            raise ValueError("Either 'merge_method' or 'metric_method' must be provided.")
+        if self.merge_method and self.metric_method:
+            raise ValueError("Only one of 'merge_method' or 'metric_method' can be provided, not both.")
+        return self
 
     @model_validator(mode="after")
     def validate_tokenizer(self):

mergekit/graph.py

@@ -37,6 +37,7 @@ class Task(ABC, BaseModel, Generic[ValueT], frozen=True):
     Abstract base class representing a task in a computational graph.
 
     This class should be extended to define specific tasks. Each task can have arguments (dependencies) and a defined execution strategy.
+    Pydantic BaseModel requires that all attributes are defined in the class initialisation, and cannot be changed after. 
 
     Attributes:
         Generic[ValueT] (TypeVar): The type of the value that the task returns upon execution.
@@ -106,7 +107,6 @@ def uses_accelerator(self) -> bool:
         """
         return False
 
-
 class Executor:
     """
     Schedules and executes a set of tasks and their dependencies.
@@ -241,13 +241,20 @@ def _make_schedule(self, targets: List[Task]) -> List[Task]:
             # they will be included in the final schedule
             edge_tups.append((Executor.DUMMY_TASK_VALUE, task))
 
+        def _pad_numbers(s):
+            parts = s.split('.')
+            for i, part in enumerate(parts):
+                if part.isdigit():
+                    parts[i] = part.zfill(3)
+            return '.'.join(parts)
+
         def _compare_key(task: Union[Task, str]):
             if task == Executor.DUMMY_TASK_VALUE:
                 return ("", 0)
-            return (
-                task.group_label() or "",
-                -task.priority(),
-            )
+            group_label = task.group_label() or ""
+            padded_label = _pad_numbers(group_label)
+            priority = -task.priority()
+            return (padded_label, priority)
 
         graph = networkx.DiGraph(edge_tups)
         res = [

mergekit/merge.py

@@ -92,11 +92,17 @@ def run_merge(
         storage_device="cuda" if options.low_cpu_memory else "cpu",
     )
 
+    metrics_out = []
     tokenizer = None
     for _task, value in exec.run(quiet=options.quiet):
+        if merge_config.metric_method is not None:
+            metrics_out.append((_task, value))
         if isinstance(value, TokenizerInfo):
             tokenizer = value.tokenizer
 
+    if metrics_out:
+        return metrics_out
+
     if tokenizer:
         _update_config_vocab(cfg_out, tokenizer)
 

mergekit/metric_methods/__init__.py

@@ -0,0 +1,29 @@
+# Copyright (C) 2024 Charles O. Goddard
+#
+# This software is free software: you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This software is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program. If not, see http://www.gnu.org/licenses/.
+
+from mergekit.metric_methods.base import MetricMethod
+from mergekit.metric_methods.all_metrics import AllMetric
+
+
+def get(method: str) -> MetricMethod:
+    if method == "all":
+        return AllMetric()
+    raise RuntimeError(f"Unimplemented metric method {method}")
+
+
+__all__ = [
+    "MetricMethod",
+    "get",
+]