fix: load the right mtp lm head quantization (#30976)

2026-07-15 12:12:04 -07:00
parent 7e7129acd7
commit ab627e5d75
1 changed files with 5 additions and 0 deletions
@@ -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]