alejandro-rodriguez/tiny-random-nemotron-h

Tiny random NemotronH model for the OpenVINO tests of optimum-intel. The tokenizer is the one of nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.

Here is the code to create this tiny model:

"""Create a tiny random NemotronH model with every block type (Mamba2, MoE, attention, MLP).

Built with transformers 5.15.1. The layer types are saved as `hybrid_override_pattern` instead of
`layers_block_type`: transformers 5.13 renamed them ("mamba" -> "linear_attention", "attention" -> "full_attention")
and 5.6 to 5.12 only accept the old names, while every version with NemotronH converts the pattern.

Usage: python create_tiny_random_nemotron_h.py [output_dir]
"""

import json
import sys
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, NemotronHConfig


# Mamba2, MoE, attention, Mamba2, MLP, attention
HYBRID_OVERRIDE_PATTERN = "ME*M-*"

output_dir = Path(sys.argv[1] if len(sys.argv) > 1 else "./tiny-random-nemotron-h")

tokenizer = AutoTokenizer.from_pretrained("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16")
config = NemotronHConfig(
    vocab_size=len(tokenizer),
    hidden_size=64,
    hybrid_override_pattern=HYBRID_OVERRIDE_PATTERN,
    num_attention_heads=4,
    num_key_value_heads=2,
    # differs from hidden_size // num_attention_heads, as in NVIDIA-Nemotron-3-Nano-4B
    head_dim=32,
    intermediate_size=128,
    mamba_num_heads=8,
    mamba_head_dim=8,
    n_groups=2,
    ssm_state_size=16,
    conv_kernel=4,
    # several chunks for the reference SSD scan of the prompts used in tests
    chunk_size=8,
    n_routed_experts=8,
    num_experts_per_tok=2,
    moe_intermediate_size=32,
    moe_shared_expert_intermediate_size=64,
    routed_scaling_factor=2.5,
    n_group=1,
    topk_group=1,
    tie_word_embeddings=False,
)
config.dtype = "float32"
torch.manual_seed(42)
model = AutoModelForCausalLM.from_config(config).float().eval()

output_dir.mkdir(parents=True, exist_ok=True)
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

config_file = output_dir / "config.json"
saved_config = json.loads(config_file.read_text(encoding="utf-8"))
saved_config.pop("layers_block_type")
saved_config.pop("mtp_layers_block_type", None)
saved_config["hybrid_override_pattern"] = HYBRID_OVERRIDE_PATTERN
config_file.write_text(json.dumps(saved_config, indent=2, sort_keys=True) + "\n", encoding="utf-8")

The model.safetensors here is the output of this script with transformers 5.15.1, with the embeddings stored as backbone.embeddings.weight, as in NVIDIA's checkpoints. transformers 5.15 saves them as backbone.embedding.weight, which 5.16 and later no longer load (huggingface/transformers#48075).