Unverified Commit e74c17e1 authored by Netanel Haber's avatar Netanel Haber Committed by GitHub
Browse files

Enable `NemotronHPuzzle` + `NemotronHMTP` (#37803)

parent eaf49786
...@@ -280,7 +280,7 @@ class SpeculativeConfig: ...@@ -280,7 +280,7 @@ class SpeculativeConfig:
) )
if ( 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 hasattr(hf_config, "num_nextn_predict_layers")
and hf_config.num_nextn_predict_layers > 0 and hf_config.num_nextn_predict_layers > 0
): ):
......
...@@ -395,13 +395,16 @@ class NemotronHMTP(nn.Module, SupportsPP): ...@@ -395,13 +395,16 @@ class NemotronHMTP(nn.Module, SupportsPP):
] ]
expert_params_mapping = [] 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( expert_params_mapping = FusedMoE.make_expert_params_mapping(
self, self,
ckpt_gate_proj_name="up_proj", ckpt_gate_proj_name="up_proj",
ckpt_down_proj_name="down_proj", ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="", # Empty - non-gated MoE 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, num_redundant_experts=self.num_redundant_experts,
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment