Add Glm4V-MoE architecture

Author: turboderp

examples/common.py

@@ -75,4 +75,4 @@ def get_stop_conditions(prompt_format, tokenizer):
         case "mistral":
             return [tokenizer.eos_token_id]
         case "glmv":
-            return [tokenizer.eos_token_id, "</answer>"]
+            return [tokenizer.eos_token_id, "</answer>", "<|user|>"]

examples/multimodal.py

@@ -24,7 +24,7 @@
         model_dir = "/mnt/str/models/qwen3-vl-30b-a3b-instruct/exl3/5.00bpw"
     case "glm":
         prompt_format = "glmv"
-        model_dir = "/mnt/str/models/glm4.1v-9b-thinking/exl3/2.0bpw"
+        model_dir = "/mnt/str/models/glm4.5v/exl3/4.00bpw"
 
 images = [
     # Cat

exllamav3/architecture/architectures.py

@@ -12,6 +12,7 @@
 from .glm4 import Glm4Model
 from .glm4_moe import Glm4MoeModel
 from .glm4v import Glm4VModel
+from .glm4v_moe import Glm4VMoeModel
 from .llama import LlamaModel
 from .mimo import MiMoModel
 from .minimax_m2 import MiniMaxM2Model
@@ -49,6 +50,7 @@
         Glm4Model,
         Glm4MoeModel,
         Glm4VModel,
+        Glm4VMoeModel,
         LlamaModel,
         MiMoModel,
         MiniMaxM2Model,

exllamav3/architecture/glm4_moe.py

@@ -7,6 +7,10 @@
 from ..modules import RMSNorm, Embedding, TransformerBlock, Attention, GatedMLP, Linear, BlockSparseMLP
 from ..modules.attn import prepare_for_attn
 
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from .glm4v_moe import Glm4VMoeConfig
+
 class Glm4MoeConfig(Config):
     arch_string = "Glm4MoeForCausalLM"
 
@@ -58,15 +62,16 @@ class Glm4MoeModel(Model):
 
     def __init__(
         self,
-        config: Glm4MoeConfig,
+        config: Glm4MoeConfig | Glm4VMoeConfig,
+        key_prefix: str = "model",
         **kwargs
     ):
         super().__init__(config, **kwargs)
 
         self.modules += [
             Embedding(
                 config = config,
-                key = "model.embed_tokens",
+                key = f"{key_prefix}.embed_tokens",
                 vocab_size = config.vocab_size,
                 hidden_size = config.hidden_size,
             )
@@ -77,15 +82,15 @@ def __init__(
         self.modules += [
             TransformerBlock(
                 config = config,
-                key = f"model.layers.{idx}",
+                key = f"{key_prefix}.layers.{idx}",
                 attn_norm = RMSNorm(
                     config = config,
-                    key = f"model.layers.{idx}.input_layernorm",
+                    key = f"{key_prefix}.layers.{idx}.input_layernorm",
                     rms_norm_eps = config.rms_norm_eps,
                 ),
                 attn = Attention(
                     config = config,
-                    key = f"model.layers.{idx}.self_attn",
+                    key = f"{key_prefix}.layers.{idx}.self_attn",
                     layer_idx = idx,
                     hidden_size = config.hidden_size,
                     head_dim = config.head_dim,
@@ -100,25 +105,25 @@ def __init__(
                     qmap = "block.attn",
                     q_norm = RMSNorm(
                         config = config,
-                        key = f"model.layers.{idx}.self_attn.q_norm",
+                        key = f"{key_prefix}.layers.{idx}.self_attn.q_norm",
                         rms_norm_eps = config.rms_norm_eps,
                     ) if config.use_qk_norm else None,
                     k_norm = RMSNorm(
                         config = config,
-                        key = f"model.layers.{idx}.self_attn.k_norm",
+                        key = f"{key_prefix}.layers.{idx}.self_attn.k_norm",
                         rms_norm_eps = config.rms_norm_eps,
                     ) if config.use_qk_norm else None,
                     out_dtype = torch.float
                 ),
                 mlp_norm = RMSNorm(
                     config = config,
-                    key = f"model.layers.{idx}.post_attention_layernorm",
+                    key = f"{key_prefix}.layers.{idx}.post_attention_layernorm",
                     rms_norm_eps = config.rms_norm_eps,
                 ),
                 mlp = (
                     GatedMLP(
                         config = config,
-                        key = f"model.layers.{idx}.mlp",
+                        key = f"{key_prefix}.layers.{idx}.mlp",
                         hidden_size = config.hidden_size,
                         intermediate_size = config.intermediate_size,
                         key_up = "up_proj",
@@ -131,7 +136,7 @@ def __init__(
                     if idx < config.first_k_dense_replace else
                     BlockSparseMLP(
                         config = config,
-                        key = f"model.layers.{idx}.mlp",
+                        key = f"{key_prefix}.layers.{idx}.mlp",
                         hidden_size = config.hidden_size,
                         intermediate_size = config.moe_intermediate_size,
                         num_experts = config.num_experts,
@@ -150,7 +155,7 @@ def __init__(
                         topk_group = 1,
                         shared_experts = GatedMLP(
                             config = config,
-                            key = f"model.layers.{idx}.mlp.shared_experts",
+                            key = f"{key_prefix}.layers.{idx}.mlp.shared_experts",
                             hidden_size = config.hidden_size,
                             intermediate_size = config.moe_intermediate_size * config.num_shared_experts,
                             key_up = "up_proj",
@@ -170,12 +175,12 @@ def __init__(
 
         head_alt_key = None
         if config.tie_word_embeddings and not self.config.stc.has_tensor("lm_head"):
-            head_alt_key = "model.embed_tokens"
+            head_alt_key = f"{key_prefix}.embed_tokens"
 
         self.modules += [
             RMSNorm(
                 config = config,
-                key = "model.norm",
+                key = f"{key_prefix}.norm",
                 rms_norm_eps = config.rms_norm_eps,
                 out_dtype = torch.half,
             ),

exllamav3/architecture/glm4v.py

@@ -28,6 +28,10 @@
 from PIL import Image
 import os, json
 
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from .glm4v_moe import Glm4VMoeConfig
+
 class Glm4VConfig(Config):
     arch_string = "Glm4vForConditionalGeneration"
 
@@ -262,7 +266,7 @@ def get_additional_compiled_tensors(config: Glm4VConfig) -> dict:
 
     def __init__(
         self,
-        config: Glm4VConfig,
+        config: Glm4VConfig | Glm4VMoeConfig,
         key_prefix = "model.visual",
         **kwargs
     ):

exllamav3/architecture/glm4v_moe.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+from typing_extensions import override
+import numpy as np
+import torch
+import torch.nn.functional as F
+from ..model.config import Config, no_default
+from ..model.model import Model
+from ..util.rope import RopeStyle, position_embedding_grid_2d, RopeSettings
+from ..util.file import read_dict, no_value, no_default
+from ..util.vision import convert_to_rgb, normalize_image, smart_resize
+from ..modules import (
+    Module,
+    RMSNorm,
+    Embedding,
+    TransformerBlock,
+    Attention,
+    GatedMLP,
+    Linear,
+    Conv,
+    LayerNorm,
+    Glm4VPosEmbedding,
+    MLP
+)
+from ..modules.attn import prepare_for_attn
+from .glm4_moe import Glm4MoeModel
+from types import SimpleNamespace
+from ..tokenizer import Tokenizer, MMEmbedding
+from PIL import Image
+import os, json
+from .glm4v import read_glm4v_vision_config, read_glm4v_pp_config, Glm4VVisionModel
+
+class Glm4VMoeConfig(Config):
+    arch_string = "Glm4vMoeForConditionalGeneration"
+
+    def __init__(
+        self,
+        directory: str,
+        **kwargs,
+    ):
+        super().__init__(
+            directory,
+            {"text": Glm4VMoeModel, "vision": Glm4VVisionModel},
+            **kwargs
+        )
+
+        # Attention params
+        self.head_dim = self.read_cfg(int, "text_config->head_dim", None)
+        self.hidden_size = self.read_cfg(int, "text_config->hidden_size", no_default)
+        self.num_q_heads = self.read_cfg(int, "text_config->num_attention_heads", no_default)
+        self.num_kv_heads = self.read_cfg(int, "text_config->num_key_value_heads", self.num_q_heads)
+        self.use_qk_norm = self.read_cfg(bool, "text_config->use_qk_norm", False)
+
+        if not self.head_dim:
+            self.head_dim = self.hidden_size // self.num_q_heads
+
+        # MLP params
+        self.assert_cfg(str, "text_config->hidden_act", "silu", True)
+        self.assert_cfg(bool, "text_config->norm_topk_prob", True, True)
+        self.intermediate_size = self.read_cfg(int, "text_config->intermediate_size", no_default)
+        self.moe_intermediate_size = self.read_cfg(int, "text_config->moe_intermediate_size", no_default)
+        self.num_shared_experts = self.read_cfg(int, "text_config->n_shared_experts", 1)
+        self.num_experts = self.read_cfg(int, "text_config->n_routed_experts", 128)
+        self.num_experts_per_tok = self.read_cfg(int, "text_config->num_experts_per_tok", 8)
+        self.first_k_dense_replace = self.read_cfg(int, "text_config->first_k_dense_replace", 3)
+        self.routed_scaling_factor = self.read_cfg(float, "text_config->routed_scaling_factor", 2.5)
+
+        # Norms
+        self.rms_norm_eps = self.read_cfg(float, "text_config->rms_norm_eps", no_default)
+
+        # Layers
+        self.num_hidden_layers = self.read_cfg(int, "text_config->num_hidden_layers", no_default)
+        self.tie_word_embeddings = self.read_cfg(bool, "tie_word_embeddings", False)
+
+        # RoPE
+        self.rope_settings = self.read_rope_settings_default(
+            RopeStyle.NEOX,
+            default_rope_theta = 10000,
+            config_dict = self.read_cfg(dict, "text_config", no_default)
+        )
+
+        # Vision model settings
+        read_vision_config = self.read_cfg(dict, "vision_config", no_default)
+        self.vision = read_glm4v_vision_config(read_vision_config)
+
+        prep_path = os.path.join(self.directory, "preprocessor_config.json")
+        with open(prep_path, encoding = "utf8") as f:
+            read_prep_config = json.load(f)
+        self.vision_pp = read_glm4v_pp_config(read_prep_config)
+
+        self.vision_start_token_id = self.read_cfg(int, "image_start_token_id", 151339)
+        self.vision_end_token_id = self.read_cfg(int, "image_end_token_id", 151340)
+
+
+class Glm4VMoeModel(Glm4MoeModel):
+    config_class = Glm4VMoeConfig
+
+    def __init__(
+        self,
+        config: Glm4VMoeConfig,
+        **kwargs
+    ):
+        super().__init__(config, key_prefix = "model.language_model", **kwargs)