extend modelbuilder to build Olmo3, SmolLM3 and other models

.github/workflows/modelbuilder_fast_tests.yml

@@ -0,0 +1,61 @@
+name: Fast Tests
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+permissions:
+  contents: read
+
+jobs:
+  tests:
+    name: ci (${{ matrix.os }}, py${{ matrix.python-version }}, torch-${{ matrix.torch-version || 'stable' }}, transformers-${{ matrix.transformers-version || 'latest' }})
+    runs-on: ${{ matrix.os }}
+    permissions:
+      contents: read
+      id-token: write
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version: ["3.13"]
+        transformers-version: ["4.57", "5.5"]
+        torch-version: ["2.11"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      # we install torch first to avoid downloading any CUDA dependency
+
+      - name: Install nightly pytorch (cpu)
+        if: matrix.torch-version == 'nightly'
+        run: pip install --pre --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu
+
+      - name: Install pytorch 2.11 (Linux/Windows)
+        if: matrix.torch-version == '2.11' && matrix.os != 'macos-latest'
+        run: pip install torch==2.11.0+cpu torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+
+      - name: Install dependencies
+        run: pip install torch "transformers==${{ matrix.transformers-version }}" tokenizers pandas openpyxl pytest huggingface_hub onnx onnxscript optree onnxruntime onnxruntime-genai peft py-cpuinfo pytest-cov pytest-subtests tabulate
+
+      - name: Run fast tests
+        run: DOCLEAN=1 PYTHONPATH=./src/python/py pytest test/python/models/fast -v --cov=modelbuilder --cov-report=xml --cov-report=term-missing
+
+      # Code to upload coverage report.
+      # - name: Upload coverage reports to Codecov
+      #   uses: codecov/codecov-action@v5
+      #   with:
+      #    token: ${{ secrets.CODECOV_TOKEN }}
+      #     flags: modelbuilder-fast-tests
+
+      - name: Upload stats artifact
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: stats-${{ matrix.os }}-py${{ matrix.python-version }}-torch${{ matrix.torch-version }}-transformers${{ matrix.transformers-version }}
+          path: stats/
+          if-no-files-found: ignore

.gitignore

@@ -16,6 +16,8 @@ src/csharp/AssemblyInfo.cs
 *.csv
 .idea
 cache_dir
+dump_models
+stats
 example-models
 *.onnx
 *.onnx.data

src/python/py/models/builders/__init__.py

@@ -7,57 +7,5 @@
 # Portions of this file consist of AI generated content.
 # -------------------------------------------------------------------------
 from .base import Model
-from .chatglm import ChatGLMModel
-from .ernie import ErnieModel
-from .gemma import Gemma2Model, Gemma3Model, GemmaModel
-from .gptoss import GPTOSSModel
-from .granite import GraniteModel
-from .internlm import InternLM2Model
-from .llama import LlamaModel
-from .mistral import MistralModel
-from .nemotron import NemotronModel
-from .olmo import OLMoModel
-from .phi import (
-    Phi3MiniLongRoPEModel,
-    Phi3MiniModel,
-    Phi3MoELongRoPEModel,
-    Phi3SmallLongRoPEModel,
-    Phi3SmallModel,
-    Phi3VModel,
-    Phi4MMModel,
-    PhiModel,
-)
-from .qwen import Qwen3Model, Qwen3VLTextModel, Qwen25VLTextModel, Qwen35TextModel, QwenModel
-from .smollm import SmolLM3Model
-from .whisper import WhisperModel
 
-__all__ = [
-    "ChatGLMModel",
-    "ErnieModel",
-    "GPTOSSModel",
-    "Gemma2Model",
-    "Gemma3Model",
-    "GemmaModel",
-    "GraniteModel",
-    "InternLM2Model",
-    "LlamaModel",
-    "MistralModel",
-    "Model",
-    "NemotronModel",
-    "OLMoModel",
-    "Phi3MiniLongRoPEModel",
-    "Phi3MiniModel",
-    "Phi3MoELongRoPEModel",
-    "Phi3SmallLongRoPEModel",
-    "Phi3SmallModel",
-    "Phi3VModel",
-    "Phi4MMModel",
-    "PhiModel",
-    "Qwen3Model",
-    "Qwen3VLTextModel",
-    "Qwen25VLTextModel",
-    "Qwen35TextModel",
-    "QwenModel",
-    "SmolLM3Model",
-    "WhisperModel",
-]
+__all__ = ["Model"]

src/python/py/models/builders/ernie.py

@@ -13,6 +13,14 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         # Ernie uses interleaved rotary position embeddings.
         self.rope_attrs["interleaved"] = 1
 
+        # Ernie4_5Config stores rope_theta inside a `rope_parameters` dict rather
+        # than as a top-level attribute.  The base Model.__init__ falls back to
+        # 10000 when `config.rope_theta` is absent, so we patch the value here.
+        if not hasattr(config, "rope_theta") and hasattr(config, "rope_parameters"):
+            rope_theta = config.rope_parameters.get("rope_theta")
+            if rope_theta is not None:
+                self.rope_attrs["theta"] = rope_theta
+
         # Ernie uses a `compression_ratio` for its RoPE scaling.
         # The original RoPE logic in ernie is: position_ids / compression_ratio,
         # which is equivalent to scaling the frequencies (inv_freq) by 1 / compression_ratio.

src/python/py/models/builders/gemma.py

@@ -124,7 +124,25 @@ class Gemma3Model(Gemma2Model):
     def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
 
-        self.rope_local_theta = config.rope_local_base_freq
+        # Remember whether this was originally a conditional-generation config
+        # (Gemma3ForConditionalGeneration) before builder.py overrides model_type
+        # to "gemma3_text". load_weights needs to know which HF class to instantiate.
+        self._original_architecture = config.architectures[0]
+
+        if hasattr(config, "rope_local_base_freq"):
+            # Older transformers: rope_local_base_freq and rope_theta are top-level fields
+            self.rope_local_theta = config.rope_local_base_freq
+        elif hasattr(config, "rope_parameters") and isinstance(config.rope_parameters, dict):
+            # Newer transformers (v5+): rope info lives in a nested rope_parameters dict
+            sliding_params = config.rope_parameters.get("sliding_attention", {})
+            self.rope_local_theta = sliding_params.get("rope_theta", 10000.0)
+            # Update the global theta (full-attention layers) which the base class
+            # could not infer because config.rope_theta is absent in this format.
+            full_params = config.rope_parameters.get("full_attention", {})
+            self.rope_attrs["theta"] = full_params.get("rope_theta", self.rope_attrs["theta"])
+        else:
+            # Default local RoPE theta matching Gemma3's original rope_local_base_freq
+            self.rope_local_theta = 10000.0
         self.make_rotary_embedding_multi_cache()
 
     def is_local(self, layer_id):
@@ -136,7 +154,7 @@ def make_attention_init(self):
         super().make_attention_init()
 
     def make_rotary_embedding_multi_cache(self):
-        self.cos_cache_global_name, self.sin_cache_global_name = "cos_cache_global", "sin_cache_global"
+        self.cos_cache_global_name, self.sin_cache_global_name = ("cos_cache_global", "sin_cache_global")
         super().make_rotary_embedding_caches(
             cos_cache_name=self.cos_cache_global_name, sin_cache_name=self.sin_cache_global_name
         )
@@ -145,16 +163,31 @@ def make_rotary_embedding_multi_cache(self):
         self.rope_attrs["create_caches"] = True
         self.rope_attrs["theta"] = self.rope_local_theta
 
-        self.cos_cache_local_name, self.sin_cache_local_name = "cos_cache_local", "sin_cache_local"
+        self.cos_cache_local_name, self.sin_cache_local_name = ("cos_cache_local", "sin_cache_local")
         super().make_rotary_embedding_caches(
             cos_cache_name=self.cos_cache_local_name, sin_cache_name=self.sin_cache_local_name
         )
 
+    def load_weights(self, input_path):
+        # Gemma3ForConditionalGeneration (VLM) does not accept the
+        # ``num_hidden_layers`` keyword argument that the base class would
+        # normally forward to ``AutoModelForCausalLM.from_pretrained``.
+        # Load it directly here instead.
+        if self._original_architecture == "Gemma3ForConditionalGeneration":
+            if self.quant_type is not None or input_path.endswith(".gguf"):
+                return super().load_weights(input_path)
+            from transformers import Gemma3ForConditionalGeneration as _HFModel
+
+            return _HFModel.from_pretrained(
+                self.model_name_or_path, cache_dir=self.cache_dir, token=self.hf_token, trust_remote_code=self.hf_remote
+            )
+        return super().load_weights(input_path)
+
     def make_rotary_embedding_caches(self, **kwargs):
         cos_cache_name = kwargs.get(
-            "cos_cache_name", self.cos_cache_global_name if self.window_size == -1 else self.cos_cache_local_name
+            "cos_cache_name", (self.cos_cache_global_name if self.window_size == -1 else self.cos_cache_local_name)
         )
         sin_cache_name = kwargs.get(
-            "sin_cache_name", self.sin_cache_global_name if self.window_size == -1 else self.sin_cache_local_name
+            "sin_cache_name", (self.sin_cache_global_name if self.window_size == -1 else self.sin_cache_local_name)
         )
         return super().make_rotary_embedding_caches(cos_cache_name=cos_cache_name, sin_cache_name=sin_cache_name)

src/python/py/models/builders/gptoss.py

@@ -53,12 +53,14 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
         super().make_layernorm(layer_id, layernorm, skip, simple, location)
 
     def make_rotary_embedding_caches_from_scratch(self):
-        inv_freq = self.rope_attrs["theta"] ** (torch.arange(0, self.head_size, 2, dtype=torch.float) / self.head_size)
+        inv_freq = 1.0 / (
+            self.rope_attrs["theta"] ** (torch.arange(0, self.head_size, 2, dtype=torch.float) / self.head_size)
+        )
         inv_freq = self.make_inv_freq_rescaled(inv_freq)
 
         t = torch.arange(self.rope_attrs["cache_length"], dtype=torch.float32)
         freqs = torch.einsum("i,j->ij", t, inv_freq)
-        cos_cache, sin_cache = freqs.cos() * self.rope_attrs["mscale"], freqs.sin() * self.rope_attrs["mscale"]
+        cos_cache, sin_cache = (freqs.cos() * self.rope_attrs["mscale"], freqs.sin() * self.rope_attrs["mscale"])
         return cos_cache, sin_cache
 
     def make_attention(self, layer_id, attention, root_input, **kwargs):
@@ -702,41 +704,57 @@ def has_quark_experts(self, experts):
 
     def combine_quark_gate_up_biases_from_experts(self, experts):
         """Combine Quark gate_proj and up_proj biases from individual experts"""
-        assert(self.has_quark_experts(experts))
+        assert self.has_quark_experts(experts)
         combined_biases = []
 
         for expert_id in sorted(experts.keys()):
             expert = experts[expert_id]
 
             if expert.gate_up_proj.qweight is not None:
                 # Fused gate_up projection
-                gate_up_proj = expert.gate_up_proj.bias if hasattr(expert.gate_up_proj, 'bias') and expert.gate_up_proj.bias is not None else torch.zeros(expert.gate_up_proj.qweight.shape[0])
+                gate_up_proj = (
+                    expert.gate_up_proj.bias
+                    if hasattr(expert.gate_up_proj, "bias") and expert.gate_up_proj.bias is not None
+                    else torch.zeros(expert.gate_up_proj.qweight.shape[0])
+                )
                 combined_biases.append(gate_up_proj)
             else:
                 # Get biases from individual projections
-                gate_bias = expert.gate_proj.bias if hasattr(expert.gate_proj, 'bias') and expert.gate_proj.bias is not None else torch.zeros(expert.gate_proj.qweight.shape[0])
-                up_bias = expert.up_proj.bias if hasattr(expert.up_proj, 'bias') and expert.up_proj.bias is not None else torch.zeros(expert.up_proj.qweight.shape[0])
+                gate_bias = (
+                    expert.gate_proj.bias
+                    if hasattr(expert.gate_proj, "bias") and expert.gate_proj.bias is not None
+                    else torch.zeros(expert.gate_proj.qweight.shape[0])
+                )
+                up_bias = (
+                    expert.up_proj.bias
+                    if hasattr(expert.up_proj, "bias") and expert.up_proj.bias is not None
+                    else torch.zeros(expert.up_proj.qweight.shape[0])
+                )
 
                 # Combine gate and up biases (interleaved pattern: even=gate, odd=up)
                 gate_out_dim = gate_bias.shape[0]
                 up_out_dim = up_bias.shape[0]
 
                 combined_bias = torch.zeros(gate_out_dim + up_out_dim, dtype=gate_bias.dtype, device="cpu")
-                combined_bias[::2] = gate_bias   # Even indices = gate
-                combined_bias[1::2] = up_bias    # Odd indices = up
+                combined_bias[::2] = gate_bias  # Even indices = gate
+                combined_bias[1::2] = up_bias  # Odd indices = up
 
                 combined_biases.append(combined_bias)
 
         return torch.stack(combined_biases, dim=0)
 
     def combine_quark_down_biases_from_experts(self, experts):
         """Combine Quark down_proj biases from individual experts"""
-        assert(self.has_quark_experts(experts))
+        assert self.has_quark_experts(experts)
         combined_biases = []
 
         for expert_id in sorted(experts.keys()):
             expert = experts[expert_id]
-            down_bias = expert.down_proj.bias if hasattr(expert.down_proj, 'bias') and expert.down_proj.bias is not None else torch.zeros(expert.down_proj.qweight.shape[0])
+            down_bias = (
+                expert.down_proj.bias
+                if hasattr(expert.down_proj, "bias") and expert.down_proj.bias is not None
+                else torch.zeros(expert.down_proj.qweight.shape[0])
+            )
             combined_biases.append(down_bias)
 
         return torch.stack(combined_biases, dim=0)