nemotron

nemotron

Nemotron (Nemotron-4 / Minitron) (NemotronForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward. The norms are NemotronLayerNorm1P (a LayerNorm whose weight is stored minus one); the rotary covers partial_rotary_factor of each head; the MLP has no gate (down_proj(relu2(up_proj(x)))).

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm'}

ENVOYS module-attribute

ENVOYS = {NemotronDecoderLayer: Layer, NemotronAttention: Attention, NemotronMLP: Mlp}