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")