NemotronH-0.3B-A0.3B
This is a tiny version of nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-Base-BF16 created for testing and development.
Model Details
- Base Model: nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-Base-BF16
- Architecture: nemotron_h (hybrid Mamba-2 / attention / mixture-of-experts)
- Total Parameters: 0.347B
- Activated Parameters: ~0.309B
The nemotron_h architecture is a hybrid that interleaves three block types. This tiny
model preserves at least one of each so the full architecture is exercised:
linear_attention— Mamba-2 SSM mixerfull_attention— grouped-query attention mixermoe— mixture-of-experts FFN (non-gated experts + shared expert + latent projection)
The layer pattern is ["linear_attention", "moe", "linear_attention", "full_attention", "moe"].
Configuration Changes
The following parameters were reduced from the original model:
| Parameter | Original | Tiny |
|---|---|---|
| num layers (layers_block_type) | 108 | 5 |
| hidden_size | 8192 | 1024 |
| num_attention_heads | 64 | 8 |
| num_key_value_heads | 2 | 2 |
| head_dim | 128 | 128 |
| mamba_num_heads | 256 | 32 |
| mamba_head_dim | 64 | 64 |
| n_groups | 8 | 8 |
| intermediate_size | 5120 | 1536 |
| n_routed_experts | 512 | 32 |
| num_experts_per_tok | 22 | 8 |
| moe_intermediate_size | 5120 | 768 |
| moe_shared_expert_intermediate_size | 10240 | 1536 |
| moe_latent_size | 2048 | 512 |
| num_nextn_predict_layers | 1 | 0 |
| vocab_size | 131072 | 131072 (unchanged) |
Checkpoint Structure
Single-file checkpoint (model.safetensors). The tensor layout matches the original
model: the backbone. prefix is used, and routed experts are stored as individual
per-expert 2D tensors (backbone.layers.N.mixer.experts.M.up_proj.weight /
down_proj.weight) rather than stacked 3D parameters. All non-MTP tensor names are
identical (modulo layer/expert indices) to the original checkpoint. Multi-token-prediction
(mtp.*) layers are intentionally omitted (num_nextn_predict_layers = 0).
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("NemotronH-0.3B-A0.3B", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("NemotronH-0.3B-A0.3B")
input_ids = tokenizer("According to all known laws", return_tensors="pt").input_ids.to(model.device)
output = model.generate(input_ids, max_new_tokens=20)
print(tokenizer.decode(output[0]))
Creation Process
This model was created using the llm-compressor create-tiny-model claude skill.
- Reduced
layers_block_type, hidden/attention/mamba sizes, and MoE sizes as above. - Re-initialized all weights, then fine-tuned on a small toy dataset until the perplexity target was reached (validates the model can learn).
- Verified the saved checkpoint structure matches the original (excluding MTP layers).
Validation
Success: 1.0079255104064941 <= 10.0
==================================================
Generating sample text:
According to all known laws of aviation, there is no way a bee should be able to fly.
==================================================
Notes
- The Mamba mixer intermediate size is
mamba_num_heads * mamba_head_dim(= 2048), andn_groupsmust divide it. use_mamba_kernels=True; if the optionalmamba_ssm/causal_conv1dpackages are not installed, transformers falls back to a correct (slower) reference implementation.- Because the large 131072-token vocabulary embedding dominates the parameter budget, total and activated parameter counts are both ~0.3B despite the MoE routing.