Enable NemotronHPuzzle + NemotronHMTP (#37803)
This commit is contained in:
@@ -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
|
||||
):
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user