minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

GLM-4.7 and GLM-4.7-Flash Loading and import format (#21851)

Browse Source
This commit is contained in:
Yuxuan Zhang
2026-04-03 20:44:08 -07:00
committed by GitHub
parent db3d4f4b76
commit b7ae3b5a9a
2 changed files with 139 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/glm4_moe.py
+116 -43
View File
@@ -15,6 +15,7 @@
"""Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 model compatible with HuggingFace weights"""
import logging
import re
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
import torch
@@ -22,6 +23,7 @@ import torch.nn.functional as F
from torch import nn
from transformers import PretrainedConfig
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import (
get_moe_expert_parallel_world_size,
@@ -63,6 +65,7 @@ from sglang.srt.layers.moe import (
)
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.utils import (
RoutingMethodType,
@@ -183,7 +186,7 @@ class Glm4MoeAttention(nn.Module):
num_heads: int,
num_kv_heads: int,
layer_id: int = 0,
rope_theta: float = 10000,
rope_theta: float = 1000000,
partial_rotary_factor: float = 0.5,
rope_scaling: Optional[Dict[str, Any]] = None,
max_position_embeddings: int = 8192,
@@ -439,9 +442,12 @@ class Glm4MoeSparseMoeBlock(nn.Module):
fused_shared_experts_scaling_factor=1,
)
# shared expert
self.shared_experts_is_int8 = False
self.shared_experts_is_fp8 = False
self.shared_experts_weight_block_size = None
if config.n_shared_experts is not None and self.num_fused_shared_experts == 0:
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
# disable tp for shared experts when enable deepep moe, or with fp4 allgather
self.shared_experts = Glm4MoeMLP(
hidden_size=config.hidden_size,
intermediate_size=intermediate_size,
@@ -453,16 +459,54 @@ class Glm4MoeSparseMoeBlock(nn.Module):
dict(tp_rank=0, tp_size=1)
if get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
or should_use_flashinfer_cutlass_moe_fp4_allgather()
else {}
),
)
is_packed_weight = hasattr(
self.shared_experts.gate_up_proj.quant_method, "quant_config"
) and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() in {
"awq",
"awq_marlin",
"moe_wna16",
}
self.shared_experts_is_int8 = (
not is_packed_weight
and self.shared_experts.gate_up_proj.weight.dtype == torch.int8
)
self.shared_experts_is_fp8 = (
not is_packed_weight
and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn
)
if self.shared_experts_is_fp8:
if (
_use_aiter
and config.quantization_config.get("quant_method")
== "compressed-tensors"
):
# For compressed-tensors ptpc model, don't need to check the weight_block_size
pass
else:
assert (
self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
== self.shared_experts.down_proj.quant_method.quant_config.weight_block_size
)
self.shared_experts_weight_block_size = (
self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
)
self.top_k = config.num_experts_per_tok
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
):
# TODO: we will support tp < ep in the future
self.ep_size = get_moe_expert_parallel_world_size()
@@ -483,7 +527,11 @@ class Glm4MoeSparseMoeBlock(nn.Module):
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_ascend_fuseep()
or get_moe_a2a_backend().is_flashinfer()
)
self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
def get_moe_weights(self):
return [
@@ -502,8 +550,7 @@ class Glm4MoeSparseMoeBlock(nn.Module):
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if not get_moe_a2a_backend().is_deepep():
if not self._enable_a2a_moe:
if (
self.alt_stream is not None
and self.num_fused_shared_experts == 0
@@ -511,11 +558,15 @@ class Glm4MoeSparseMoeBlock(nn.Module):
and get_is_capture_mode()
):
return self.forward_normal_dual_stream(
hidden_states, should_allreduce_fusion, use_reduce_scatter
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
else:
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
)
else:
return self.forward_deepep(hidden_states, forward_batch)
@@ -534,20 +585,12 @@ class Glm4MoeSparseMoeBlock(nn.Module):
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
topk_output = self.topk(hidden_states, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output)
if not _is_cuda and not _use_aiter:
# fused in biased_grouped_topk so we can skip here
if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
final_hidden_states *= self.routed_scaling_factor
current_stream.wait_stream(self.alt_stream)
with use_symmetric_memory(
parallel_state.get_tp_group(), disabled=not is_allocation_symmetric()
):
final_hidden_states_out = torch.empty_like(final_hidden_states)
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
final_hidden_states = final_hidden_states_out
final_hidden_states += shared_output
if (
self.tp_size > 1
and not should_allreduce_fusion
@@ -1090,27 +1133,32 @@ class Glm4MoeForCausalLM(nn.Module):
# For EAGLE3 support
self.capture_aux_hidden_states = False
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
def determine_num_fused_shared_experts(self):
if get_global_server_args().disable_shared_experts_fusion:
return
disable_reason = None
if not getattr(self.config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config."
elif not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices."
elif _is_cuda and (_device_sm is not None) and (_device_sm < 80):
disable_reason = "Shared experts fusion requires SM80 or newer GPUs."
elif get_moe_expert_parallel_world_size() > 1:
disable_reason = "Shared experts fusion is not supported together with expert parallelism yet."
elif get_moe_a2a_backend().is_deepep():
disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled."
if (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and (
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
):
disable_reason = (
"Only GLM-4.5 on NV-platform with capability >= 80 "
"or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
)
elif get_moe_expert_parallel_world_size() > 1 and (
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
):
disable_reason = "Only GLM-4.5 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."
elif disable_reason is None and (
get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori()
):
disable_reason = "GLM-4.5 cannot use shared experts fusion optimization under deepep expert parallelism."
elif self.quant_config and self.quant_config.get_name() == "w4afp8":
disable_reason = "GLM-4.5 W4AFP8 model uses different quant method for routed experts and shared experts."
if disable_reason is not None:
get_global_server_args().disable_shared_experts_fusion = True
self.num_fused_shared_experts = 0
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
@@ -1118,10 +1166,9 @@ class Glm4MoeForCausalLM(nn.Module):
return
self.num_fused_shared_experts = self.config.n_shared_experts
assert (
self.num_fused_shared_experts == 1
), "Only 1 fused shared expert is supported for Glm4MoeForCausalLM"
log_info_on_rank0(logger, "Shared experts fusion optimization enabled.")
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
@torch.no_grad()
def forward(
@@ -1154,7 +1201,12 @@ class Glm4MoeForCausalLM(nn.Module):
def end_layer(self):
return self.model.end_layer
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False):
def load_weights(
self,
weights: Iterable[Tuple[str, torch.Tensor]],
is_nextn=False,
params_dict=None,
):
if is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = self.config.num_nextn_predict_layers
@@ -1177,6 +1229,28 @@ class Glm4MoeForCausalLM(nn.Module):
("gate_up_proj", "up_proj", 1),
]
if self.num_fused_shared_experts > 0:
assert self.num_fused_shared_experts == 1
def iter_weights_with_fused_shared_experts(
weights: Iterable[Tuple[str, torch.Tensor]],
) -> Iterable[Tuple[str, torch.Tensor]]:
pattern = re.compile(
r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$"
)
for name, weight in weights:
match = pattern.match(name)
if match:
layer_id = int(match.group(1))
suffix = match.group(2)
name = f"model.layers.{layer_id}.mlp.experts.{self.config.n_routed_experts}.{suffix}"
yield name, weight
weights = iter_weights_with_fused_shared_experts(weights)
# Params for weights, fp8 weight scales, fp8 activation scales
# (param_name, weight_name, expert_id, shard_id)
expert_params_mapping = FusedMoE.make_expert_params_mapping(
ckpt_gate_proj_name="gate_proj",
ckpt_down_proj_name="down_proj",
@@ -1192,20 +1266,17 @@ class Glm4MoeForCausalLM(nn.Module):
"enorm",
"hnorm",
]
else:
nextn_layer_prefix = None
nextn_spec_weight_names = []
if params_dict is None:
params_dict = dict(self.named_parameters())
weight_names = []
for name, loaded_weight in weights:
weight_names.append(name)
if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
# Map shared expert weights to the last expert slot
# Shared expert becomes expert ID = n_routed_experts
name = name.replace(
"mlp.shared_experts",
f"mlp.experts.{self.config.n_routed_experts}",
)
if not is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = self.config.num_nextn_predict_layers
@@ -1217,9 +1288,10 @@ class Glm4MoeForCausalLM(nn.Module):
):
continue
else:
if not name.startswith(nextn_layer_prefix):
if nextn_layer_prefix and not name.startswith(nextn_layer_prefix):
continue
if nextn_layer_prefix is not None: # mtp
# Use shared head and embed weights from target model
if "shared_head.head" in name or "embed_tokens" in name:
continue
@@ -1295,6 +1367,7 @@ class Glm4MoeForCausalLM(nn.Module):
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
if name not in params_dict:
continue
python/sglang/srt/models/glm4_moe_lite.py
+9 -29
View File
@@ -12,13 +12,15 @@
# limitations under the License.
# ==============================================================================
"""Inference-only GLM-Lite model compatible with HuggingFace weights"""
"""Inference-only GLM-4.7-Flash model compatible with HuggingFace weights"""
import logging
import re
from typing import Iterable, Optional, Tuple
import torch
import torch.nn.functional as F
from sgl_kernel import dsv3_router_gemm
from torch import nn
from transformers import PretrainedConfig
@@ -29,12 +31,14 @@ from sglang.srt.distributed import (
get_tensor_model_parallel_world_size,
)
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.communicator import (
LayerCommunicator,
LayerScatterModes,
enable_moe_dense_fully_dp,
)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
is_dp_attention_enabled,
)
@@ -72,6 +76,7 @@ from sglang.srt.utils import (
log_info_on_rank0,
make_layers,
)
from sglang.srt.utils.hf_transformers_utils import get_rope_config
_is_cuda = is_cuda()
_device_sm = get_device_sm()
@@ -183,7 +188,6 @@ class Glm4MoeLiteGate(nn.Module):
and self.weight.shape[0] == 256
and _device_sm >= 90
):
from sgl_kernel import dsv3_router_gemm
logits = dsv3_router_gemm(hidden_states, self.weight).to(
hidden_states.dtype
@@ -335,12 +339,8 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer):
nn.Module.__init__(self)
self.hidden_size = config.hidden_size
self.config = config
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
rope_theta = 1000000
rope_scaling = None
rope_theta, rope_scaling = get_rope_config(config)
max_position_embeddings = getattr(config, "max_position_embeddings", 202752)
self.layer_id = layer_id
@@ -429,8 +429,6 @@ class Glm4MoeLiteModel(DeepseekV2Model):
self.pp_group = get_pp_group()
# DeepseekV2Model.forward expects these attributes to exist.
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
self.cp_size = get_attention_tp_size() if self.nsa_enable_prefill_cp else None
self.gemm_output_zero_allocator_size = 0
@@ -501,15 +499,8 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
)
self.capture_aux_hidden_states = False
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.nsa_enable_prefill_cp:
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
)
self.cp_rank = get_attention_tp_rank()
self.cp_size = get_attention_tp_size()
else:
@@ -549,7 +540,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
weights: Iterable[Tuple[str, torch.Tensor]],
is_nextn=False,
params_dict=None,
is_eagle=False,
):
if is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"):
@@ -579,7 +569,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
def iter_weights_with_fused_shared_experts(
weights: Iterable[Tuple[str, torch.Tensor]],
) -> Iterable[Tuple[str, torch.Tensor]]:
import re
pattern = re.compile(
r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$"
@@ -621,13 +610,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
nextn_layer_prefix = None
nextn_spec_weight_names = []
eagle_ignore_weight_names = []
if is_eagle:
eagle_ignore_weight_names = [
"eagle_draft_tokens_map",
"eagle_lm_head.weight",
]
if params_dict is None:
params_dict = dict(self.named_parameters())
@@ -635,7 +617,7 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
for name, loaded_weight in weights:
weight_names.append(name)
if not is_nextn and not is_eagle:
if not is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = self.config.num_nextn_predict_layers
if num_nextn_layers > 0 and name.startswith("model.layers"):
@@ -725,8 +707,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
if name in eagle_ignore_weight_names:
continue
# GLM NOTE: for MLA
if fuse_qkv_a_proj and (
@@ -797,7 +777,7 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
# DeepseekV2AttentionMLA.forward_* expects post_load_weights() to populate
# per-layer packed weights like `w_kc`/`w_vc` (used during CUDA graph capture).
# GLM-Lite configs may not set `config.mla`, but this model always uses
# GLM-4.7-Flash configs not set `config.mla`, but this model always uses
# DeepseekV2AttentionMLA, so we must run the post-load processing.
# Use weight_names=None to ensure we always process all layers. Some checkpoints /
# naming schemes may not include "kv_b_proj" in `weight_names`, but `w_kc`/`w_vc`