@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 92ff41abea
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Model] Modify the gate implementation of glm4_moe (#22832)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Aug 14, 2025 at 20:28 UTC 92ff41abea130a3217faa54abb89ccc27aef3f06
2 files changed +11 -11
docs/models/supported_models.md
+1 -1
@@ -615,7 +615,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
615 | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | ✅︎ |
616 | `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | ✅︎ |
617 | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | ✅︎ |
618 -| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | | ✅︎ | ✅︎ |
618 +| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ | ✅︎ |
619 | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ | ✅︎ |
620 | `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ | ✅︎ |
621 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | | ✅︎ |
vllm/model_executor/models/glm4_moe.py
+10 -10
@@ -41,7 +41,6 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
41 from vllm.model_executor.layers.layernorm import RMSNorm
42 from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
43 QKVParallelLinear,
44 - ReplicatedLinear,
44 RowParallelLinear)
45 from vllm.model_executor.layers.logits_processor import LogitsProcessor
46 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -118,14 +117,15 @@ class Glm4MoE(nn.Module):
117 if config.hidden_act != "silu":
118 raise ValueError(f"Unsupported activation: {config.hidden_act}. "
119 "Only silu is supported for now.")
121 -
122 - self.gate = ReplicatedLinear(config.hidden_size,
123 - config.n_routed_experts,
124 - bias=False,
125 - quant_config=None,
126 - params_dtype=torch.float32,
127 - prefix=f"{prefix}.gate")
128 -
120 + # NOTE In the transformers implementation, the gate isn't an nn.Linear,
121 + # so we cannot use ReplicatedLinear here.
122 + # See: https://github.com/huggingface/transformers/blob/v4.55.1/src/transformers/models/glm4_moe/modeling_glm4_moe.py#L260
123 + self.gate = nn.Linear(
124 + config.hidden_size,
125 + config.n_routed_experts,
126 + bias=False,
127 + dtype=torch.float32,
128 + )
129 self.gate.e_score_correction_bias = nn.Parameter(
130 torch.empty(config.n_routed_experts, dtype=torch.float32))
131
@@ -181,7 +181,7 @@ class Glm4MoE(nn.Module):
181
182 if self.n_shared_experts is not None:
183 shared_output = self.shared_experts(hidden_states)
184 - router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
184 + router_logits = self.gate(hidden_states.to(dtype=torch.float32))
185 final_hidden_states = self.experts(
186 hidden_states=hidden_states,
187 router_logits=router_logits) * self.routed_scaling_factor