nemotron_h_model_descriptor

Classes

NemotronHExpertRemovalLayerDescriptor

NemotronHFFNIntermediateLayerDescriptor

NemotronHKVHeadsLayerDescriptor

NemotronHMoELayerDescriptor

NemotronHMambaLayerDescriptor

NemotronHModelDescriptor

class NemotronHExpertRemovalLayerDescriptor

Bases: ExpertRemovalLayerDescriptor

__init__(target_name='mixer.gate', moe_prefix_name='model.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 = 'model.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 NemotronHFFNIntermediateLayerDescriptor

Bases: FFNIntermediateLayerDescriptor

__init__(down_proj_name='mixer.down_proj', ffn_prefix_name='model.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 = 'model.layers.{layer_idx}.mixer'
linear_weight_names: List[str]
class NemotronHKVHeadsLayerDescriptor

Bases: KVHeadsLayerDescriptor

__init__(o_proj_name='mixer.o_proj', attn_prefix_name='model.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 = 'model.layers.{layer_idx}.mixer'
o_proj_name: str = 'mixer.o_proj'
qkvo_weight_names: List[str]
class NemotronHMambaLayerDescriptor

Bases: MambaLayerDescriptor

__init__(target_name='mixer.in_proj', mamba_prefix_name='model.layers.{layer_idx}.mixer', in_proj_name='in_proj', out_proj_name='out_proj')
Parameters:

  • target_name (str)

  • mamba_prefix_name (str)

  • in_proj_name (str)

  • out_proj_name (str)

Return type:

None

mamba_prefix_name: str = 'model.layers.{layer_idx}.mixer'
class NemotronHMoELayerDescriptor

Bases: MoELayerDescriptor

__init__(target_name='mixer', moe_prefix_name='model.layers.{layer_idx}.mixer', gate_name='gate', experts_name='experts', shared_experts_name='shared_experts', latent_fc1_name='fc1_latent_proj', latent_fc2_name='fc2_latent_proj', require_attrs=(), expert_down_proj_regex='\\.experts\\.(\\d+)\\.down_proj$')
Parameters:

  • target_name (str)

  • moe_prefix_name (str)

  • gate_name (str)

  • experts_name (str)

  • shared_experts_name (str)

  • latent_fc1_name (str)

  • latent_fc2_name (str)

  • require_attrs (tuple[str, ...])

  • expert_down_proj_regex (str)

Return type:

None

moe_prefix_name: str = 'model.layers.{layer_idx}.mixer'
class NemotronHModelDescriptor

Bases: ModelDescriptor

classmethod adapt_loaded_state_dict_for_model(state_dict, *, model, config, checkpoint_to_model_key=None)
Parameters:

  • state_dict (dict[str, Any])

  • model (Module)

  • config (Any)

  • checkpoint_to_model_key (dict[str, str] | None)

Return type:

dict[str, Any]

classmethod adapt_materialized_state_dict_for_model(state_dict, *, model, config)
Parameters:

  • state_dict (dict[str, Any])

  • model (Module)

  • config (Any)

Return type:

dict[str, Any]

classmethod adapt_module_name_for_model(module_name, model)
Parameters:

  • module_name (str)

  • model (Module)

Return type:

str

classmethod anymodel_arch_info()
Return type:

dict

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

block_config (BlockConfig)

classmethod checkpoint_equivalence_tolerances()

Account for measured BF16 reduction-order drift after a residual permutation.

Structural-only sorting moved LM loss by 3.89e-4. Adding the exact global hidden-basis permutation moved it by 2.714e-3 while KL remained 3.088e-3 and top-1 agreement 0.9829. The transform is tensor-exact; the additional output drift comes from GEMM accumulation order across 52 hybrid layers.

Return type:

dict[str, float]

classmethod checkpoint_key_candidates_for_model_key(model_key, *, model, config)
Parameters:

  • model_key (str)

  • model (Module)

  • config (Any)

Return type:

tuple[str, …]

classmethod create_dummy_block(original_layer, block_index)
Parameters:

  • original_layer (Module)

  • block_index (int)

Return type:

Module

classmethod create_runtime_benchmark_model(runtime_config, block_configs)
Parameters:

block_configs (list[BlockConfig])

static decoder_layer_cls()
classmethod embedding_pruning_spec(config, *, widths, alignment)

Describe Nemotron-H’s shared residual width across hybrid layers.

The consolidated checkpoint uses backbone plus split experts while native AutoModel uses model plus fused expert tensors. Rules cover both representations so ranking, runtime prefix views, and physical HF materialization share one descriptor-owned contract.

Parameters:

alignment (int)

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]

classmethod local_kd_subblock_module_paths(block_config, *, layer_idx)

Map each hybrid sublayer to Nemotron-H’s exclusive mixer module.

Parameters:
Return type:

dict[tuple[str, str], str]

static mlp_no_op_post_init(decoder_layer)
static output_embedding_name()
classmethod patch_pipeline_model_part(model_part)

Install transient aliases expected by NeMo’s generic HF PP forward.

The aliases live only on the stage-local model chunk: embed_tokens points at NemotronH’s embeddings, norm points at norm_f, and rotary_emb is explicitly absent. This lets NeMo’s patched pipeline forward drive NemotronH without changing the checkpoint or the NeMo fork.

Parameters:

model_part (Module)

Return type:

bool

classmethod pipeline_module_fqns_per_model_part(config, *, pp_size, pipeline_config=None)

Use NemotronH’s real HF remote-code names for AutoModel PP splitting.

NeMo’s built-in HF splitter assumes model.embed_tokens and model.norm. NemotronH exposes the token embedding and final norm as model.embeddings and model.norm_f. If we let the generic split run, stage 0 drops the embedding and the first PP forward receives integer token ids with no embedding module. Keeping these names in the descriptor makes the PP layout explicit for NemotronH and gives future remote-code families one place to declare their own split names.

Parameters:

  • config (Any)

  • pp_size (int)

  • pipeline_config (dict[str, Any] | None)

Return type:

list[list[str]] | None

classmethod postprocess_runtime_benchmark_checkpoint(output_dir)
Parameters:

output_dir (Any)

Return type:

None

static pruning_mixins()
Return type:

Dict[str, PruningMixIn]

static puzzletron_capabilities(config)
static requires_trust_remote_code()
Return type:

bool

classmethod runtime_benchmark_base_block_config(runtime_config)
Return type:

BlockConfig

classmethod runtime_benchmark_config_fields(lm_config)
Return type:

dict[str, Any]

classmethod runtime_benchmark_scaffold_policy(block_config)

Keep one attention cache anchor per PP stage for cacheless hybrid candidates.

Parameters:

block_config (BlockConfig)

Return type:

str

classmethod runtime_benchmark_sublayers_are_exclusive()

Nemotron-H hybrid layers select one of attention, Mamba, MLP, or MoE.

Return type:

bool

classmethod runtime_vllm_benchmark_args(config)
Parameters:

config (Any)

Return type:

list[str]

classmethod set_block_configs(model_config, block_configs)
Parameters:

block_configs (list[BlockConfig | dict])

Return type:

None

static sorted_teacher_layout_kwargs(lm_config)
Return type:

dict

static truncate_pattern_for_subblock(lm_config, parent_layer_index=None)
Parameters:

parent_layer_index (int | None)

Return type:

None

classmethod update_runtime_benchmark_config(config_data)
Parameters:

config_data (dict[str, Any])

Return type:

None