diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 827f60e0f..c5ed94b47 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -54,7 +54,7 @@ from sglang.srt.layers.moe.utils import get_moe_a2a_backend from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope -from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -967,6 +967,7 @@ class MiniMaxM2ForCausalLM(nn.Module): self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) + self.pp_group = get_pp_group() # For EAGLE3 self.capture_aux_hidden_states = False @@ -999,17 +1000,26 @@ class MiniMaxM2ForCausalLM(nn.Module): positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - # _print_tensor_info(input_ids, "input_ids") - hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states - return self.logits_processor( - input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states - ) + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states + ) + else: + return hidden_states def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): """Load model weights with proper mapping for MiniMax architecture.""" @@ -1038,6 +1048,17 @@ class MiniMaxM2ForCausalLM(nn.Module): if "rotary_emb.inv_freq" in name: continue + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is not None: continue # skip spec decode layers for main model @@ -1056,7 +1077,10 @@ class MiniMaxM2ForCausalLM(nn.Module): 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: + if name not in params_dict: + continue + + if name.endswith(".bias"): continue param = params_dict[name] @@ -1070,6 +1094,8 @@ class MiniMaxM2ForCausalLM(nn.Module): continue name = name.replace(weight_name, param_name) + if name not in params_dict: + continue param = params_dict[name] weight_loader = param.weight_loader weight_loader( @@ -1090,6 +1116,8 @@ class MiniMaxM2ForCausalLM(nn.Module): if name is None: continue + if name not in params_dict: + continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader