inference-optimization/NemotronH-0.3B-A0.3B

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 mixer
• —full_attention — grouped-query attention mixer
• —moe — 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:

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.

1. 1.Reduced layers_block_type, hidden/attention/mamba sizes, and MoE sizes as above.
2. 2.Re-initialized all weights, then fine-tuned on a small toy dataset until the perplexity target was reached (validates the model can learn).
3. 3.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), and n_groups must divide it.
• —use_mamba_kernels=True; if the optional mamba_ssm / causal_conv1d packages 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.