mergekit-evolve QoL tweaks (#298)

A few quality of life tweaks for `mergekit-evolve`:
1. Make flash attention optional
2. Write out the final merge instead of just the config
(`--no-save-final-model` to disable)
3. Add option to not reshard input models (`--no-reshard`) for lower
disk use at cost of merge speed

Author: cg123

mergekit/evo/actors.py

@@ -27,6 +27,7 @@
 import ray.util.scheduling_strategies
 import torch
 import transformers
+from transformers.utils import is_flash_attn_2_available
 
 try:
     import vllm
@@ -39,7 +40,11 @@
 from mergekit.evo.config import EvolMergeConfiguration
 from mergekit.evo.genome import ModelGenome
 from mergekit.evo.helpers import _eval_model, evaluate_model, merge_model
-from mergekit.evo.monkeypatch import NoInit, monkeypatch_lmeval_shuffle
+from mergekit.evo.monkeypatch import (
+    NoInit,
+    monkeypatch_lmeval_shuffle,
+    monkeypatch_lmeval_vllm,
+)
 from mergekit.graph import Executor
 from mergekit.io.tasks import LoaderCache, ReturnTensor
 from mergekit.merge import _model_out_config
@@ -72,6 +77,7 @@ def __init__(
             monkeypatch_lmeval_shuffle()
 
         # monkeypatch_tqdm()
+        monkeypatch_lmeval_vllm()
 
 
 @ray.remote(num_cpus=1, num_gpus=1.0)
@@ -164,13 +170,18 @@ def _maybe_init_model(self, config: MergeConfiguration):
             if not different:
                 return
 
+        model_kwargs = {
+            "trust_remote_code": self.merge_options.trust_remote_code,
+            "torch_dtype": torch.bfloat16,
+        }
+        if is_flash_attn_2_available():
+            model_kwargs["attn_implementation"] = "flash_attention_2"
+
         with NoInit():
             inner_model = (
                 transformers.AutoModelForCausalLM.from_config(
                     cfg_out,
-                    trust_remote_code=self.merge_options.trust_remote_code,
-                    attn_implementation="flash_attention_2",
-                    torch_dtype=torch.bfloat16,
+                    **model_kwargs,
                 )
                 .bfloat16()
                 .cuda()
@@ -203,11 +214,14 @@ def _maybe_init_model(self, config: MergeConfiguration):
                     max_model_len = 8192
                     logging.warn(f"Clipping sequence length to {max_model_len}")
 
+                mem_util = (
+                    0.7 if self.merge_options.cuda else 0.9
+                )  # reduce memory usage if we're also using cuda for the merge
                 self.model = lm_eval.models.vllm_causallms.VLLM(
                     pretrained=tempdir,
                     batch_size=self.batch_size or "auto",
                     max_model_len=max_model_len,
-                    gpu_memory_utilization=0.7,  # can't do 0.9 because the merge will OOM
+                    gpu_memory_utilization=mem_util,
                     dtype="bfloat16",
                     device="cuda",
                     trust_remote_code=self.merge_options.trust_remote_code,
@@ -279,6 +293,7 @@ def evaluate(self, genotype: torch.Tensor) -> dict:
             num_fewshot=self.config.num_fewshot,
             limit=self.config.limit,
             task_manager=self.task_manager,
+            batch_size=self.batch_size,
         )
 
     def evaluate_genotype(

mergekit/evo/helpers.py

@@ -30,6 +30,7 @@
 
 from mergekit.evo.config import TaskConfiguration
 from mergekit.evo.genome import ModelGenome
+from mergekit.evo.monkeypatch import monkeypatch_lmeval_vllm
 from mergekit.merge import run_merge
 from mergekit.options import MergeOptions
 
@@ -68,6 +69,7 @@ def evaluate_model(
     task_manager: Optional[lm_eval.tasks.TaskManager] = None,
 ) -> float:
     # monkeypatch_tqdm()
+    monkeypatch_lmeval_vllm()
     try:
         model_args = {
             "pretrained": merged_path,

mergekit/evo/monkeypatch.py

@@ -100,6 +100,15 @@ def _patch_lm_eval():
         mergekit.tokenizer.tqdm = fake_module
 
 
+def monkeypatch_lmeval_vllm():
+    # HACK: fix crash on some tasks due to unset AUTO_MODEL_CLASS for vLLM
+    import lm_eval.models.vllm_causallms
+
+    lm_eval.models.vllm_causallms.VLLM.AUTO_MODEL_CLASS = (
+        transformers.AutoModelForCausalLM
+    )
+
+
 class NoInit:
     def __enter__(self):
         def noop(*args, **kwargs):

mergekit/evo/strategy.py

@@ -89,6 +89,7 @@ def __init__(
                     self.merge_options,
                     model_storage_path=self.model_storage_path,
                     vllm=vllm,
+                    batch_size=self.batch_size,
                     task_manager=self.task_manager,
                 )
                 for _ in range(self.num_gpus)

mergekit/scripts/evolve.py

@@ -21,6 +21,7 @@
 import cma
 import numpy as np
 import pandas
+import ray
 import torch
 import tqdm
 import transformers
@@ -44,6 +45,7 @@
     BufferedRayEvaluationStrategy,
     SerialEvaluationStrategy,
 )
+from mergekit.merge import run_merge
 from mergekit.options import MergeOptions
 
 
@@ -93,6 +95,18 @@
     default=False,
     help="Allow benchmark tasks as objectives",
 )
+@click.option(
+    "--save-final-model/--no-save-final-model",
+    is_flag=True,
+    default=True,
+    help="Save the final merged model",
+)
+@click.option(
+    "--reshard/--no-reshard",
+    is_flag=True,
+    default=True,
+    help="Convert models to single-shard safetensors for faster merge",
+)
 def main(
     genome_config_path: str,
     max_fevals: int,
@@ -112,6 +126,8 @@ def main(
     wandb_entity: Optional[str],
     task_search_path: List[str],
     allow_benchmark_tasks: bool,
+    save_final_model: bool,
+    reshard: bool,
 ):
     config = EvolMergeConfiguration.model_validate(
         yaml.safe_load(open(genome_config_path, "r", encoding="utf-8"))
@@ -146,21 +162,28 @@ def main(
     )
 
     # convert models to single-shard safetensors
-    resharded_models = []
-    resharded_base = None
-    for model in tqdm.tqdm(config.genome.models, desc="Resharding models"):
-        resharded_models.append(
-            _reshard_model(
-                model, storage_path, merge_options.lora_merge_cache, trust_remote_code
+    if reshard:
+        resharded_models = []
+        resharded_base = None
+        for model in tqdm.tqdm(config.genome.models, desc="Resharding models"):
+            resharded_models.append(
+                _reshard_model(
+                    model,
+                    storage_path,
+                    merge_options.lora_merge_cache,
+                    trust_remote_code,
+                )
             )
-        )
-    if config.genome.base_model is not None:
-        resharded_base = _reshard_model(
-            config.genome.base_model,
-            storage_path,
-            merge_options.lora_merge_cache,
-            trust_remote_code,
-        )
+        if config.genome.base_model is not None:
+            resharded_base = _reshard_model(
+                config.genome.base_model,
+                storage_path,
+                merge_options.lora_merge_cache,
+                trust_remote_code,
+            )
+    else:
+        resharded_models = config.genome.models
+        resharded_base = config.genome.base_model
 
     genome = ModelGenome(
         ModelGenomeDefinition.model_validate(
@@ -289,16 +312,22 @@ def parallel_evaluate(x: List[np.ndarray]) -> List[float]:
         )
         xbest_cost = es.result.fbest
     except KeyboardInterrupt:
-        pass
+        ray.shutdown()
 
     print("!!! OPTIMIZATION COMPLETE !!!")
     print(f"Best cost: {xbest_cost:.4f}")
     print()
 
-    best_config = genome.genotype_merge_config(xbest)
+    # save the best merge configuration using original model references
+    genome_pretty = ModelGenome(config.genome, trust_remote_code=trust_remote_code)
+    best_config = genome_pretty.genotype_merge_config(xbest)
     print("Best merge configuration:")
     print(best_config.to_yaml())
 
+    if save_final_model:
+        print("Saving final model...")
+        run_merge(best_config, os.path.join(storage_path, "final_model"), merge_options)
+
 
 def _reshard_model(
     model: ModelReference, storage_path: str, merge_cache: str, trust_remote_code: bool
@@ -322,6 +351,7 @@ def _reshard_model(
         revision=merged.model.revision,
         trust_remote_code=trust_remote_code,
         torch_dtype=torch.bfloat16,
+        cache_dir=os.path.join(storage_path, "transformers_cache"),
     )
     model_hf.save_pretrained(
         out_path, safe_serialization=True, out_shard_size=1_000_000_000_000

mergekit/scripts/extract_lora.py

@@ -9,7 +9,7 @@
 from peft.tuners.lora import QuantLinear
 from safetensors.torch import save_file
 from tqdm import tqdm
-from transformers import AutoConfig, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM
 from transformers.modeling_utils import PreTrainedModel
 
 from mergekit.card import generate_card_lora
@@ -216,8 +216,6 @@ def main(
     base_model_ref = ModelReference.parse(base_model)
     finetuned_model_ref = ModelReference.parse(finetuned_model)
 
-    base_model_config = AutoConfig.from_pretrained(base_model_ref.model.path)
-
     linear_module_names = get_linear_module_names(base_model_ref.model.path)
     finetuned_model_linear_module_names = get_linear_module_names(
         finetuned_model_ref.model.path

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = ["black~=24.2.0", "isort~=5.13.2", "pre-commit~=3.6.2"]
 test = ["pytest~=8.0.1"]
-evolve = ["ray", "cma", "lm_eval", "flash-attn", "wandb"]
+evolve = ["ray", "cma", "lm_eval", "wandb"]
 vllm = ["vllm==0.3.2", "lm_eval[vllm]"]
 
 [project.urls]