| ---
|
| library_name: transformers
|
| license: other
|
| license_name: nvidia-nemotron-open-model-license
|
| license_link: >-
|
| https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-nemotron-open-model-license/
|
| ---
|
|
|
| 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](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16).
|
|
|
| Here is the code to create this tiny model:
|
|
|
| ```python
|
| """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")
|
| ```
|
|
|
