diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 61b1e3bba..d38de9b9d 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -369,5 +369,10 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM): ): super().load_weights(weights, is_mtp=True) + def set_lm_head_from_target(self, target_lm_head: nn.Module) -> None: + if self.config.tie_word_embeddings: + return + self.lm_head = target_lm_head + EntryClass = [NemotronHForCausalLMMTP]