import logging from collections.abc import Iterable from typing import Optional import torch import torch.nn as nn from transformers import PretrainedConfig from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.step3p5 import Step3p5DecoderLayer, Step3p5ForCausalLM from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) def get_spec_layer_idx_from_weight_name( config: PretrainedConfig, weight_name: str ) -> Optional[int]: """Return MTP/nextn layer index if this weight belongs to spec layers. Step3p5 MTP/nextn checkpoints append extra layers after the main decoder: model.layers.[num_hidden_layers ... num_hidden_layers + num_nextn_predict_layers) """ if hasattr(config, "num_nextn_predict_layers") and ( getattr(config, "num_nextn_predict_layers", 0) > 0 ): base = config.num_hidden_layers for i in range(config.num_nextn_predict_layers): if weight_name.startswith(f"model.layers.{base + i}."): return base + i return None class SharedHead(nn.Module): def __init__( self, config, quant_config=None, ) -> None: super().__init__() self.norm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps) self.head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=quant_config ) self.lm_head = self.head def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.norm(hidden_states) class Step3p5AMultiTokenPredictor(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, ) self.mtp_start_layer_idx = config.num_hidden_layers self.num_mtp_layers = config.num_nextn_predict_layers layer_id = 45 # FIXME self.enorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps) self.hnorm = GemmaRMSNorm(config.hidden_size, config.rms_norm_eps) self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) self.shared_head = SharedHead(config=config, quant_config=quant_config) self.mtp_block = Step3p5DecoderLayer( config=config, layer_id=layer_id, prefix=f"{prefix}.mtp_block" ) self.lm_head = self.shared_head.head def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, ) -> torch.Tensor: if input_embeds is None: hidden_states = self.embed_tokens(input_ids) else: hidden_states = input_embeds if hidden_states.shape[0] > 0: hidden_states = self.eh_proj( torch.cat( ( self.enorm(hidden_states), self.hnorm(forward_batch.spec_info.hidden_states), ), dim=-1, ) ) hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, residual=None, ) hidden_states_before_norm = None if not forward_batch.forward_mode.is_idle(): # if forward_batch.return_hidden_states_before_norm: hidden_states_before_norm = ( hidden_states if residual is None else hidden_states + residual ) if residual is not None: hidden_states, _ = self.shared_head.norm(hidden_states, residual) else: hidden_states = self.shared_head.norm(hidden_states) return hidden_states, hidden_states_before_norm def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) # Chain-style multi-layer MTP (standard Step-3.5 Flash design): # each MTP layer consumes the hidden states produced by the preceding MTP layer, # while layer-0 consumes the hidden states from the target model. # The chain propagation is driven by MultiLayerEagleDraftWorker via the # ``chain_mtp_hidden_states`` flag: between speculative steps it overwrites # ``forward_batch.spec_info.hidden_states`` (and the CUDA-graph hidden_states # buffer in the draft-extend graph) with the previous layer's # ``hidden_states_before_norm`` returned by ``Step3p5AMultiTokenPredictor``. class Step3p5MTP(Step3p5ForCausalLM): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, draft_model_idx: Optional[int] = None, prefix: str = "", ) -> None: nn.Module.__init__(self) self.config = config self.tp_size = get_parallel().tp_size self.quant_config = quant_config self.draft_model_idx = draft_model_idx self.model = Step3p5AMultiTokenPredictor( config=config, quant_config=quant_config, prefix=add_prefix("model", prefix) ) self.logits_processor = LogitsProcessor(config) self.lm_head = self.model.lm_head def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: hidden_states, hidden_states_before_norm = self.model( input_ids, positions, forward_batch ) return self.logits_processor( input_ids, hidden_states, self.model.shared_head.head, forward_batch, hidden_states_before_norm=hidden_states_before_norm, ) def get_embed_and_head(self): return self.model.embed_tokens.weight, self.model.shared_head.head.weight def set_embed_and_head(self, embed, head): return def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), ] expert_params_mapping = [ (".moe.experts.w13_weight", ".moe.gate_proj.weight", "w1"), (".moe.experts.w13_weight", ".moe.up_proj.weight", "w3"), (".moe.experts.w2_weight", ".moe.down_proj.weight", "w2"), ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is not None and spec_layer != ( self.config.num_hidden_layers + self.draft_model_idx ): continue if "embed_tokens" not in name and spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: continue # We have mlp.experts[0].gate_proj in the checkpoint. # Since we handle the experts below in expert_params_mapping, # we need to skip here BEFORE we update the name, otherwise # name will be updated to mlp.experts[0].gate_up_proj, which # will then be updated below in expert_params_mapping # for mlp.experts[0].gate_gate_up_proj, which breaks load. if ("mlp.experts." in name) and name not in params_dict: continue if "experts" in name or "moe" in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if ( name.endswith(".bias") or name.endswith("_bias") ) and name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader for expert_id in range(loaded_weight.shape[0]): loaded_weight_expert = loaded_weight[expert_id] weight_loader( param, loaded_weight_expert, name, shard_id=shard_id, expert_id=expert_id, ) loaded_params.add(name) break else: # Skip loading extra bias for GPTQ models. if ( name.endswith(".bias") and name not in params_dict or "tok_embeddings" in name ): continue if "shared_head" in name: name = name.replace("shared_head.output", "shared_head.head") if "embed_tokens" in name: assert ( hasattr(self.config, "num_nextn_predict_layers") and self.config.num_nextn_predict_layers > 0 ) name = "model.embed_tokens.weight" param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(name) params_need_to_load = set(params_dict.keys()) if params_need_to_load != loaded_params: missing_params = list(params_need_to_load - loaded_params) param_name_example = missing_params[0] raise RuntimeError( f"Some parameters like {param_name_example} are not in the checkpoint and will falsely use random initialization" ) return loaded_params def _rewrite_spec_layer_name(self, spec_layer: Optional[int], name: str) -> str: """ Rewrite the weight name to match the format of the original model. Add .mtp_block for modules in transformer layer block for spec layer """ if spec_layer is None: return name # Some checkpoints place MTP weights under "model.layers..transformer.*". # Our modules use "model.layers..*", so drop the ".transformer." segment. transformer_prefix = f"model.layers.{spec_layer}.transformer." if name.startswith(transformer_prefix): name = name.replace(".transformer.", ".", 1) spec_layer_weight_names = [ "embed_tokens", "enorm", "hnorm", "eh_proj", "shared_head", ] spec_layer_weight = False for weight_name in spec_layer_weight_names: if weight_name in name: spec_layer_weight = True break if not spec_layer_weight: # treat rest weights as weights for transformer layer block name = name.replace( f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block." ) # NEW: drop "layers.." from the rewritten name (minimal change). layers_prefix = f"model.layers.{spec_layer}." if name.startswith(layers_prefix): name = name.replace(layers_prefix, "model.", 1) return name EntryClass = [Step3p5MTP]