Bring over graph/multigpu changes from tokensurgeon experiments

Author: cg123

mergekit/architecture/__init__.py

@@ -22,7 +22,7 @@
 if TYPE_CHECKING:
     from mergekit.config import MergeConfiguration
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture]:
@@ -44,11 +44,11 @@ def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture
         for c in candidates:
             if c.expected_model_type == config.model_type:
                 return c
-        logger.warning(
+        LOG.warning(
             f"Multiple architectures for {arch_name}, none match model type {config.model_type}"
         )
 
-    logger.warning(f"No JSON architecture found for {arch_name}")
+    LOG.warning(f"No JSON architecture found for {arch_name}")
     return None
 
 

mergekit/architecture/auto.py

@@ -21,7 +21,7 @@
 
 RE_LAYER_INDEX = re.compile(r"\.(\d+)\.")
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 def get_model_tensor_names(model: ModelReference, options: MergeOptions) -> List[str]:

mergekit/evo/actors.py

@@ -41,7 +41,7 @@
 from mergekit.options import MergeOptions
 from mergekit.plan import MergePlanner
 
-logger = logging.getLogger(__name__)
+LOG = logging.getLogger(__name__)
 
 
 class MergeActorBase:
@@ -91,18 +91,18 @@ def evaluate_genotype(
     ) -> dict:
         gc.collect()
         torch.cuda.empty_cache()
-        logger.info("Merging model")
+        LOG.info("Merging model")
         merged_path = merge_model(
             genotype, self.genome, self.model_storage_path, self.merge_options
         )
         if not merged_path:
-            logger.error("Model merge failed")
+            LOG.error("Model merge failed")
             return {"score": None, "results": None}
 
         model_kwargs = {}
         if self.quantization_config is not None:
             model_kwargs["quantization_config"] = self.quantization_config
-        logger.info(f"Model merged to {merged_path}")
+        LOG.info(f"Model merged to {merged_path}")
         return evaluate_model(
             merged_path,
             self.config.tasks,
@@ -167,7 +167,7 @@ def _maybe_init_model(self, config: MergeConfiguration):
                     continue
 
                 if getattr(cfg_out, key) != getattr(self.arch_info.config, key, None):
-                    logger.warning(f"Config key {key} changed, reinitializing model")
+                    LOG.warning(f"Config key {key} changed, reinitializing model")
                     different = True
                     break
 
@@ -206,7 +206,7 @@ def _maybe_init_model(self, config: MergeConfiguration):
                 del inner_model
                 tokenizer_donor = self.genome.definition.base_model
                 if tokenizer_donor is None:
-                    logger.warning(
+                    LOG.warning(
                         "Base model not set, using tokenizer from first model in genome"
                     )
                     tokenizer_donor = self.genome.definition.models[0]
@@ -224,7 +224,7 @@ def _maybe_init_model(self, config: MergeConfiguration):
                     max_model_len = min(max_model_len or 1024, window_sz)
                 if max_model_len and max_model_len > 8192:
                     max_model_len = 8192
-                    logger.warning(f"Clipping sequence length to {max_model_len}")
+                    LOG.warning(f"Clipping sequence length to {max_model_len}")
 
                 mem_util = (
                     0.7 if self.merge_options.cuda else 0.9
@@ -248,13 +248,13 @@ def _maybe_init_model(self, config: MergeConfiguration):
             if ai
             else None
         )
-        logger.info("Model initialized")
+        LOG.info("Model initialized")
 
     def evaluate(self, genotype: torch.Tensor) -> dict:
         try:
             config = self.genome.genotype_merge_config(genotype)
         except InvalidGenotypeError as e:
-            logger.error("Invalid genotype", exc_info=e)
+            LOG.error("Invalid genotype", exc_info=e)
             return {"score": None, "results": None}
 
         self._maybe_init_model(config)