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

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
+130 -57
View File
@@ -15,6 +15,7 @@
"""Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 model compatible with HuggingFace weights""" """Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 model compatible with HuggingFace weights"""
import logging import logging
import re
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
import torch import torch
@@ -22,6 +23,7 @@ import torch.nn.functional as F
from torch import nn from torch import nn
from transformers import PretrainedConfig 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.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import ( from sglang.srt.distributed import (
get_moe_expert_parallel_world_size, 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.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE 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.topk import TopK
from sglang.srt.layers.moe.utils import ( from sglang.srt.layers.moe.utils import (
RoutingMethodType, RoutingMethodType,
@@ -183,7 +186,7 @@ class Glm4MoeAttention(nn.Module):
num_heads: int, num_heads: int,
num_kv_heads: int, num_kv_heads: int,
layer_id: int = 0, layer_id: int = 0,
rope_theta: float = 10000, rope_theta: float = 1000000,
partial_rotary_factor: float = 0.5, partial_rotary_factor: float = 0.5,
rope_scaling: Optional[Dict[str, Any]] = None, rope_scaling: Optional[Dict[str, Any]] = None,
max_position_embeddings: int = 8192, max_position_embeddings: int = 8192,
@@ -439,9 +442,12 @@ class Glm4MoeSparseMoeBlock(nn.Module):
fused_shared_experts_scaling_factor=1, 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: 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 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( self.shared_experts = Glm4MoeMLP(
hidden_size=config.hidden_size, hidden_size=config.hidden_size,
intermediate_size=intermediate_size, intermediate_size=intermediate_size,
@@ -453,16 +459,54 @@ class Glm4MoeSparseMoeBlock(nn.Module):
dict(tp_rank=0, tp_size=1) dict(tp_rank=0, tp_size=1)
if get_moe_a2a_backend().is_deepep() if get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake() 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 get_moe_a2a_backend().is_flashinfer()
or should_use_flashinfer_cutlass_moe_fp4_allgather() or should_use_flashinfer_cutlass_moe_fp4_allgather()
else {} 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 ( if (
get_moe_a2a_backend().is_deepep() get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl() 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 # TODO: we will support tp < ep in the future
self.ep_size = get_moe_expert_parallel_world_size() self.ep_size = get_moe_expert_parallel_world_size()
@@ -483,7 +527,11 @@ class Glm4MoeSparseMoeBlock(nn.Module):
get_moe_a2a_backend().is_deepep() get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_nixl() 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): def get_moe_weights(self):
return [ return [
@@ -502,8 +550,7 @@ class Glm4MoeSparseMoeBlock(nn.Module):
should_allreduce_fusion: bool = False, should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False, use_reduce_scatter: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
if not self._enable_a2a_moe:
if not get_moe_a2a_backend().is_deepep():
if ( if (
self.alt_stream is not None self.alt_stream is not None
and self.num_fused_shared_experts == 0 and self.num_fused_shared_experts == 0
@@ -511,11 +558,15 @@ class Glm4MoeSparseMoeBlock(nn.Module):
and get_is_capture_mode() and get_is_capture_mode()
): ):
return self.forward_normal_dual_stream( return self.forward_normal_dual_stream(
hidden_states, should_allreduce_fusion, use_reduce_scatter hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
) )
else: else:
return self.forward_normal( return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter hidden_states,
should_allreduce_fusion,
use_reduce_scatter,
) )
else: else:
return self.forward_deepep(hidden_states, forward_batch) return self.forward_deepep(hidden_states, forward_batch)
@@ -534,20 +585,12 @@ class Glm4MoeSparseMoeBlock(nn.Module):
# router_logits: (num_tokens, n_experts) # router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states) router_logits = self.gate(hidden_states)
topk_output = self.topk(hidden_states, router_logits) topk_output = self.topk(hidden_states, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output) final_hidden_states = self.experts(hidden_states, topk_output)
if not _is_cuda and not _use_aiter: if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
# fused in biased_grouped_topk so we can skip here
final_hidden_states *= self.routed_scaling_factor final_hidden_states *= self.routed_scaling_factor
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
final_hidden_states += shared_output
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
if ( if (
self.tp_size > 1 self.tp_size > 1
and not should_allreduce_fusion and not should_allreduce_fusion
@@ -1090,27 +1133,32 @@ class Glm4MoeForCausalLM(nn.Module):
# For EAGLE3 support # For EAGLE3 support
self.capture_aux_hidden_states = False 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): def determine_num_fused_shared_experts(self):
if get_global_server_args().disable_shared_experts_fusion: if get_global_server_args().disable_shared_experts_fusion:
return return
disable_reason = None disable_reason = None
if not getattr(self.config, "n_shared_experts", None): if (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and (
disable_reason = "No shared experts are defined in the config." not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
elif not _is_cuda: ):
disable_reason = "Shared experts fusion currently requires CUDA devices." disable_reason = (
elif _is_cuda and (_device_sm is not None) and (_device_sm < 80): "Only GLM-4.5 on NV-platform with capability >= 80 "
disable_reason = "Shared experts fusion requires SM80 or newer GPUs." "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
elif get_moe_expert_parallel_world_size() > 1: )
disable_reason = "Shared experts fusion is not supported together with expert parallelism yet." elif get_moe_expert_parallel_world_size() > 1 and (
elif get_moe_a2a_backend().is_deepep(): not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled." ):
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: if disable_reason is not None:
get_global_server_args().disable_shared_experts_fusion = True get_global_server_args().disable_shared_experts_fusion = True
self.num_fused_shared_experts = 0
log_info_on_rank0( log_info_on_rank0(
logger, logger,
f"{disable_reason} Shared experts fusion optimization is disabled.", f"{disable_reason} Shared experts fusion optimization is disabled.",
@@ -1118,10 +1166,9 @@ class Glm4MoeForCausalLM(nn.Module):
return return
self.num_fused_shared_experts = self.config.n_shared_experts self.num_fused_shared_experts = self.config.n_shared_experts
assert (
self.num_fused_shared_experts == 1 def get_input_embeddings(self) -> nn.Embedding:
), "Only 1 fused shared expert is supported for Glm4MoeForCausalLM" return self.model.embed_tokens
log_info_on_rank0(logger, "Shared experts fusion optimization enabled.")
@torch.no_grad() @torch.no_grad()
def forward( def forward(
@@ -1154,7 +1201,12 @@ class Glm4MoeForCausalLM(nn.Module):
def end_layer(self): def end_layer(self):
return self.model.end_layer 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 is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"): if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = 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), ("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( expert_params_mapping = FusedMoE.make_expert_params_mapping(
ckpt_gate_proj_name="gate_proj", ckpt_gate_proj_name="gate_proj",
ckpt_down_proj_name="down_proj", ckpt_down_proj_name="down_proj",
@@ -1192,20 +1266,17 @@ class Glm4MoeForCausalLM(nn.Module):
"enorm", "enorm",
"hnorm", "hnorm",
] ]
else:
nextn_layer_prefix = None
nextn_spec_weight_names = []
if params_dict is None:
params_dict = dict(self.named_parameters())
params_dict = dict(self.named_parameters())
weight_names = [] weight_names = []
for name, loaded_weight in weights: for name, loaded_weight in weights:
weight_names.append(name) 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 not is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"): if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = self.config.num_nextn_predict_layers num_nextn_layers = self.config.num_nextn_predict_layers
@@ -1217,23 +1288,24 @@ class Glm4MoeForCausalLM(nn.Module):
): ):
continue continue
else: else:
if not name.startswith(nextn_layer_prefix): if nextn_layer_prefix and not name.startswith(nextn_layer_prefix):
continue continue
# Use shared head and embed weights from target model if nextn_layer_prefix is not None: # mtp
if "shared_head.head" in name or "embed_tokens" in name: # Use shared head and embed weights from target model
continue if "shared_head.head" in name or "embed_tokens" in name:
continue
is_decoder = True is_decoder = True
# For nextn specific weights # For nextn specific weights
for weight_name in nextn_spec_weight_names: for weight_name in nextn_spec_weight_names:
if weight_name in name: if weight_name in name:
name = name.replace(nextn_layer_prefix, "model") name = name.replace(nextn_layer_prefix, "model")
is_decoder = False is_decoder = False
break break
# For decoder layer weights # For decoder layer weights
if is_decoder: if is_decoder:
name = name.replace(nextn_layer_prefix, "model.decoder") name = name.replace(nextn_layer_prefix, "model.decoder")
if "rotary_emb.inv_freq" in name: if "rotary_emb.inv_freq" in name:
continue continue
@@ -1295,6 +1367,7 @@ class Glm4MoeForCausalLM(nn.Module):
# Skip loading extra bias for GPTQ models. # Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict: if name.endswith(".bias") and name not in params_dict:
continue continue
if name not in params_dict: if name not in params_dict:
continue continue
+9 -29
View File
@@ -12,13 +12,15 @@
# limitations under the License. # 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 logging
import re
from typing import Iterable, Optional, Tuple from typing import Iterable, Optional, Tuple
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from sgl_kernel import dsv3_router_gemm
from torch import nn from torch import nn
from transformers import PretrainedConfig from transformers import PretrainedConfig
@@ -29,12 +31,14 @@ from sglang.srt.distributed import (
get_tensor_model_parallel_world_size, get_tensor_model_parallel_world_size,
) )
from sglang.srt.layers.activation import SiluAndMul 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 ( from sglang.srt.layers.communicator import (
LayerCommunicator, LayerCommunicator,
LayerScatterModes, LayerScatterModes,
enable_moe_dense_fully_dp, enable_moe_dense_fully_dp,
) )
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size, get_attention_tp_size,
is_dp_attention_enabled, is_dp_attention_enabled,
) )
@@ -72,6 +76,7 @@ from sglang.srt.utils import (
log_info_on_rank0, log_info_on_rank0,
make_layers, make_layers,
) )
from sglang.srt.utils.hf_transformers_utils import get_rope_config
_is_cuda = is_cuda() _is_cuda = is_cuda()
_device_sm = get_device_sm() _device_sm = get_device_sm()
@@ -183,7 +188,6 @@ class Glm4MoeLiteGate(nn.Module):
and self.weight.shape[0] == 256 and self.weight.shape[0] == 256
and _device_sm >= 90 and _device_sm >= 90
): ):
from sgl_kernel import dsv3_router_gemm
logits = dsv3_router_gemm(hidden_states, self.weight).to( logits = dsv3_router_gemm(hidden_states, self.weight).to(
hidden_states.dtype hidden_states.dtype
@@ -335,12 +339,8 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer):
nn.Module.__init__(self) nn.Module.__init__(self)
self.hidden_size = config.hidden_size self.hidden_size = config.hidden_size
self.config = config 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() self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
rope_theta = 1000000 rope_theta, rope_scaling = get_rope_config(config)
rope_scaling = None
max_position_embeddings = getattr(config, "max_position_embeddings", 202752) max_position_embeddings = getattr(config, "max_position_embeddings", 202752)
self.layer_id = layer_id self.layer_id = layer_id
@@ -429,8 +429,6 @@ class Glm4MoeLiteModel(DeepseekV2Model):
self.pp_group = get_pp_group() self.pp_group = get_pp_group()
# DeepseekV2Model.forward expects these attributes to exist. # 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.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.cp_size = get_attention_tp_size() if self.nsa_enable_prefill_cp else None
self.gemm_output_zero_allocator_size = 0 self.gemm_output_zero_allocator_size = 0
@@ -501,15 +499,8 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
) )
self.capture_aux_hidden_states = False 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() self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
if self.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_rank = get_attention_tp_rank()
self.cp_size = get_attention_tp_size() self.cp_size = get_attention_tp_size()
else: else:
@@ -549,7 +540,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
weights: Iterable[Tuple[str, torch.Tensor]], weights: Iterable[Tuple[str, torch.Tensor]],
is_nextn=False, is_nextn=False,
params_dict=None, params_dict=None,
is_eagle=False,
): ):
if is_nextn: if is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"): if hasattr(self.config, "num_nextn_predict_layers"):
@@ -579,7 +569,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
def iter_weights_with_fused_shared_experts( def iter_weights_with_fused_shared_experts(
weights: Iterable[Tuple[str, torch.Tensor]], weights: Iterable[Tuple[str, torch.Tensor]],
) -> Iterable[Tuple[str, torch.Tensor]]: ) -> Iterable[Tuple[str, torch.Tensor]]:
import re
pattern = re.compile( pattern = re.compile(
r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$" r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$"
@@ -621,13 +610,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
nextn_layer_prefix = None nextn_layer_prefix = None
nextn_spec_weight_names = [] 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: if params_dict is None:
params_dict = dict(self.named_parameters()) params_dict = dict(self.named_parameters())
@@ -635,7 +617,7 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
for name, loaded_weight in weights: for name, loaded_weight in weights:
weight_names.append(name) weight_names.append(name)
if not is_nextn and not is_eagle: if not is_nextn:
if hasattr(self.config, "num_nextn_predict_layers"): if hasattr(self.config, "num_nextn_predict_layers"):
num_nextn_layers = 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"): if num_nextn_layers > 0 and name.startswith("model.layers"):
@@ -725,8 +707,6 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
# Skip loading extra bias for GPTQ models. # Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict: if name.endswith(".bias") and name not in params_dict:
continue continue
if name in eagle_ignore_weight_names:
continue
# GLM NOTE: for MLA # GLM NOTE: for MLA
if fuse_qkv_a_proj and ( if fuse_qkv_a_proj and (
@@ -797,7 +777,7 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
# DeepseekV2AttentionMLA.forward_* expects post_load_weights() to populate # DeepseekV2AttentionMLA.forward_* expects post_load_weights() to populate
# per-layer packed weights like `w_kc`/`w_vc` (used during CUDA graph capture). # 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. # DeepseekV2AttentionMLA, so we must run the post-load processing.
# Use weight_names=None to ensure we always process all layers. Some checkpoints / # 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` # naming schemes may not include "kv_b_proj" in `weight_names`, but `w_kc`/`w_vc`