Enable NemotronHPuzzle + NemotronHMTP (#37803)

This commit is contained in:
Netanel Haber
2026-03-22 17:13:58 +02:00
committed by GitHub
parent eaf4978621
commit e74c17e153
2 changed files with 6 additions and 3 deletions

View File

@@ -280,7 +280,7 @@ class SpeculativeConfig:
)
if (
hf_config.model_type == "nemotron_h"
hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"}
and hasattr(hf_config, "num_nextn_predict_layers")
and hf_config.num_nextn_predict_layers > 0
):

View File

@@ -395,13 +395,16 @@ class NemotronHMTP(nn.Module, SupportsPP):
]
expert_params_mapping = []
if hasattr(self.config, "n_routed_experts") and self.config.n_routed_experts:
num_experts = getattr(self.config, "n_routed_experts", None)
if getattr(self.config, "model_type", None) == "nemotron_h_puzzle":
num_experts = self.config.mtp_n_routed_experts
if num_experts is not None:
expert_params_mapping = FusedMoE.make_expert_params_mapping(
self,
ckpt_gate_proj_name="up_proj",
ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="", # Empty - non-gated MoE
num_experts=self.config.n_routed_experts,
num_experts=num_experts,
num_redundant_experts=self.num_redundant_experts,
)