Add --num-threads option

mergekit/options.py

@@ -2,10 +2,13 @@
 # SPDX-License-Identifier: BUSL-1.1
 
 import functools
+import logging
 import typing
 from typing import Any, Callable, Optional, Union
 
 import click
+import torch
+import transformers
 from click.core import Context, Parameter
 from pydantic import BaseModel
 
@@ -26,9 +29,19 @@ class MergeOptions(BaseModel, frozen=True):
     lazy_unpickle: bool = False
     write_model_card: bool = True
     safe_serialization: bool = True
+    verbose: bool = False
     quiet: bool = False
     read_to_gpu: bool = False
     multi_gpu: bool = False
+    num_threads: Optional[int] = None
+
+    def apply_global_options(self):
+        logging.basicConfig(level=logging.INFO if self.verbose else logging.WARNING)
+        if self.random_seed is not None:
+            transformers.trainer_utils.set_seed(self.random_seed)
+        if self.num_threads is not None:
+            torch.set_num_threads(self.num_threads)
+            torch.set_num_interop_threads(self.num_threads)
 
 
 OPTION_HELP = {
@@ -48,6 +61,8 @@ class MergeOptions(BaseModel, frozen=True):
     "quiet": "Suppress progress bars and other non-essential output",
     "read_to_gpu": "Read model weights directly to GPU",
     "multi_gpu": "Use multi-gpu parallel graph execution engine",
+    "num_threads": "Number of threads to use for parallel CPU operations",
+    "verbose": "Enable verbose logging",
 }
 
 
@@ -88,10 +103,15 @@ def wrapper(*args, **kwargs):
             arg_str = f"--{arg_name}/--no-{arg_name}"
         else:
             arg_str = f"--{arg_name}"
+        param_decls = [arg_str]
+        if field_name == "verbose":
+            param_decls = ["--verbose/--no-verbose", "-v"]
+        if field_name == "num_threads":
+            param_decls = ["--num-threads", "-j"]
 
         help_str = OPTION_HELP.get(field_name, None)
         wrapper = click.option(
-            arg_str,
+            *param_decls,
             type=field_type,
             default=info.default,
             help=help_str,

mergekit/scripts/extract_lora.py

@@ -85,12 +85,6 @@
     multiple=True,
     help="Include modules matching the specified regex",
 )
-@click.option(
-    "--verbose",
-    "-v",
-    is_flag=True,
-    help="Verbose logging",
-)
 @click.option(
     "--sv-epsilon",
     type=float,
@@ -109,11 +103,10 @@ def main(
     modules_to_save: List[str],
     exclude_regexes: List[str],
     include_regexes: List[str],
-    verbose: bool,
     sv_epsilon: float,
     merge_options: MergeOptions,
 ):
-    logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO)
+    merge_options.apply_global_options()
 
     if not modules_to_save:
         modules_to_save = []

mergekit/scripts/merge_raw_pytorch.py

@@ -236,6 +236,8 @@ def main(
     Uses similar configuration syntax to `mergekit-yaml`, minus the
     `slices` sections. Each input model should be the path on disk to a
     pytorch pickle file or safetensors file."""
+    merge_options.apply_global_options()
+
     with open(config_path, "r", encoding="utf-8") as file:
         config_source = file.read()
 

mergekit/scripts/moe.py

@@ -180,12 +180,12 @@ def main(
     load_in_4bit: bool,
     load_in_8bit: bool,
     device: str,
-    merge_options: MergeOptions,
     verbose: bool,
     i_understand_this_is_not_useful_without_training: bool,
+    merge_options: MergeOptions,
 ):
     """Create a Mixture of Experts model by combining the pretrained weights of multiple models."""
-    logging.basicConfig(level=logging.INFO if verbose else logging.WARNING)
+    merge_options.apply_global_options()
 
     if merge_options.cuda:
         logging.warning(

mergekit/scripts/multimerge.py

@@ -75,9 +75,6 @@ def execute(self, **kwargs):
     required=True,
     help="Directory to store intermediate merges",
 )
-@click.option(
-    "--verbose", "-v", type=bool, default=False, is_flag=True, help="Verbose logging"
-)
 @click.option(
     "--lazy/--no-lazy",
     default=True,
@@ -88,7 +85,6 @@ def main(
     config_file: str,
     intermediate_dir: str,
     out_path: Optional[str],
-    verbose: bool,
     lazy: bool,
     merge_options: MergeOptions,
 ):
@@ -102,7 +98,7 @@ def main(
     Any merge configuration with a `name` field will be saved to this
     directory. If an unnamed merge configuration is present, it will be
     saved to `out_path` (which is required in this case)."""
-    logging.basicConfig(level=logging.INFO if verbose else logging.WARNING)
+    merge_options.apply_global_options()
     os.makedirs(intermediate_dir, exist_ok=True)
 
     with open(config_file, "r", encoding="utf-8") as file:

mergekit/scripts/run_yaml.py

@@ -1,7 +1,6 @@
 # Copyright (C) 2025 Arcee AI
 # SPDX-License-Identifier: BUSL-1.1
 
-import logging
 
 import click
 import yaml
@@ -14,17 +13,13 @@
 @click.command("mergekit-yaml")
 @click.argument("config_file")
 @click.argument("out_path")
-@click.option(
-    "--verbose", "-v", type=bool, default=False, is_flag=True, help="Verbose logging"
-)
 @add_merge_options
 def main(
     merge_options: MergeOptions,
     config_file: str,
     out_path: str,
-    verbose: bool,
 ):
-    logging.basicConfig(level=logging.INFO if verbose else logging.WARNING)
+    merge_options.apply_global_options()
 
     with open(config_file, "r", encoding="utf-8") as file:
         config_source = file.read()