Correctly set vocab size in Gemma3 merges

Author: cg123

mergekit/_data/architectures/gemma3vl.json

@@ -8,6 +8,7 @@
         "preprocessor_config.json",
         "processor_config.json"
     ],
+    "vocab_size_config_key": "text_config.vocab_size",
     "modules": {
         "text_decoder": {
             "weight_prefix": "language_model.",

mergekit/architecture/base.py

@@ -128,6 +128,7 @@ class ModelArchitecture(BaseModel, frozen=True):
     architectures: List[str]
     expected_model_type: str = Field(alias="model_type")
     tagalong_files: Optional[List[str]] = None
+    vocab_size_config_key: Optional[str] = None
 
     def all_weights(self, config: PretrainedConfig) -> List[WeightInfo]:
         res = []

mergekit/architecture/json_definitions.py

@@ -105,6 +105,7 @@ class JsonModularArchitectureDefinition(BaseModel, frozen=True):
     architectures: List[str]
     expected_model_type: str = Field(alias="model_type")
     tagalong_files: Optional[List[str]] = None
+    vocab_size_config_key: Optional[str] = None
 
 
 class TemplateWithArithmetic(string.Template):
@@ -154,6 +155,7 @@ def _load_architecture_json(name: str) -> ModelArchitecture:
             architectures=parsed.architectures,
             model_type=parsed.expected_model_type,
             tagalong_files=parsed.tagalong_files,
+            vocab_size_config_key=parsed.vocab_size_config_key,
         )
     elif data.get("kind", "module") == "module":
         module = JsonModuleArchitecture(

mergekit/merge.py

@@ -90,7 +90,9 @@ def run_merge(
         pad_to_multiple_of = None
         if merge_config.tokenizer and merge_config.tokenizer.pad_to_multiple_of:
             pad_to_multiple_of = merge_config.tokenizer.pad_to_multiple_of
-        _update_config_vocab(cfg_out, tokenizer, pad_to_multiple_of=pad_to_multiple_of)
+        _update_config_vocab(
+            cfg_out, arch_info, tokenizer, pad_to_multiple_of=pad_to_multiple_of
+        )
 
     logger.info("Saving config")
     cfg_out.save_pretrained(out_path)
@@ -308,14 +310,15 @@ def _model_out_config(
 
 def _update_config_vocab(
     config: transformers.PretrainedConfig,
+    arch_info: ModelArchitecture,
     tokenizer: transformers.PreTrainedTokenizerBase,
     pad_to_multiple_of: Optional[int] = None,
 ):
     vocab_size = len(tokenizer.get_vocab())
     if pad_to_multiple_of and vocab_size % pad_to_multiple_of:
         vocab_size = vocab_size + pad_to_multiple_of - (vocab_size % pad_to_multiple_of)
     try:
-        config.vocab_size = vocab_size
+        setattr(config, arch_info.vocab_size_config_key or "vocab_size", vocab_size)
     except Exception as e:
         logger.warning(
             "Unable to set vocabulary size in output config - you may need to manually correct it.",

mergekit/scripts/tokensurgeon.py

@@ -199,7 +199,11 @@ def main(
     tokenizer.save_pretrained(out_path)
     cfg_out = arch_info.config
     try:
-        cfg_out.vocab_size = new_embed.shape[0]
+        setattr(
+            cfg_out,
+            arch_info.info.vocab_size_config_key or "vocab_size",
+            new_embed.shape[0],
+        )
     except AttributeError:
         LOG.error(
             "Could not set vocab size in config.json - you may need to update it manually."

mergekit/tokenizer/build.py

@@ -13,7 +13,8 @@
 from pydantic import BaseModel
 from typing_extensions import Literal
 
-from mergekit.common import ModelPath, ModelReference
+from mergekit.architecture import arch_info_for_config
+from mergekit.common import ModelPath, ModelReference, get_config_value
 from mergekit.graph import Task
 
 logger = logging.getLogger(__name__)
@@ -26,7 +27,8 @@ def get_vocab_size(model_path: ModelPath, trust_remote_code: bool) -> Optional[i
             revision=model_path.revision,
             trust_remote_code=trust_remote_code,
         )
-        return cfg.vocab_size
+        arch_info = arch_info_for_config(cfg)
+        return get_config_value(cfg, arch_info.vocab_size_config_key or "vocab_size")
     except Exception as e:
         logger.warning(f"Unable to get vocab size for {model_path}", exc_info=e)