nemotron_h_v2_model_descriptor#

Classes

NemotronHV2FFNIntermediateLayerDescriptor

NemotronHV2FFNIntermediateLayerDescriptor(down_proj_name: str = 'mixer.down_proj', ffn_prefix_name: str = 'backbone.layers.{layer_idx}.mixer', linear_weight_names: List[str] = <factory>)

NemotronHV2KVHeadsLayerDescriptor

NemotronHV2KVHeadsLayerDescriptor(o_proj_name: str = 'mixer.o_proj', attn_prefix_name: str = 'backbone.layers.{layer_idx}.mixer', qkvo_weight_names: List[str] = <factory>)

NemotronHV2ExpertRemovalLayerDescriptor

NemotronHV2ExpertRemovalLayerDescriptor(target_name: str = 'mixer.gate', moe_prefix_name: str = 'backbone.layers.{layer_idx}.mixer', expert_prefix_name: str = 'experts.{expert_idx}', router_weights: List[str] = <factory>, router_biases: List[str] = <factory>, expert_weights: List[str] = <factory>, expert_biases: List[str] = <factory>, is_fused_experts: bool = False, fused_expert_weights: List[str] = <factory>)

NemotronHV2ModelDescriptor

class NemotronHV2ExpertRemovalLayerDescriptor#

Bases: ExpertRemovalLayerDescriptor

NemotronHV2ExpertRemovalLayerDescriptor(target_name: str = ‘mixer.gate’, moe_prefix_name: str = ‘backbone.layers.{layer_idx}.mixer’, expert_prefix_name: str = ‘experts.{expert_idx}’, router_weights: List[str] = <factory>, router_biases: List[str] = <factory>, expert_weights: List[str] = <factory>, expert_biases: List[str] = <factory>, is_fused_experts: bool = False, fused_expert_weights: List[str] = <factory>)

__init__(target_name='mixer.gate', moe_prefix_name='backbone.layers.{layer_idx}.mixer', expert_prefix_name='experts.{expert_idx}', router_weights=<factory>, router_biases=<factory>, expert_weights=<factory>, expert_biases=<factory>, is_fused_experts=False, fused_expert_weights=<factory>)#
Parameters:
  • target_name (str)

  • moe_prefix_name (str)

  • expert_prefix_name (str)

  • router_weights (List[str])

  • router_biases (List[str])

  • expert_weights (List[str])

  • expert_biases (List[str])

  • is_fused_experts (bool)

  • fused_expert_weights (List[str])

Return type:

None

expert_prefix_name: str = 'experts.{expert_idx}'#

Expert prefix relative to moe_prefix with {expert_idx} placeholder, e.g. experts.{expert_idx}.

expert_weights: List[str]#

Per-expert weight names relative to expert_prefix (per-expert format).

get_modules_names_to_hook(model)#
Return type:

List[Tuple[int, str]]

moe_prefix_name: str = 'backbone.layers.{layer_idx}.mixer'#

MoE prefix layer name with {layer_idx} placeholder, e.g. model.layers.{layer_idx}.moe.

router_biases: List[str]#

Router bias names relative to moe_prefix.

router_weights: List[str]#

Router weight names relative to moe_prefix.

target_name: str = 'mixer.gate'#

Module name for hook registration; supports regex: prefix.

class NemotronHV2FFNIntermediateLayerDescriptor#

Bases: FFNIntermediateLayerDescriptor

NemotronHV2FFNIntermediateLayerDescriptor(down_proj_name: str = ‘mixer.down_proj’, ffn_prefix_name: str = ‘backbone.layers.{layer_idx}.mixer’, linear_weight_names: List[str] = <factory>)

__init__(down_proj_name='mixer.down_proj', ffn_prefix_name='backbone.layers.{layer_idx}.mixer', linear_weight_names=<factory>)#
Parameters:
  • down_proj_name (str)

  • ffn_prefix_name (str)

  • linear_weight_names (List[str])

Return type:

None

down_proj_name: str = 'mixer.down_proj'#
ffn_prefix_name: str = 'backbone.layers.{layer_idx}.mixer'#
linear_weight_names: List[str]#
class NemotronHV2KVHeadsLayerDescriptor#

Bases: KVHeadsLayerDescriptor

NemotronHV2KVHeadsLayerDescriptor(o_proj_name: str = ‘mixer.o_proj’, attn_prefix_name: str = ‘backbone.layers.{layer_idx}.mixer’, qkvo_weight_names: List[str] = <factory>)

__init__(o_proj_name='mixer.o_proj', attn_prefix_name='backbone.layers.{layer_idx}.mixer', qkvo_weight_names=<factory>)#
Parameters:
  • o_proj_name (str)

  • attn_prefix_name (str)

  • qkvo_weight_names (List[str])

Return type:

None

attn_prefix_name: str = 'backbone.layers.{layer_idx}.mixer'#
o_proj_name: str = 'mixer.o_proj'#
qkvo_weight_names: List[str]#
class NemotronHV2ModelDescriptor#

Bases: ModelDescriptor

static attn_no_op_post_init(decoder_layer)#
static block_config_to_layer_overrides(block_config)#
Parameters:

block_config (BlockConfig)

classmethod create_dummy_block(original_layer, block_index)#
Parameters:
  • original_layer (Module)

  • block_index (int)

Return type:

Module

static decoder_layer_cls()#
static final_norm_name()#
classmethod get_weight_groups(layer_names, num_hidden_layers)#

Problem with NemotronH is that norm.weight can be in both block_{i}_ffn and block_{i}_attention. duplicate groups with norm.weight should be removed.

Parameters:
  • layer_names (Iterable[str])

  • num_hidden_layers (int)

Return type:

Dict[str, List[str]]

static init_rotary_embedding(model, runtime)#

NemotronH has no positional embeddings

static input_embedding_name()#
static layer_block_name(index)#
Parameters:

index (int)

static layer_name_predicates(num_layers)#
Parameters:

num_layers (int)

Return type:

Dict[str, Pattern]

static mlp_no_op_post_init(decoder_layer)#
static output_embedding_name()#
static pruning_mixins()#
Return type:

Dict[str, PruningMixIn]

static requires_trust_remote_code()#
Return type:

bool