diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 4d88bf325..a782c98e9 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -300,6 +300,10 @@ class ModelConfig: self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 + if is_draft_model and self.hf_config.architectures[0] == "NemotronHForCausalLM": + self.hf_config.architectures[0] = "NemotronHForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + def _derive_hybrid_model(self): # Use self.context_len after it has been initialized to prevent using context_len which may be None. self.is_hybrid_swa = is_hybrid_swa_model(self.hf_config.architectures) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6752ad5f3..e0de14115 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1485,6 +1485,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): @property def mamba2_config(self): config = self.model_config.hf_config + if isinstance(config, NemotronHConfig) and self.is_draft_worker: + # NemotronH MTP draft models have no Mamba layers (pattern like "*E") + # so they shouldn't use HybridLinearAttnBackend + pattern = getattr(config, "mtp_hybrid_override_pattern", None) + if pattern is not None and "M" not in pattern: + return None if isinstance(config, FalconH1Config | NemotronHConfig): return config if isinstance(config, NemotronH_Nano_VL_V2_Config): diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 76fa41e3c..a42b1f2c0 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -728,7 +728,20 @@ class NemotronHForCausalLM(nn.Module): def get_seqlen_agnostic_capture_inputs(self, batch_size: int): return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False + ) -> None: updated_weights = [] for name, loaded_weight in weights: name = replace_prefix(name, self.remap_prefix) @@ -749,6 +762,20 @@ class NemotronHForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) for name, loaded_weight in updated_weights: + if is_mtp: + if "mtp" not in name: + continue + + name = name.replace("mtp.layers.", "model.layers.") + + if "embeddings" in name: + name = name.replace("embeddings", "model.embed_tokens") + if name.startswith("backbone."): + name = name.replace("backbone.", "") + + if not is_mtp and "mtp" in name: + continue + if "scale" in name: name = maybe_remap_kv_scale_name(name, params_dict) if name is None: diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py new file mode 100644 index 000000000..0d42cf168 --- /dev/null +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -0,0 +1,340 @@ +# Copyright 2023-2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from collections.abc import Iterable + +import torch +from torch import nn + +from sglang.srt.configs import NemotronHConfig +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ColumnParallelLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization 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.models.nemotron_h import ( + NemotronHAttentionDecoderLayer, + NemotronHForCausalLM, + NemotronHMoEDecoderLayer, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix + + +class NemotronHMTPAttentionDecoderLayer(NemotronHAttentionDecoderLayer): + def __init__( + self, + config: NemotronHConfig, + layer_idx: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + has_start_projections: bool = False, + has_end_norm: bool = False, + ) -> None: + super().__init__( + config=config, + layer_idx=layer_idx, + quant_config=quant_config, + prefix=prefix, + ) + self.has_start_projections = has_start_projections + self.has_end_norm = has_end_norm + + if has_start_projections: + self.enorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.hnorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + # Fusion layer to combine embeddings with target hidden states + self.eh_proj = ColumnParallelLinear( + input_size=config.hidden_size * 2, + output_size=config.hidden_size, + bias=False, + gather_output=True, + params_dtype=( + config.dtype if hasattr(config, "dtype") else torch.bfloat16 + ), + quant_config=quant_config, + prefix=f"{prefix}.eh_proj", + ) + + if has_end_norm: + self.final_layernorm = RMSNorm( + config.hidden_size, + eps=getattr(config, "layer_norm_epsilon", 1e-5), + ) + + def forward( + self, + *, + inputs_embeds: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None = None, + forward_batch: ForwardBatch, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.has_start_projections: + inputs_embeds_normed = self.enorm(inputs_embeds) + previous_hidden_states_normed = self.hnorm(hidden_states) + + fused = torch.cat( + [inputs_embeds_normed, previous_hidden_states_normed], dim=-1 + ) + hidden_states, _ = self.eh_proj(fused) + + hidden_states, residual = super().forward( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + ) + + if self.has_end_norm: + if residual is not None: + hidden_states = hidden_states + residual + residual = None + + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states, residual + + +class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer): + def __init__( + self, + config: NemotronHConfig, + layer_idx: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + has_start_projections: bool = False, + has_end_norm: bool = False, + ) -> None: + super().__init__( + config=config, + layer_idx=layer_idx, + quant_config=quant_config, + prefix=prefix, + ) + self.has_start_projections = has_start_projections + self.has_end_norm = has_end_norm + + if has_start_projections: + self.enorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.hnorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + self.eh_proj = ColumnParallelLinear( + input_size=config.hidden_size * 2, + output_size=config.hidden_size, + bias=False, + gather_output=True, + params_dtype=( + config.dtype if hasattr(config, "dtype") else torch.bfloat16 + ), + quant_config=quant_config, + prefix=f"{prefix}.eh_proj", + ) + + if has_end_norm: + self.final_layernorm = RMSNorm( + config.hidden_size, + eps=getattr(config, "layer_norm_epsilon", 1e-5), + ) + + def forward( + self, + *, + inputs_embeds: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None = None, + forward_batch: ForwardBatch, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.has_start_projections: + inputs_embeds_normed = self.enorm(inputs_embeds) + previous_hidden_states_normed = self.hnorm(hidden_states) + + fused = torch.cat( + [inputs_embeds_normed, previous_hidden_states_normed], dim=-1 + ) + hidden_states, _ = self.eh_proj(fused) + + hidden_states, residual = super().forward( + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + ) + + if self.has_end_norm: + if residual is not None: + hidden_states = hidden_states + residual + residual = None + + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states, residual + + +class NemotronHMultiTokenPredictor(nn.Module): + def __init__( + self, + config: NemotronHConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.vocab_size = config.vocab_size + self.org_vocab_size = config.vocab_size + + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1) + assert ( + self.num_mtp_layers == 1 + ), "Only one MTP layer is supported for NemotronH-MTP" + + self.pattern_str = config.mtp_hybrid_override_pattern + self.pattern_len = len(self.pattern_str) + assert self.pattern_len > 0 + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + ) + + # Build flat list of layers + self.layers = nn.ModuleDict() + + # Total number of physical layers = num_steps * pattern_len + total_layers = self.num_mtp_layers * self.pattern_len + for i in range(total_layers): + step_rel_idx = i % self.pattern_len + + char = self.pattern_str[step_rel_idx] + + is_start_of_step = step_rel_idx == 0 + is_end_of_step = step_rel_idx == self.pattern_len - 1 + + layer_prefix = f"{prefix}.layers.{i}" + + common_kwargs = dict( + config=config, + layer_idx=i, + quant_config=quant_config, + prefix=layer_prefix, + has_start_projections=is_start_of_step, + has_end_norm=is_end_of_step, + ) + + if char == "*": + self.layers[str(i)] = NemotronHMTPAttentionDecoderLayer(**common_kwargs) + elif char == "E": + self.layers[str(i)] = NemotronHMTPMoEDecoderLayer(**common_kwargs) + else: + raise NotImplementedError( + f"Pattern char '{char}' in {self.pattern_str} not implemented" + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + assert ( + self.embed_tokens is not None + ), "embed_tokens not initialized - must be shared from target model" + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings(input_ids) + + residual = None + + for i in range(self.pattern_len): + hidden_states, residual = self.layers[str(i)]( + inputs_embeds=inputs_embeds, + hidden_states=hidden_states, + residual=residual, + forward_batch=forward_batch, + ) + return hidden_states + + +class NemotronHForCausalLMMTP(NemotronHForCausalLM): + def __init__( + self, + config: NemotronHConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + nn.Module.__init__(self) + self.config = config + self.quant_config = quant_config + # Required for parent's load_weights + self.pp_group = get_pp_group() + + # Override config for MTP pattern (which has no Mamba layers) + config.num_hidden_layers = len(config.mtp_hybrid_override_pattern) + # Set hybrid_override_pattern to MTP pattern so attention backend + # doesn't use Mamba2AttnBackend (MTP has no Mamba layers) + config.hybrid_override_pattern = config.mtp_hybrid_override_pattern + + self.model = NemotronHMultiTokenPredictor( + config=config, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = forward_batch.spec_info.hidden_states + + hidden_states = self.model( + input_ids, + hidden_states, + forward_batch, + input_embeds, + ) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False + ): + super().load_weights(weights, is_mtp=True) + + +EntryClass = [NemotronHForCausalLMMTP]